mxnet-commits mailing list archives

Site index · List index
Message view « Date » · « Thread »
Top « Date » · « Thread »
From GitBox <...@apache.org>
Subject [GitHub] [incubator-mxnet] eric-haibin-lin commented on a change in pull request #17010: [API] unified API for custom kvstores
Date Sat, 14 Dec 2019 23:53:39 GMT
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_r357944836
 
 

 ##########
 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):
 
 Review comment:
   I will add more comments on this function for clarification with allreduce backends. 

----------------------------------------------------------------
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

Mime
View raw message