allow seq_number == global_rank for TCP backend

Differential Revision: D48130088

Pull Request resolved: #407

Author: XilunWu

gloo/common/utils.cc

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

gloo/common/utils.h

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

gloo/transport/tcp/context.cc

@@ -49,15 +49,20 @@ void Context::createAndConnectAllPairs(IStore& store) {
   int localRank = 0;
   bool localRankSet = false;
   auto localHostName = getHostname();
+  bool useRankAsSeqNum = useRankAsSeqNumber();
 
   // We will create all the pairs including self
   // the self pair will not be connected
   // it's just to keep the later seq num matching logic simple
   std::vector<ssize_t> pairIdentifiers;
   for (int i = 0; i < size; i++) {
-    auto& pair = createPair(i);
-    pairIdentifiers.emplace_back(
-        static_cast<Pair*>(pair.get())->address().getSeq());
+    const auto& pair = createPair(i, useRankAsSeqNum);
+    if (!useRankAsSeqNum) {
+      // Need to preserve the order of the pair identifiers if we are not using
+      // the rank as seq number
+      pairIdentifiers.emplace_back(
+          static_cast<Pair*>(pair.get())->address().getSeq());
+    }
   }
 
   // Obtain the pair object for this rank
@@ -105,8 +110,9 @@ void Context::createAndConnectAllPairs(IStore& store) {
 
     const auto& pair = getPair(i);
     auto remoteDeviceAddr = Address(remoteRankInfo.addressBytes).getSockaddr();
-    auto remoteAddr =
-        Address(remoteDeviceAddr, remoteRankInfo.pairIdentifiers[rank]);
+    auto remoteAddr = Address(
+        remoteDeviceAddr,
+        useRankAsSeqNum ? (ssize_t)rank : remoteRankInfo.pairIdentifiers[rank]);
     pair->connect(remoteAddr.bytes());
   }
 
@@ -124,7 +130,15 @@ void Context::createAndConnectAllPairs(IStore& store) {
 
 std::unique_ptr<transport::Pair>& Context::createPair(int rank) {
   pairs_[rank] = std::unique_ptr<transport::Pair>(
-      new tcp::Pair(this, device_.get(), rank, getTimeout()));
+      new tcp::Pair(this, device_.get(), rank, getTimeout(), false));
+  return pairs_[rank];
+}
+
+std::unique_ptr<transport::Pair>& Context::createPair(
+    int rank,
+    bool useRankAsSeqNumber = false) {
+  pairs_[rank] = std::unique_ptr<transport::Pair>(new tcp::Pair(
+      this, device_.get(), rank, getTimeout(), useRankAsSeqNumber));
   return pairs_[rank];
 }
 
@@ -305,14 +319,16 @@ Rank::Rank(const std::vector<char>& bytes) {
   bytesOffset += sizeof(addrSz) + addrSz;
   // pair identifiers
   size_t pairIdChunkSz = bytes.size() - bytesOffset;
-  GLOO_ENFORCE_EQ(
-      pairIdChunkSz % sizeof(ssize_t),
-      0,
-      "Remaining bytes do not map to entire chunk of pair identifiers");
-  size_t numPairs = pairIdChunkSz / sizeof(ssize_t);
-  pairIdentifiers.resize(numPairs);
-  std::memcpy(
-      pairIdentifiers.data(), bytes.data() + bytesOffset, pairIdChunkSz);
+  if (pairIdChunkSz) {
+    GLOO_ENFORCE_EQ(
+        pairIdChunkSz % sizeof(ssize_t),
+        0,
+        "Remaining bytes do not map to entire chunk of pair identifiers");
+    size_t numPairs = pairIdChunkSz / sizeof(ssize_t);
+    pairIdentifiers.resize(numPairs);
+    std::memcpy(
+        pairIdentifiers.data(), bytes.data() + bytesOffset, pairIdChunkSz);
+  }
 }
 
 std::vector<char> Rank::bytes() const {
@@ -336,7 +352,9 @@ std::vector<char> Rank::bytes() const {
   std::memcpy(bufOffset, addressBytes.data(), addressBytes.size());
   bufOffset += addrSz;
   // pair identifiers
-  std::memcpy(bufOffset, pairIdentifiers.data(), pairIdChunkSz);
+  if (pairIdChunkSz) {
+    std::memcpy(bufOffset, pairIdentifiers.data(), pairIdChunkSz);
+  }
   return buf;
 }
 

gloo/transport/tcp/context.h

@@ -39,6 +39,9 @@ class Context : public ::gloo::transport::Context,
   virtual void createAndConnectAllPairs(IStore& store) override;
 
   std::unique_ptr<transport::Pair>& createPair(int rank) override;
+  std::unique_ptr<transport::Pair>& createPair(
+      int rank,
+      bool useRankAsSeqNumber);
 
   std::unique_ptr<transport::UnboundBuffer> createUnboundBuffer(
       void* ptr,

gloo/transport/tcp/device.cc

@@ -253,6 +253,10 @@ Address Device::nextAddress() {
   return listener_->nextAddress();
 }
 
+Address Device::nextAddress(int seq) {
+  return listener_->nextAddress(seq);
+}
+
 bool Device::isInitiator(const Address& local, const Address& remote) const {
   int rv = 0;
   // The remote side of a pair will be called with the same

gloo/transport/tcp/device.h

@@ -68,6 +68,15 @@ class Device : public ::gloo::transport::Device,
   //
   Address nextAddress();
 
+  // Return a new `Address` instance using the provided sequence number.
+  //
+  // This is called by the constructor of the `Pair` class. It gives
+  // the pair a uniquely identifying address even though the device
+  // uses a shared listening socket. Caller must provide a unique sequence
+  // number
+  //
+  Address nextAddress(int);
+
   // Connect a pair to a remote.
   //
   // This is performed by the device instance because we use a single

gloo/transport/tcp/listener.cc

@@ -13,6 +13,7 @@
 
 #include <gloo/common/common.h>
 #include <gloo/common/logging.h>
+#include <gloo/common/utils.h>
 #include <gloo/transport/tcp/helpers.h>
 
 namespace gloo {
@@ -26,6 +27,7 @@ Listener::Listener(std::shared_ptr<Loop> loop, const attr& attr)
   listener_->bind(attr.ai_addr);
   listener_->listen(kBacklog);
   addr_ = listener_->sockName();
+  useRankAsSeqNumber_ = useRankAsSeqNumber();
 
   // Register with loop for readability events.
   loop_->registerDescriptor(listener_->fd(), EPOLLIN, this);
@@ -76,9 +78,19 @@ void Listener::handleEvents(int /* unused */) {
 
 Address Listener::nextAddress() {
   std::lock_guard<std::mutex> guard(mutex_);
+  GLOO_ENFORCE(
+      !useRankAsSeqNumber_,
+      "Listener cannot use internal sequence with enabled option to use rank as sequence number");
   return Address(addr_.getSockaddr(), seq_++);
 }
 
+Address Listener::nextAddress(int seq) {
+  GLOO_ENFORCE(
+      useRankAsSeqNumber_,
+      "Listener must be setup to use rank as sequence number");
+  return Address(addr_.getSockaddr(), seq);
+}
+
 void Listener::waitForConnection(sequence_number_t seq, connect_callback_t fn) {
   std::unique_lock<std::mutex> lock(mutex_);
 

gloo/transport/tcp/listener.h

@@ -42,6 +42,8 @@ class Listener final : public Handler {
 
   Address nextAddress();
 
+  Address nextAddress(int);
+
   // Wait for connection with sequence number `seq`. The callback is
   // always called from a different thread (the event loop thread),
   // even if the connection is already available.
@@ -67,6 +69,12 @@ class Listener final : public Handler {
 
   // Sockets by sequence number (while waiting for a pair to call).
   std::unordered_map<sequence_number_t, std::shared_ptr<Socket>> seqToSocket_;
+
+  // Option to use rank as sequence number and avoid pair identifiers
+  // to the store during rendezvous. Experimental, disabled by default.
+  // Can be enabled by setting the environment variable
+  // GLOO_ENABLE_RANK_AS_SEQUENCE_NUMBER=1.
+  bool useRankAsSeqNumber_{false};
 };
 
 } // namespace tcp

gloo/transport/tcp/pair.cc

@@ -46,7 +46,8 @@ Pair::Pair(
     Context* context,
     Device* device,
     int rank,
-    std::chrono::milliseconds timeout)
+    std::chrono::milliseconds timeout,
+    bool useRankAsSeqNumber)
     : context_(context),
       device_(device),
       rank_(rank),
@@ -56,7 +57,9 @@ Pair::Pair(
       busyPoll_(false),
       fd_(FD_INVALID),
       sendBufferSize_(0),
-      self_(device_->nextAddress()),
+      self_(
+          useRankAsSeqNumber ? device_->nextAddress(rank)
+                             : device_->nextAddress()),
       ex_(nullptr) {}
 
 // Destructor performs a "soft" close.

gloo/transport/tcp/pair.h

@@ -95,7 +95,8 @@ class Pair : public ::gloo::transport::Pair, public Handler {
       Context* context,
       Device* device,
       int rank,
-      std::chrono::milliseconds timeout);
+      std::chrono::milliseconds timeout,
+      bool useRankAsSeqNumber = false);
 
   virtual ~Pair();