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_r357944814
########## 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): Review comment: For APIs like allgather, they are not involved in the usual data parallel training practice using either module or gluon trainer. The native kvstore does not support it either. However, users are free to call hvd.allgather when they use hvd backend to gather any tensor they want. ---------------------------------------------------------------- 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: [email protected] With regards, Apache Git Services
