Use multi_get for store that has extended API support. (#408)

Summary:
Pull Request resolved: #408

The TCP store has API v2 support we can reduce the network overhead of Gloo rendezvous significantly by fetching a batch of key instead of doing them one by one.
Initial testing shows ~15X improvement for 4k jobs.

Gloo process group init:

Baseline ( fbcode trunk):
2k job (https://fburl.com/mlhub/x1prxu89) : ~82sec (~1.4 min)
4k job (https://fburl.com/mlhub/v1djk4n5) : ~393 sec (~6.6min)
8k job (https://fburl.com/mlhub/cagqrs7m): (~55mins)

With optimizations (D48130088 + D52083376):

2k job (https://fburl.com/mlhub/x0cskdag)  : ~18 sec ( ~5x faster)
4k job (https://fburl.com/mlhub/xzmvkm4j) : ~ 25 sec (~15x faster)
8k job (https://fburl.com/mlhub/gdyeizv9)   : ~ 85 sec (~35x faster)

Reviewed By: xunnanxu

Differential Revision: D52083376

Author: Boris Sarana

gloo/common/store.h

@@ -25,6 +25,21 @@ class IStore {
   virtual void wait(
       const std::vector<std::string>& keys,
       const std::chrono::milliseconds& timeout) = 0;
+
+  // Extended 2.0 API support
+  virtual bool has_v2_support() = 0;
+
+  virtual std::vector<std::vector<char>> multi_get(
+      const std::vector<std::string>& keys) = 0;
+
+  virtual void multi_set(
+      const std::vector<std::string>& keys,
+      const std::vector<std::vector<char>>& values) = 0;
+
+  virtual void append(
+      const std::string& key,
+      const std::vector<char>& value) = 0;
+  virtual int64_t add(const std::string& key, int64_t value) = 0;
 };
 
 } // namespace gloo

gloo/common/utils.cc

@@ -36,4 +36,10 @@ bool useRankAsSeqNumber() {
       (std::string(res) == "True" || std::string(res) == "1");
 }
 
+bool isStoreExtendedApiEnabled() {
+  const auto& res = std::getenv("GLOO_ENABLE_STORE_V2_API");
+  return res != nullptr &&
+      (std::string(res) == "True" || std::string(res) == "1");
+}
+
 } // namespace gloo

gloo/common/utils.h

@@ -16,4 +16,6 @@ std::string getHostname();
 
 bool useRankAsSeqNumber();
 
+bool isStoreExtendedApiEnabled();
+
 } // namespace gloo

gloo/transport/tcp/context.cc

@@ -8,10 +8,12 @@
 
 #include "gloo/transport/tcp/context.h"
 
+#include <algorithm>
+#include <cstdint>
 #include <cstring>
 #include <iostream>
+#include <string>
 
-#include "gloo/common/error.h"
 #include "gloo/common/logging.h"
 #include "gloo/common/utils.h"
 #include "gloo/transport/tcp/device.h"
@@ -22,6 +24,8 @@ namespace gloo {
 namespace transport {
 namespace tcp {
 
+constexpr int kDefaultBatchSize = 128;
+
 Context::Context(std::shared_ptr<Device> device, int rank, int size)
     : ::gloo::transport::Context(rank, size), device_(std::move(device)) {}
 
@@ -78,12 +82,36 @@ void Context::createAndConnectAllPairs(IStore& store) {
   // which does not have the rank info hosted at a higher `Pair` level).
   // So better safe than sorry for now we try to minimize the changeset needed.
   const auto& currentRankPair = getPair(rank);
-  auto deviceAddress = Address(
+  const auto& deviceAddress = Address(
       static_cast<const Pair*>(currentRankPair.get())->address().getSockaddr());
   Rank currentRankInfo(
       localHostName, deviceAddress.bytes(), std::move(pairIdentifiers));
   store.set(std::to_string(rank), currentRankInfo.bytes());
 
+  std::vector<std::vector<char>> remoteRankInfos;
+  int key = 0;
+  if (isStoreExtendedApiEnabled() && store.has_v2_support()) {
+    auto sizeRemaining = size;
+    while (sizeRemaining > 0) {
+      const auto batchKeys = std::min(kDefaultBatchSize, sizeRemaining);
+      std::vector<std::string> keys(batchKeys);
+      std::generate_n(
+          keys.begin(), batchKeys, [&] { return std::to_string(key++); });
+      const auto& batchRemoteInfos = store.multi_get(keys);
+      remoteRankInfos.insert(
+          remoteRankInfos.end(),
+          batchRemoteInfos.begin(),
+          batchRemoteInfos.end());
+      sizeRemaining -= batchKeys;
+    }
+  } else {
+    std::generate_n(std::back_inserter(remoteRankInfos), size, [&] {
+      const auto& keyStr = std::to_string(key++);
+      store.wait({keyStr.c_str()}, getTimeout());
+      return store.get(keyStr);
+    });
+  }
+
   // Connect every pair
   for (int i = 0; i < size; i++) {
     if (i == rank) {
@@ -95,24 +123,18 @@ void Context::createAndConnectAllPairs(IStore& store) {
       continue;
     }
 
-    // Wait for address of other side of this pair to become available
-    std::ostringstream key;
-    key << i;
-    store.wait({key.str()}, getTimeout());
+    Rank remoteRankInfo(remoteRankInfos[i]);
 
-    // Connect to other side of this pair
-    std::vector<char> rankInfoBytes = store.get(key.str());
-    Rank remoteRankInfo(rankInfoBytes);
-    const auto& remoteHostname = remoteRankInfo.hostname;
-    if (!localRankSet && remoteHostname == localHostName) {
+    if (!localRankSet && remoteRankInfo.hostname == localHostName) {
       ++localRank;
     }
 
     const auto& pair = getPair(i);
     auto remoteDeviceAddr = Address(remoteRankInfo.addressBytes).getSockaddr();
     auto remoteAddr = Address(
         remoteDeviceAddr,
-        useRankAsSeqNum ? (ssize_t)rank : remoteRankInfo.pairIdentifiers[rank]);
+        useRankAsSeqNum ? (sequence_number_t)rank
+                        : remoteRankInfo.pairIdentifiers[rank]);
     pair->connect(remoteAddr.bytes());
   }