gloo: use shared Stores (#423)

Summary:
X-link: pytorch/pytorch#150230


This modifies `connectFullMesh` to take in a shared_ptr<IStore> instead of a reference. This is an API breaking change but fairly easy to work around.

To have backwards compatibility in PyTorch during the commit phase we add a new ifdef `GLOO_SHARED_STORE` which can provide backwards compatibility until we update the pinned Gloo version in pytorch OSS repo.

This also adds a new `wait_get` method to `IStore` which will allow us to do a more efficient operation in PyTorch TCPStore. PyTorch's `Store::get` automatically waits so we want to make sure we can avoid waiting twice to reduce network traffic.

This change will land simultaneously in PyTorch and Gloo repos.

Reviewed By: fduwjj

Differential Revision: D72084111

Author: d4l3k

gloo/benchmark/runner.cc

@@ -187,8 +187,10 @@ void Runner::rendezvousRedis() {
     return;
   }
 
-  rendezvous::RedisStore redisStore(options_.redisHost, options_.redisPort);
-  rendezvous::PrefixStore prefixStore(options_.prefix, redisStore);
+  auto redisStore = std::make_shared<rendezvous::RedisStore>(
+      options_.redisHost, options_.redisPort);
+  auto prefixStore =
+      std::make_shared<rendezvous::PrefixStore>(options_.prefix, redisStore);
   auto backingContext = std::make_shared<rendezvous::Context>(
       options_.contextRank, options_.contextSize);
   backingContext->connectFullMesh(prefixStore, transportDevices_.front());
@@ -221,14 +223,15 @@ void Runner::rendezvousFileSystem() {
     return;
   }
 
-  rendezvous::FileStore fileStore(options_.sharedPath);
-  rendezvous::PrefixStore prefixStore(options_.prefix, fileStore);
+  auto fileStore = std::make_shared<rendezvous::FileStore>(options_.sharedPath);
+  auto prefixStore =
+      std::make_shared<rendezvous::PrefixStore>(options_.prefix, fileStore);
   auto backingContext = std::make_shared<rendezvous::Context>(
       options_.contextRank, options_.contextSize);
   backingContext->connectFullMesh(prefixStore, transportDevices_.front());
   // After connectFullMesh is called, the rendezvous files will have been
   // generated so we need to fetch them from the FileStore
-  keyFilePaths_ = fileStore.getAllKeyFilePaths();
+  keyFilePaths_ = fileStore->getAllKeyFilePaths();
   contextFactory_ =
       std::make_shared<rendezvous::ContextFactory>(backingContext);
 }

gloo/common/store.h

@@ -9,9 +9,12 @@
 #pragma once
 
 #include <chrono>
+#include <memory>
 #include <string>
 #include <vector>
 
+#define GLOO_SHARED_STORE
+
 namespace gloo {
 
 class IStore {
@@ -22,6 +25,13 @@ class IStore {
 
   virtual std::vector<char> get(const std::string& key) = 0;
 
+  virtual std::vector<char> wait_get(
+      const std::string& key,
+      const std::chrono::milliseconds& timeout) {
+    wait({key}, timeout);
+    return get(key);
+  }
+
   virtual void wait(
       const std::vector<std::string>& keys,
       const std::chrono::milliseconds& timeout) = 0;

gloo/examples/example1.cc

@@ -82,14 +82,15 @@ int main(void) {
   // Below, we instantiate rendezvous using the filesystem, given that
   // this example uses multiple processes on a single machine.
   //
-  auto fileStore = gloo::rendezvous::FileStore("/tmp");
+  auto fileStore = std::make_shared<gloo::rendezvous::FileStore>("/tmp");
 
   // To be able to reuse the same store over and over again and not have
   // interference between runs, we scope it to a unique prefix with the
   // PrefixStore. This wraps another store and prefixes every key before
   // forwarding the call to the underlying store.
   std::string prefix = getenv("PREFIX");
-  auto prefixStore = gloo::rendezvous::PrefixStore(prefix, fileStore);
+  auto prefixStore =
+      std::make_shared<gloo::rendezvous::PrefixStore>(prefix, fileStore);
 
   // Using this store, we can now create a Gloo context. The context
   // holds a reference to every communication pair involving this

gloo/rendezvous/context.cc

@@ -11,6 +11,7 @@
 #include <memory>
 
 #include "gloo/common/logging.h"
+#include "gloo/rendezvous/store.h"
 #include "gloo/transport/address.h"
 
 namespace gloo {
@@ -22,12 +23,12 @@ Context::Context(int rank, int size, int base)
 Context::~Context() {}
 
 void Context::connectFullMesh(
-    rendezvous::Store& store,
+    std::shared_ptr<rendezvous::Store> store,
     std::shared_ptr<transport::Device>& dev) {
   auto transportContext = dev->createContext(rank, size);
   transportContext->setTimeout(getTimeout());
 
-  transportContext->createAndConnectAllPairs(store);
+  transportContext->createAndConnectAllPairs(std::move(store));
 
   device_ = dev;
   transportContext_ = std::move(transportContext);

gloo/rendezvous/context.h

@@ -29,7 +29,9 @@ class Context : public ::gloo::Context {
   Context(int rank, int size, int base = 2);
   virtual ~Context();
 
-  void connectFullMesh(Store& store, std::shared_ptr<transport::Device>& dev);
+  void connectFullMesh(
+      std::shared_ptr<Store> store,
+      std::shared_ptr<transport::Device>& dev);
 
  protected:
   friend class ContextFactory;

gloo/rendezvous/prefix_store.cc

@@ -13,8 +13,10 @@
 namespace gloo {
 namespace rendezvous {
 
-PrefixStore::PrefixStore(const std::string& prefix, Store& store)
-    : prefix_(prefix), store_(store) {}
+PrefixStore::PrefixStore(
+    const std::string& prefix,
+    std::shared_ptr<Store> store)
+    : prefix_(prefix), store_(std::move(store)) {}
 
 std::string PrefixStore::joinKey(const std::string& key) {
   std::stringstream ss;
@@ -23,11 +25,11 @@ std::string PrefixStore::joinKey(const std::string& key) {
 }
 
 void PrefixStore::set(const std::string& key, const std::vector<char>& data) {
-  store_.set(joinKey(key), data);
+  store_->set(joinKey(key), data);
 }
 
 std::vector<char> PrefixStore::get(const std::string& key) {
-  return store_.get(joinKey(key));
+  return store_->get(joinKey(key));
 }
 
 void PrefixStore::wait(
@@ -38,56 +40,56 @@ void PrefixStore::wait(
   for (const auto& key : keys) {
     joinedKeys.push_back(joinKey(key));
   }
-  store_.wait(joinedKeys, timeout);
+  store_->wait(joinedKeys, timeout);
 }
 
 bool PrefixStore::has_v2_support() {
-  return store_.has_v2_support();
+  return store_->has_v2_support();
 }
 
 std::vector<std::vector<char>> PrefixStore::multi_get(
     const std::vector<std::string>& keys) {
-  if (!store_.has_v2_support()) {
+  if (!store_->has_v2_support()) {
     GLOO_THROW_INVALID_OPERATION_EXCEPTION(
         "underlying store doesn't support multi_get");
   }
   std::vector<std::string> prefixed_keys;
   for (auto& key : keys) {
     prefixed_keys.push_back(joinKey(key));
   }
-  return store_.multi_get(prefixed_keys);
+  return store_->multi_get(prefixed_keys);
 }
 
 void PrefixStore::multi_set(
     const std::vector<std::string>& keys,
     const std::vector<std::vector<char>>& values) {
-  if (!store_.has_v2_support()) {
+  if (!store_->has_v2_support()) {
     GLOO_THROW_INVALID_OPERATION_EXCEPTION(
         "underlying store doesn't support multi_set");
   }
   std::vector<std::string> prefixed_keys;
   for (auto& key : keys) {
     prefixed_keys.push_back(joinKey(key));
   }
-  return store_.multi_set(prefixed_keys, values);
+  return store_->multi_set(prefixed_keys, values);
 }
 
 void PrefixStore::append(
     const std::string& key,
     const std::vector<char>& data) {
-  if (!store_.has_v2_support()) {
+  if (!store_->has_v2_support()) {
     GLOO_THROW_INVALID_OPERATION_EXCEPTION(
         "underlying store doesn't support append");
   }
-  store_.append(joinKey(key), data);
+  store_->append(joinKey(key), data);
 }
 
 int64_t PrefixStore::add(const std::string& key, int64_t value) {
-  if (!store_.has_v2_support()) {
+  if (!store_->has_v2_support()) {
     GLOO_THROW_INVALID_OPERATION_EXCEPTION(
         "underlying store doesn't support append");
   }
-  return store_.add(joinKey(key), value);
+  return store_->add(joinKey(key), value);
 }
 
 } // namespace rendezvous

gloo/rendezvous/prefix_store.h

@@ -17,7 +17,7 @@ namespace rendezvous {
 
 class PrefixStore : public Store {
  public:
-  PrefixStore(const std::string& prefix, Store& store);
+  PrefixStore(const std::string& prefix, std::shared_ptr<Store> store);
 
   virtual ~PrefixStore() {}
 
@@ -46,7 +46,7 @@ class PrefixStore : public Store {
 
  protected:
   const std::string prefix_;
-  Store& store_;
+  std::shared_ptr<Store> store_;
 
   std::string joinKey(const std::string& key);
 };

gloo/test/base_test.h

@@ -106,7 +106,7 @@ class BaseTest : public ::testing::Test {
       std::function<void(std::shared_ptr<Context>)> fn,
       int base = 2) {
     Barrier barrier(size);
-    ::gloo::rendezvous::HashStore store;
+    auto store = std::make_shared<::gloo::rendezvous::HashStore>();
 
     spawnThreads(size, [&](int rank) {
       auto context =

gloo/test/multiproc_test.h

@@ -101,14 +101,14 @@ class MultiProcWorker {
     auto context = std::make_shared<::gloo::rendezvous::Context>(rank, size);
     auto device = createDevice(transport);
     context->setTimeout(std::chrono::milliseconds(kMultiProcTimeout));
-    context->connectFullMesh(*store_, device);
+    context->connectFullMesh(store_, device);
     device.reset();
     sem_post(semaphore_);
     fn(std::move(context));
   }
 
  protected:
-  std::unique_ptr<::gloo::rendezvous::Store> store_;
+  std::shared_ptr<::gloo::rendezvous::Store> store_;
   sem_t* semaphore_;
 };
 

gloo/transport/context.cc

@@ -23,7 +23,7 @@ std::unique_ptr<transport::Pair>& Context::getPair(int rank_2) {
   return pairs_.at(rank_2);
 }
 
-void Context::createAndConnectAllPairs(IStore& store) {
+void Context::createAndConnectAllPairs(std::shared_ptr<IStore> store) {
   // this is the default un-optimized version of the rendezvous protocol
   // where each rank would write N pairs to the store
   // and then for each remote peer load the N addresses
@@ -40,15 +40,15 @@ void Context::createAndConnectAllPairs(IStore& store) {
   // hostname mapping to compute local ranks.
   std::string localKey("rank_" + std::to_string(rank));
   const std::vector<char> value(localHostName.begin(), localHostName.end());
-  store.set(localKey, value);
+  store->set(localKey, value);
 
   for (int i = 0; i < size; i++) {
     if (i == rank) {
       break;
     }
 
     std::string key("rank_" + std::to_string(i));
-    auto val = store.get(key);
+    auto val = store->get(key);
     auto hostName = std::string((const char*)val.data(), val.size());
 
     if (hostName == localHostName) {
@@ -68,7 +68,7 @@ void Context::createAndConnectAllPairs(IStore& store) {
     allBytes.insert(allBytes.end(), addrBytes.begin(), addrBytes.end());
   }
 
-  store.set(std::to_string(rank), allBytes);
+  store->set(std::to_string(rank), allBytes);
 
   // Connect every pair
   for (int i = 0; i < size; i++) {
@@ -79,10 +79,10 @@ void Context::createAndConnectAllPairs(IStore& store) {
     // Wait for address of other side of this pair to become available
     std::ostringstream key;
     key << i;
-    store.wait({key.str()}, getTimeout());
+    store->wait({key.str()}, getTimeout());
 
     // Connect to other side of this pair
-    auto allAddrs = store.get(key.str());
+    auto allAddrs = store->get(key.str());
     auto addr = extractAddress(allAddrs, i);
     getPair(i)->connect(addr);
   }