pytorch · osalpekar · Sep 1, 2020 · Sep 11, 2020 · Sep 11, 2020 · Sep 18, 2020
diff --git a/caffe2/distributed/file_store_handler.cc b/caffe2/distributed/file_store_handler.cc
@@ -122,6 +122,11 @@ int64_t FileStoreHandler::add(
   return 0;
 }
 
+int64_t FileStoreHandler::getNumKeys() {
+  CHECK(false) << "getNumKeys not implemented for FileStoreHandler";
+  return 0;
+}
+
 bool FileStoreHandler::check(const std::vector<std::string>& names) {
   std::vector<std::string> paths;
   for (const auto& name : names) {

diff --git a/caffe2/distributed/file_store_handler.h b/caffe2/distributed/file_store_handler.h
@@ -17,6 +17,8 @@ class CAFFE2_API FileStoreHandler : public StoreHandler {
 
   virtual int64_t add(const std::string& name, int64_t value) override;
 
+  virtual int64_t getNumKeys() override;
+
   virtual bool check(const std::vector<std::string>& names) override;
 
   virtual void wait(

diff --git a/caffe2/distributed/redis_store_handler.cc b/caffe2/distributed/redis_store_handler.cc
@@ -76,6 +76,11 @@ int64_t RedisStoreHandler::add(const std::string& name, int64_t value) {
   return reply->integer;
 }
 
+int64_t RedisStoreHandler::getNumKeys() {
+  CHECK(false) << "getNumKeys not implemented for RedisStoreHandler";
+  return 0;
+}
+
 bool RedisStoreHandler::check(const std::vector<std::string>& names) {
   std::vector<std::string> args;
   args.push_back("EXISTS");

diff --git a/caffe2/distributed/redis_store_handler.h b/caffe2/distributed/redis_store_handler.h
@@ -23,6 +23,8 @@ class CAFFE2_API RedisStoreHandler : public StoreHandler {
 
   virtual int64_t add(const std::string& name, int64_t value) override;
 
+  virtual int64_t getNumKeys() override;
+
   virtual bool check(const std::vector<std::string>& names) override;
 
   virtual void wait(

diff --git a/caffe2/distributed/store_handler.h b/caffe2/distributed/store_handler.h
@@ -41,6 +41,11 @@ class CAFFE2_API StoreHandler {
    */
   virtual int64_t add(const std::string& name, int64_t value) = 0;
 
+  /*
+   * Returns the number of keys in this store.
+   */
+  virtual int64_t getNumKeys() = 0;
+
   /*
    * Check if a keys exist in the store.
    */

diff --git a/torch/csrc/distributed/c10d/init.cpp b/torch/csrc/distributed/c10d/init.cpp
@@ -92,6 +92,10 @@ class PythonStore : public ::c10d::Store {
     PYBIND11_OVERLOAD_PURE(int64_t, ::c10d::Store, add, key, value);
   }
 
+  int64_t getNumKeys() override {
+    PYBIND11_OVERLOAD_PURE(int64_t, ::c10d::Store, getNumKeys);
+  }
+
   bool check(const std::vector<std::string>& keys) override {
     PYBIND11_OVERLOAD_PURE(bool, ::c10d::Store, check, keys);
   }
@@ -301,6 +305,10 @@ They are used in specifying strategies for reduction collectives, e.g.,
               "add",
               &::c10d::Store::add,
               py::call_guard<py::gil_scoped_release>())
+          .def(
+              "num_keys",
+              &::c10d::Store::getNumKeys,
+              py::call_guard<py::gil_scoped_release>())
           .def(
               "set_timeout",
               &::c10d::Store::setTimeout,

diff --git a/torch/lib/c10d/FileStore.cpp b/torch/lib/c10d/FileStore.cpp
@@ -16,6 +16,8 @@
 #include <system_error>
 #include <thread>
 
+#include <c10/util/Exception.h>
+
 #define SYSASSERT(rv, ...)                                                 \
   if ((rv) < 0) {                                                          \
     throw std::system_error(errno, std::system_category(), ##__VA_ARGS__); \
@@ -303,6 +305,10 @@ int64_t FileStore::add(const std::string& key, int64_t value) {
   return addHelper(regKey, value);
 }
 
+int64_t FileStore::getNumKeys() {
+  TORCH_CHECK(false, "getNumKeys not implemented for FileStore");
+}
+
 bool FileStore::check(const std::vector<std::string>& keys) {
   std::unique_lock<std::mutex> l(activeFileOpLock_);
   File file(path_, O_RDONLY, timeout_);

diff --git a/torch/lib/c10d/FileStore.hpp b/torch/lib/c10d/FileStore.hpp
@@ -21,6 +21,8 @@ class FileStore : public Store {
 
   int64_t add(const std::string& key, int64_t value) override;
 
+  int64_t getNumKeys() override;
+
   bool check(const std::vector<std::string>& keys) override;
 
   void wait(const std::vector<std::string>& keys) override;

diff --git a/torch/lib/c10d/HashStore.cpp b/torch/lib/c10d/HashStore.cpp
@@ -8,6 +8,8 @@
 #include <cstdio>
 #include <system_error>
 
+#include <c10/util/Exception.h>
+
 namespace c10d {
 
 void HashStore::set(const std::string& key, const std::vector<uint8_t>& data) {
@@ -77,6 +79,10 @@ int64_t HashStore::add(const std::string& key, int64_t i) {
   return ti;
 }
 
+int64_t HashStore::getNumKeys() {
+  TORCH_CHECK(false, "getNumKeys not implemented for HashStore");
+}
+
 bool HashStore::check(const std::vector<std::string>& keys) {
   std::unique_lock<std::mutex> lock(m_);
   for (const auto& key : keys) {

diff --git a/torch/lib/c10d/HashStore.hpp b/torch/lib/c10d/HashStore.hpp
@@ -28,6 +28,8 @@ class HashStore : public Store {
 
   int64_t add(const std::string& key, int64_t value) override;
 
+  int64_t getNumKeys() override;
+
   bool check(const std::vector<std::string>& keys) override;
 
  protected:

diff --git a/torch/lib/c10d/PrefixStore.cpp b/torch/lib/c10d/PrefixStore.cpp
@@ -35,6 +35,10 @@ int64_t PrefixStore::add(const std::string& key, int64_t value) {
   return store_->add(joinKey(key), value);
 }
 
+int64_t PrefixStore::getNumKeys() {
+  return store_->getNumKeys();
+}
+
 bool PrefixStore::check(const std::vector<std::string>& keys) {
   auto joinedKeys = joinKeys(keys);
   return store_->check(joinedKeys);

diff --git a/torch/lib/c10d/PrefixStore.hpp b/torch/lib/c10d/PrefixStore.hpp
@@ -17,6 +17,8 @@ class PrefixStore : public Store {
 
   int64_t add(const std::string& key, int64_t value) override;
 
+  int64_t getNumKeys() override;
+
   bool check(const std::vector<std::string>& keys) override;
 
   void wait(const std::vector<std::string>& keys) override;

diff --git a/torch/lib/c10d/Store.hpp b/torch/lib/c10d/Store.hpp
@@ -32,6 +32,8 @@ class Store {
 
   virtual bool check(const std::vector<std::string>& keys) = 0;
 
+  virtual int64_t getNumKeys() = 0;
+
   virtual void wait(const std::vector<std::string>& keys) = 0;
 
   virtual void wait(

diff --git a/torch/lib/c10d/TCPStore.cpp b/torch/lib/c10d/TCPStore.cpp
@@ -10,7 +10,7 @@ namespace c10d {
 
 namespace {
 
-enum class QueryType : uint8_t { SET, GET, ADD, CHECK, WAIT };
+enum class QueryType : uint8_t { SET, GET, ADD, CHECK, WAIT, GETNUMKEYS };
 
 enum class CheckResponseType : uint8_t { READY, NOT_READY };
 
@@ -180,6 +180,9 @@ void TCPStoreDaemon::query(int socket) {
   } else if (qt == QueryType::WAIT) {
     waitHandler(socket);
 
+  } else if (qt == QueryType::GETNUMKEYS) {
+    getNumKeysHandler(socket);
+
   } else {
     throw std::runtime_error("Unexpected query type");
   }
@@ -228,6 +231,10 @@ void TCPStoreDaemon::getHandler(int socket) const {
   tcputil::sendVector<uint8_t>(socket, data);
 }
 
+void TCPStoreDaemon::getNumKeysHandler(int socket) const {
+  tcputil::sendValue<int64_t>(socket, tcpStore_.size());
+}
+
 void TCPStoreDaemon::checkHandler(int socket) const {
   SizeType nargs;
   tcputil::recvBytes<SizeType>(socket, &nargs, 1);
@@ -364,6 +371,11 @@ int64_t TCPStore::addHelper_(const std::string& key, int64_t value) {
   return tcputil::recvValue<int64_t>(storeSocket_);
 }
 
+int64_t TCPStore::getNumKeys() {
+  tcputil::sendValue<QueryType>(storeSocket_, QueryType::GETNUMKEYS);
+  return tcputil::recvValue<int64_t>(storeSocket_);
+}
+
 bool TCPStore::check(const std::vector<std::string>& keys) {
   tcputil::sendValue<QueryType>(storeSocket_, QueryType::CHECK);
   SizeType nkeys = keys.size();

diff --git a/torch/lib/c10d/TCPStore.hpp b/torch/lib/c10d/TCPStore.hpp
@@ -26,6 +26,7 @@ class TCPStoreDaemon {
   void addHandler(int socket);
   void getHandler(int socket) const;
   void checkHandler(int socket) const;
+  void getNumKeysHandler(int socket) const;
   void waitHandler(int socket);
 
   bool checkKeys(const std::vector<std::string>& keys) const;
@@ -63,6 +64,8 @@ class TCPStore : public Store {
 
   bool check(const std::vector<std::string>& keys) override;
 
+  int64_t getNumKeys() override;
+
   void wait(const std::vector<std::string>& keys) override;
 
   void wait(

diff --git a/torch/lib/c10d/test/TCPStoreTest.cpp b/torch/lib/c10d/test/TCPStoreTest.cpp
@@ -36,6 +36,10 @@ void testHelper(const std::string& prefix = "") {
     c10d::test::check(*serverStore, "key0", "value0");
     c10d::test::check(*serverStore, "key1", "value1");
     c10d::test::check(*serverStore, "key2", "value2");
+    auto numKeys = serverStore->getNumKeys();
+    // We expect 5 keys since 3 are added above, 'counter' is added by the
+    // helper thread, and the init key to coordinate workers.
+    EXPECT_EQ(numKeys, 5);
   });
 
   // Hammer on TCPStore