[GitHub] [incubator-mxnet] eric-haibin-lin commented on a change in pull request #17010: [API] unified API for custom kvstores

[GitBox <gi...@apache.org>]

eric-haibin-lin commented on a change in pull request #17010: [API] unified API for custom kvstores
URL: https://github.com/apache/incubator-mxnet/pull/17010#discussion_r357866523
 
 

 ##########
 File path: python/mxnet/kvstore/base.py
 ##########
 @@ -0,0 +1,452 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+# coding: utf-8
+""" Key value store interface of MXNet for parameter synchronization."""
+from __future__ import absolute_import
+
+from array import array
+import ctypes
+import warnings
+from ..ndarray import NDArray
+from ..base import _LIB, c_str_array, c_handle_array, c_array, c_array_buf, c_str
+from ..base import check_call, string_types
+from ..base import KVStoreHandle
+from ..profiler import set_kvstore_handle
+
+__all__ = ['create', 'KVStoreBase']
+
+def _ctype_key_value(keys, vals):
+    """Returns ctype arrays for the key-value args, and the whether string keys are used.
+    For internal use only.
+    """
+    if isinstance(keys, (tuple, list)):
+        assert(len(keys) == len(vals))
+        c_keys = []
+        c_vals = []
+        use_str_keys = None
+        for key, val in zip(keys, vals):
+            c_key_i, c_val_i, str_keys_i = _ctype_key_value(key, val)
+            c_keys += c_key_i
+            c_vals += c_val_i
+            use_str_keys = str_keys_i if use_str_keys is None else use_str_keys
+            assert(use_str_keys == str_keys_i), "inconsistent types of keys detected."
+        c_keys_arr = c_array(ctypes.c_char_p, c_keys) if use_str_keys \
+                     else c_array(ctypes.c_int, c_keys)
+        c_vals_arr = c_array(ctypes.c_void_p, c_vals)
+        return (c_keys_arr, c_vals_arr, use_str_keys)
+
+    assert(isinstance(keys, (int,) + string_types)), \
+           "unexpected type for keys: " + str(type(keys))
+    use_str_keys = isinstance(keys, string_types)
+    if isinstance(vals, NDArray):
+        c_keys = c_str_array([keys]) if use_str_keys \
+                 else c_array_buf(ctypes.c_int, array('i', [keys]))
+        return (c_keys, c_handle_array([vals]), use_str_keys)
+    else:
+        for value in vals:
+            assert(isinstance(value, NDArray))
+        c_keys = c_str_array([keys] * len(vals)) if use_str_keys \
+                 else c_array_buf(ctypes.c_int, array('i', [keys] * len(vals)))
+        return (c_keys, c_handle_array(vals), use_str_keys)
+
+def _ctype_dict(param_dict):
+    """Returns ctype arrays for keys and values(converted to strings) in a dictionary"""
+    assert(isinstance(param_dict, dict)), \
+        "unexpected type for param_dict: " + str(type(param_dict))
+    c_keys = c_array(ctypes.c_char_p, [c_str(k) for k in param_dict.keys()])
+    c_vals = c_array(ctypes.c_char_p, [c_str(str(v)) for v in param_dict.values()])
+    return (c_keys, c_vals)
+
+class KVStoreBase(object):
+    """An abstract key-value store interface for data parallel training."""
+
+    def broadcast(self, key, value, out, priority=0):
+        """ Broadcast the `value` NDArray at rank 0 to all ranks,
+        and store the result in `out`
+
+        Parameters
+        ----------
+        key : str or int
+            The key.
+
+        value : NDArray
+            The value corresponding to the key to broadcast
+
+        out : NDArray, or list of NDArray
+            Values corresponding to the key to store the result
+
+        priority : int, optional
+            The priority of the operation.
+            Higher priority operations are likely to be executed before other actions.
+        """
+        raise NotImplementedError()
+
+    def pushpull(self, key, value, out=None, priority=0):
+        """ Performs push and pull a single value or a sequence of values from the store.
+
+        This function is coalesced form of push and pull operations.
+
+        `value` is pushed to the kvstore server for summation with the specified keys,
+        and the results are pulled from the server to `out`. If `out` is not specified
+        the pulled values are written to `value`.
+
+        Parameters
+        ----------
+        key : str or int
+            The key.
+
+        value : NDArray, or list of NDArray
+            Values corresponding to the keys.
+
+        out: NDArray, or list of NDArray
+            Values corresponding to the key.
+
+        priority : int, optional
+            The priority of the operation.
+            Higher priority operations are likely to be executed before other actions.
+        """
+        raise NotImplementedError()
+
+    def set_optimizer(self, optimizer):
+        """ Registers an optimizer with the kvstore.
+
+        When using a single machine, this function updates the local optimizer.
+        If using multiple machines and this operation is invoked from a worker node,
+        it will serialized the optimizer with pickle and send it to all servers.
+        The function returns after all servers have been updated.
+
+        Parameters
+        ----------
+        optimizer : KVStoreBase
+            The new optimizer for the store
+        """
+        raise NotImplementedError()
+
+    OPTIMIZER = 'optimizer'
+
+    @staticmethod
+    def is_capable(capability):
+        """Queries if the KVStore type supports certain capability, such as optimizer algorithm,
+        gradient compression, sparsity, etc.
+
+        Parameters
+        ----------
+        capability: str
+            The capability to query
+
+        Returns
+        -------
+        result : bool
+            Whether the capability is supported or not.
+        """
+        raise NotImplementedError()
+
+    def save_optimizer_states(self, fname, dump_optimizer=False):
+        """Saves the optimizer (updater) state to a file. This is often used when checkpointing
+        the model during training.
+
+        Parameters
+        ----------
+        fname : str
+            Path to the output states file.
+        dump_optimizer : bool, default False
+            Whether to also save the optimizer itself. This would also save optimizer
+            information such as learning rate and weight decay schedules.
+        """
+        raise NotImplementedError()
+
+    def load_optimizer_states(self, fname):
+        """Loads the optimizer (updater) state from the file.
+
+        Parameters
+        ----------
+        fname : str
+            Path to input states file.
+        """
+        raise NotImplementedError()
+
+    @property
+    def type(self):
 
 Review comment:
   what about updating the doc to be "Returns the type of this kvstore backend" ? 

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
users@infra.apache.org


With regards,
Apache Git Services