[DO NOT MERGE] Added error handling in MXNet

horovod/mxnet/__init__.py

@@ -19,6 +19,9 @@
 
 from horovod.common import check_extension
 
+check_extension('horovod.mxnet', 'HOROVOD_WITH_MXNET',
+                __file__, 'mpi_lib')
+
 from horovod.mxnet.mpi_ops import allreduce, allreduce_
 from horovod.mxnet.mpi_ops import allgather
 from horovod.mxnet.mpi_ops import broadcast, broadcast_

horovod/mxnet/adapter.cc

@@ -125,16 +125,11 @@ template <class T> Framework MXOpContext<T>::framework() const {
 }
 
 void ThrowIfError(Status status) {
-  switch (status.type()) {
-  case StatusType::OK:
+  if (StatusType::OK == status.type()) {
     return;
-  case StatusType::PRECONDITION_ERROR:
-    throw std::logic_error(status.reason());
-  case StatusType::ABORTED:
-    throw std::runtime_error(status.reason());
-  default: // Includes UNKNOWN_ERROR
-    throw std::runtime_error(status.reason());
   }
+
+  throw dmlc::Error(status.reason());
 }
 
 template class MXTensor<NDArray>;

horovod/mxnet/handle_manager.cc

@@ -0,0 +1,76 @@
+// Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+//
+// Licensed 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.
+// =============================================================================
+
+#include "handle_manager.h"
+
+namespace horovod {
+namespace mxnet {
+
+typedef ::mxnet::Engine::CallbackOnComplete Callback;
+
+int HandleManager::AllocateHandle() {
+  int handle = last_handle_.fetch_add(1) + 1;
+  std::lock_guard<std::mutex> guard(mutex_);
+  results_[handle] = nullptr;
+  return handle;
+}
+
+void HandleManager::MarkDone(int handle, const Status& status) {
+  std::lock_guard<std::mutex> guard(mutex_);
+  results_[handle] = std::make_shared<Status>(status);
+}
+
+void HandleManager::AttachCallback(int handle, Callback cb) {
+  std::lock_guard<std::mutex> guard(mutex_);
+  if (callbacks_.find(handle) == callbacks_.end()) {
+    callbacks_[handle] = std::make_shared<Callback>(cb);
+  }
+}
+
+void HandleManager::ExecuteCallback(int handle) {
+  std::unique_lock<std::mutex> lock(mutex_);
+  if (callbacks_.find(handle) == callbacks_.end()) {
+    return;
+  }
+  auto cb_ptr = callbacks_[handle];
+  lock.unlock();
+  if (cb_ptr != nullptr) {
+    (*cb_ptr)();
+  }
+}
+
+bool HandleManager::PollHandle(int handle) {
+  std::lock_guard<std::mutex> guard(mutex_);
+  if (results_.find(handle) == results_.end()) {
+    throw std::invalid_argument("Handle " + std::to_string(handle) +
+                                " was not created or has been cleared.");
+  }
+  return results_[handle] != nullptr;
+}
+
+std::shared_ptr<Status> HandleManager::ReleaseHandle(int handle) {
+  std::lock_guard<std::mutex> guard(mutex_);
+  if (results_.find(handle) == results_.end()) {
+    throw std::invalid_argument("Handle " + std::to_string(handle) +
+        " was not created or has been cleared.");
+  }
+  auto status = results_[handle];
+  results_.erase(handle);
+  callbacks_.erase(handle);
+  return status;
+}
+
+} // namespace mxnet
+} // namespace horovod

horovod/mxnet/handle_manager.h

@@ -0,0 +1,56 @@
+// Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+//
+// Licensed 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.
+// =============================================================================
+
+#ifndef HOROVOD_MXNET_HANDLE_MANAGER_H
+#define HOROVOD_MXNET_HANDLE_MANAGER_H
+
+#include <atomic>
+#include <memory>
+#include <mutex>
+#include <unordered_map>
+
+#include "../common/common.h"
+
+#include <mxnet/engine.h>
+
+namespace horovod {
+namespace mxnet {
+
+using namespace horovod::common;
+
+typedef ::mxnet::Engine Engine;
+typedef ::mxnet::NDArray NDArray;
+typedef ::mxnet::Engine::CallbackOnComplete Callback;
+
+class HandleManager {
+public:
+  int AllocateHandle();
+  void AttachCallback(int handle, Callback cb);
+  void MarkDone(int handle, const Status& status);
+  void ExecuteCallback(int handle);
+  bool PollHandle(int handle);
+  std::shared_ptr<Status> ReleaseHandle(int handle);
+
+private:
+  std::atomic_int last_handle_;
+  std::unordered_map<int, std::shared_ptr<Status>> results_;
+  std::unordered_map<int, std::shared_ptr<Callback>> callbacks_;
+  std::mutex mutex_;
+};
+
+} // namespace mxnet
+} // namespace horovod
+
+#endif // HOROVOD_MXNET_HANDLE_MANAGER_H