reduce gloo rendezvous cost for TCP backend

Summary:
The original version would repeat the local device address N times in the payload sent to store during mesh connection rendezvous.

This yields quadratic payload sent to store (N ranks, each sending N).
During later get, there's an outer for loop to grab the remote pair info, discard most of them and just use one (that matches the current rank).
In total this op is cubic which is not efficient for large scale jobs.

In reality if one device is used, for TCP based connection, only the seq number is different, the device addresses are all the same.

This change aims to reduce the payload size sent to store to linear (approx. almost constant -> 1 addr + N seq numbers).

In total this would make the original cubic op to quadratic (or more like O(1.xN) if that makes sense)

Reviewed By: bmaurer

Differential Revision: D45740631

fbshipit-source-id: aa089bfd81f3d392c0aa23a65f497bbacbdf0384

Author: xunnanxu

gloo/common/CMakeLists.txt

@@ -1,13 +1,16 @@
 set(GLOO_COMMON_SRCS
   "${CMAKE_CURRENT_SOURCE_DIR}/logging.cc"
+  "${CMAKE_CURRENT_SOURCE_DIR}/utils.cc"
   )
 
 set(GLOO_COMMON_HDRS
   "${CMAKE_CURRENT_SOURCE_DIR}/aligned_allocator.h"
   "${CMAKE_CURRENT_SOURCE_DIR}/common.h"
   "${CMAKE_CURRENT_SOURCE_DIR}/error.h"
   "${CMAKE_CURRENT_SOURCE_DIR}/logging.h"
+  "${CMAKE_CURRENT_SOURCE_DIR}/store.h"
   "${CMAKE_CURRENT_SOURCE_DIR}/string.h"
+  "${CMAKE_CURRENT_SOURCE_DIR}/utils.h"
   )
 
 if(${CMAKE_SYSTEM_NAME} STREQUAL "Linux")

gloo/common/store.h

@@ -0,0 +1,30 @@
+/**
+ * Copyright (c) 2017-present, Facebook, Inc.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <chrono>
+#include <string>
+#include <vector>
+
+namespace gloo {
+
+class IStore {
+ public:
+  virtual ~IStore() = default;
+
+  virtual void set(const std::string& key, const std::vector<char>& data) = 0;
+
+  virtual std::vector<char> get(const std::string& key) = 0;
+
+  virtual void wait(
+    const std::vector<std::string>& keys,
+    const std::chrono::milliseconds& timeout) = 0;
+};
+
+} // namespace gloo

gloo/common/utils.cc

@@ -0,0 +1,28 @@
+/**
+ * Copyright (c) 2017-present, Facebook, Inc.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <system_error>
+#include <unistd.h>
+
+#include "gloo/common/utils.h"
+
+namespace gloo {
+
+constexpr int HOSTNAME_MAX_SIZE = 192;
+
+std::string getHostname() {
+  // Get Hostname using syscall
+  char hostname[HOSTNAME_MAX_SIZE]; // NOLINT
+  int rv = gethostname(hostname, HOSTNAME_MAX_SIZE);
+  if (rv != 0) {
+    throw std::system_error(errno, std::system_category());
+  }
+  return std::string(hostname);
+}
+
+} // namespace gloo

gloo/common/utils.h

@@ -0,0 +1,17 @@
+/**
+ * Copyright (c) 2017-present, Facebook, Inc.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <string>
+
+namespace gloo {
+
+std::string getHostname();
+
+} // namespace gloo

gloo/rendezvous/context.cc

@@ -8,6 +8,8 @@
 
 #include "gloo/rendezvous/context.h"
 
+#include <memory>
+
 #include "gloo/common/logging.h"
 #include "gloo/transport/address.h"
 
@@ -30,85 +32,13 @@ Context::Context(int rank, int size, int base)
 Context::~Context() {
 }
 
-std::vector<char> Context::extractAddress(
-    std::vector<char>& allAddrs, int i) {
-  // Extract address from the list of all addresses
-  int adjRank = (rank > i ? rank - 1 : rank);
-  // Adjust for the fact that nodes do not store address for themselves
-  int addrSize = allAddrs.size() / (size - 1);
-  return std::vector<char>(allAddrs.begin() + adjRank * addrSize,
-                           allAddrs.begin() + (adjRank + 1) * addrSize);
-}
-
 void Context::connectFullMesh(
     rendezvous::Store& store,
     std::shared_ptr<transport::Device>& dev) {
-  std::vector<char> allBytes;
-  int localRank = 0;
-
-  // Get Hostname using syscall
-  char hostname[HOSTNAME_MAX_SIZE]; // NOLINT
-  int rv = gethostname(hostname, HOSTNAME_MAX_SIZE);
-  if (rv != 0) {
-    throw std::system_error(errno, std::system_category());
-  }
-
-  auto localHostName = std::string(hostname);
-  // Add global rank <> hostname pair to the Store. This store is then passed
-  // to Gloo when connectFullMesh is called, where Gloo uses the global rank <>
-  // 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);
-
-  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 hostName = std::string((const char*)val.data(), val.size());
-
-    if (hostName == localHostName) {
-      localRank++;
-    }
-  }
-
-  // Create pairs
   auto transportContext = dev->createContext(rank, size);
   transportContext->setTimeout(getTimeout());
-  for (int i = 0; i < size; i++) {
-    if (i == rank) {
-      continue;
-    }
-
-    auto& pair = transportContext->createPair(i);
-    pair->setLocalRank(localRank);
-    auto addrBytes = pair->address().bytes();
-    allBytes.insert(allBytes.end(), addrBytes.begin(), addrBytes.end());
-  }
-
-  std::ostringstream storeKey;
-  storeKey << rank;
-  store.set(storeKey.str(), allBytes);
 
-  // Connect every pair
-  for (int i = 0; i < size; i++) {
-    if (i == rank) {
-      continue;
-    }
-
-    // Wait for address of other side of this pair to become available
-    std::ostringstream key;
-    key << i;
-    store.wait({key.str()}, getTimeout());
-
-    // Connect to other side of this pair
-    auto allAddrs = store.get(key.str());
-    auto addr = extractAddress(allAddrs, i);
-    transportContext->getPair(i)->connect(addr);
-  }
+  transportContext->createAndConnectAllPairs(store);
 
   device_ = dev;
   transportContext_ = std::move(transportContext);

gloo/rendezvous/context.h

@@ -9,9 +9,11 @@
 #pragma once
 
 #include <memory>
+#include <string>
 #include <vector>
 
 #include "gloo/common/error.h"
+#include "gloo/common/store.h"
 #include "gloo/context.h"
 #include "gloo/rendezvous/store.h"
 #include "gloo/transport/address.h"
@@ -32,8 +34,6 @@ class Context : public ::gloo::Context {
       std::shared_ptr<transport::Device>& dev);
 
  protected:
-  std::vector<char> extractAddress(std::vector<char>& allAddrs, int i);
-
   friend class ContextFactory;
 };
 

gloo/rendezvous/store.h

@@ -14,14 +14,15 @@
 
 #include "gloo/common/logging.h"
 #include "gloo/common/error.h"
+#include "gloo/common/store.h"
 
 //can be used by upstream users to know whether this is available or not.
 #define GLOO_STORE_HAS_STORE_V2 1
 
 namespace gloo {
 namespace rendezvous {
 
-class Store {
+class Store: public IStore {
  public:
   static constexpr std::chrono::milliseconds kDefaultTimeout =
       std::chrono::seconds(30);

gloo/transport/address.h

@@ -22,6 +22,7 @@ class Address {
   virtual ~Address() = 0;
 
   virtual std::string str() const = 0;
+
   virtual std::vector<char> bytes() const = 0;
 };
 

gloo/transport/context.cc

@@ -6,6 +6,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+#include "gloo/common/utils.h"
 #include "gloo/transport/context.h"
 
 namespace gloo {
@@ -22,6 +23,81 @@ std::unique_ptr<transport::Pair>& Context::getPair(int rank) {
   return pairs_.at(rank);
 }
 
+void Context::createAndConnectAllPairs(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
+  // and only pick the 1 useful
+  // A more efficient version (for transport supporting multiplexing like TCP)
+  // can be seen in gloo/transport/tcp/context.cc
+
+  std::vector<char> allBytes;
+  int localRank = 0;
+
+  auto localHostName = getHostname();
+  // Add global rank <> hostname pair to the Store. This store is then passed
+  // to Gloo when connectFullMesh is called, where Gloo uses the global rank <>
+  // 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);
+
+  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 hostName = std::string((const char*)val.data(), val.size());
+
+    if (hostName == localHostName) {
+      localRank++;
+    }
+  }
+
+  // Create pairs
+  for (int i = 0; i < size; i++) {
+    if (i == rank) {
+      continue;
+    }
+
+    auto& pair = createPair(i);
+    pair->setLocalRank(localRank);
+    auto addrBytes = pair->address().bytes();
+    allBytes.insert(allBytes.end(), addrBytes.begin(), addrBytes.end());
+  }
+
+  store.set(std::to_string(rank), allBytes);
+
+  // Connect every pair
+  for (int i = 0; i < size; i++) {
+    if (i == rank) {
+      continue;
+    }
+
+    // Wait for address of other side of this pair to become available
+    std::ostringstream key;
+    key << i;
+    store.wait({key.str()}, getTimeout());
+
+    // Connect to other side of this pair
+    auto allAddrs = store.get(key.str());
+    auto addr = extractAddress(allAddrs, i);
+    getPair(i)->connect(addr);
+  }
+}
+
+std::vector<char> Context::extractAddress(
+    const std::vector<char>& allAddrs, int i) const {
+  // Extract address from the list of all addresses
+  int adjRank = (rank > i ? rank - 1 : rank);
+  // Adjust for the fact that nodes do not store address for themselves
+  int addrSize = allAddrs.size() / (size - 1);
+  return std::vector<char>(allAddrs.begin() + adjRank * addrSize,
+                           allAddrs.begin() + (adjRank + 1) * addrSize);
+}
+
 Context::LazyTally::LazyTally(std::vector<Tally>& vec, slot_t slot)
     : vec_(vec), slot_(slot), initialized_(false) {}
 

gloo/transport/context.h

@@ -17,6 +17,7 @@
 #include <unordered_map>
 #include <vector>
 
+#include "gloo/common/store.h"
 #include "gloo/transport/pair.h"
 #include "gloo/transport/unbound_buffer.h"
 
@@ -48,6 +49,8 @@ class Context {
 
   virtual std::unique_ptr<Pair>& createPair(int rank) = 0;
 
+  virtual void createAndConnectAllPairs(IStore& store);
+
   // Creates unbound buffer to be used with the ranks in this context.
   // It is not bound to a specific rank, but still bound to this
   // context. This is needed to support recv-from-any semantics, where
@@ -90,6 +93,8 @@ class Context {
   // any kind of send/recv operation.
   std::chrono::milliseconds timeout_;
 
+  std::vector<char> extractAddress(const std::vector<char>& allAddrs, int i) const;
+
  protected:
   // Keep track of pending send and recv notifications or operations
   // for a single slot.

gloo/transport/ibverbs/address.h

@@ -31,6 +31,8 @@ class Address : public ::gloo::transport::Address {
   virtual std::string str() const override;
 
  protected:
+  explicit Address(const Address&) = default;
+
   struct {
     uint32_t lid;
     uint32_t qpn;

gloo/transport/tcp/address.cc

@@ -9,6 +9,7 @@
 #include "gloo/transport/tcp/address.h"
 
 #include <arpa/inet.h>
+#include <memory>
 #include <mutex>
 #include <string.h>