gloo: async

Differential Revision: D69698406

Pull Request resolved: #418

Author: d4l3k

gloo/test/allgather_test.cc

@@ -167,6 +167,11 @@ TEST_F(AllgatherNewTest, TestTimeout) {
     AllgatherOptions opts(context);
     opts.setInput(input.getPointer(), 1);
     opts.setOutput(output.getPointer(), context->size);
+
+    // Run one operation first so we're measuring the operation timeout not
+    // connection timeout.
+    allgather(opts);
+
     opts.setTimeout(std::chrono::milliseconds(10));
     if (context->rank == 0) {
       try {

gloo/test/allgatherv_test.cc

@@ -95,6 +95,11 @@ TEST_F(AllgathervTest, TestTimeout) {
     std::vector<size_t> counts({1, 1});
     AllgathervOptions opts(context);
     opts.setOutput(output.getPointer(), counts);
+
+    // Run one operation first so we're measuring the operation timeout not
+    // connection timeout.
+    allgatherv(opts);
+
     opts.setTimeout(std::chrono::milliseconds(10));
     if (context->rank == 0) {
       try {

gloo/test/allreduce_test.cc

@@ -389,6 +389,11 @@ TEST_F(AllreduceNewTest, TestTimeout) {
     AllreduceOptions opts(context);
     opts.setOutputs(outputs.getPointers(), 1);
     opts.setReduceFunction(getFunction<uint64_t>());
+
+    // Run one operation first so we're measuring the operation timeout not
+    // connection timeout.
+    allreduce(opts);
+
     opts.setTimeout(std::chrono::milliseconds(10));
     if (context->rank == 0) {
       try {

gloo/test/barrier_test.cc

@@ -127,12 +127,18 @@ INSTANTIATE_TEST_CASE_P(
 TEST_F(BarrierNewTest, TestTimeout) {
   spawn(Transport::TCP, 2, [&](std::shared_ptr<Context> context) {
     BarrierOptions opts(context);
+
+    // Run barrier first so we're measuring the barrier timeout not connection
+    // timeout.
+    barrier(opts);
+
     opts.setTimeout(std::chrono::milliseconds(10));
     if (context->rank == 0) {
       try {
         barrier(opts);
         FAIL() << "Expected exception to be thrown";
       } catch (::gloo::IoException& e) {
+        std::cerr << e.what() << std::endl;
         ASSERT_NE(std::string(e.what()).find("Timed out"), std::string::npos);
       }
     }

gloo/test/base_test.cc

@@ -17,6 +17,7 @@ const char* kDefaultDevice = "localhost";
 // Transports that instantiated algorithms can be tested against.
 const std::vector<Transport> kTransportsForClassAlgorithms = {
     Transport::TCP,
+    Transport::TCP_LAZY,
 #if GLOO_HAVE_TRANSPORT_TCP_TLS
     Transport::TCP_TLS,
 #endif
@@ -27,6 +28,7 @@ const std::vector<Transport> kTransportsForClassAlgorithms = {
 // preferred over the instantiated style.
 const std::vector<Transport> kTransportsForFunctionAlgorithms = {
     Transport::TCP,
+    Transport::TCP_LAZY,
 #if GLOO_HAVE_TRANSPORT_TCP_TLS
     Transport::TCP_TLS,
 #endif
@@ -37,6 +39,8 @@ std::shared_ptr<::gloo::transport::Device> createDevice(Transport transport) {
 #if GLOO_HAVE_TRANSPORT_TCP
   if (transport == Transport::TCP) {
     return ::gloo::transport::tcp::CreateDevice(kDefaultDevice);
+  } else if (transport == Transport::TCP_LAZY) {
+    return ::gloo::transport::tcp::CreateLazyDevice(kDefaultDevice);
   }
 #endif
 #if GLOO_HAVE_TRANSPORT_TCP_TLS

gloo/test/base_test.h

@@ -59,6 +59,7 @@ class Barrier {
 
 enum Transport {
   TCP,
+  TCP_LAZY,
 #if GLOO_HAVE_TRANSPORT_TCP_TLS
   TCP_TLS,
 #endif

gloo/test/broadcast_test.cc

@@ -182,6 +182,11 @@ TEST_F(BroadcastTest, TestTimeout) {
     BroadcastOptions opts(context);
     opts.setOutput(output.getPointer(), 1);
     opts.setRoot(0);
+
+    // Run one operation first so we're measuring the operation timeout not
+    // connection timeout.
+    broadcast(opts);
+
     opts.setTimeout(std::chrono::milliseconds(10));
     if (context->rank == 0) {
       try {

gloo/test/gather_test.cc

@@ -76,6 +76,11 @@ TEST_F(GatherTest, TestTimeout) {
     opts.setInput(input.getPointer(), 1);
     opts.setOutput(output.getPointer(), context->size);
     opts.setRoot(0);
+
+    // Run one operation first so we're measuring the operation timeout not
+    // connection timeout.
+    gather(opts);
+
     opts.setTimeout(std::chrono::milliseconds(10));
     if (context->rank == 0) {
       try {

gloo/test/gatherv_test.cc

@@ -106,6 +106,11 @@ TEST_F(GathervTest, TestTimeout) {
     opts.setRoot(0);
     opts.setInput(input.getPointer(), 1);
     opts.setOutput(output.getPointer(), counts);
+
+    // Run one operation first so we're measuring the operation timeout not
+    // connection timeout.
+    gatherv(opts);
+
     opts.setTimeout(std::chrono::milliseconds(10));
     if (context->rank == 0) {
       try {

gloo/test/multiproc_test.h

@@ -102,11 +102,31 @@ class MultiProcWorker {
     auto device = createDevice(transport);
     context->setTimeout(std::chrono::milliseconds(kMultiProcTimeout));
     context->connectFullMesh(store_, device);
+
+    // Wait for all workers to be ready
+    ringBarrier(context);
+
     device.reset();
     sem_post(semaphore_);
     fn(std::move(context));
   }
 
+  void ringBarrier(std::shared_ptr<::gloo::rendezvous::Context>& context) {
+    int sendScratch = 0;
+    int recvScratch = 0;
+    auto sendBuf =
+        context->createUnboundBuffer(&sendScratch, sizeof(sendScratch));
+    auto recvBuf =
+        context->createUnboundBuffer(&recvScratch, sizeof(recvScratch));
+    const auto leftRank = (context->size + context->rank - 1) % context->size;
+    const auto rightRank = (context->rank + 1) % context->size;
+
+    sendBuf->send(leftRank, 0);
+    recvBuf->recv(rightRank, 0);
+    sendBuf->waitSend();
+    recvBuf->waitRecv();
+  }
+
  protected:
   std::shared_ptr<::gloo::rendezvous::Store> store_;
   sem_t* semaphore_;

gloo/test/reduce_test.cc

@@ -101,6 +101,11 @@ TEST_F(ReduceTest, TestTimeout) {
     opts.setOutput(outputs.getPointer(), 1);
     opts.setRoot(0);
     opts.setReduceFunction(getFunction<uint64_t>());
+
+    // Run one operation first so we're measuring the operation timeout not
+    // connection timeout.
+    reduce(opts);
+
     opts.setTimeout(std::chrono::milliseconds(10));
     if (context->rank == 0) {
       try {

gloo/test/scatter_test.cc

@@ -72,6 +72,11 @@ TEST_F(ScatterTest, TestTimeout) {
     opts.setInputs(input.getPointers(), 1);
     opts.setOutput(output.getPointer(), 1);
     opts.setRoot(0);
+
+    // Run one operation first so we're measuring the operation timeout not
+    // connection timeout.
+    scatter(opts);
+
     opts.setTimeout(std::chrono::milliseconds(10));
     if (context->rank == 0) {
       try {

gloo/test/tcp_test.cc

@@ -8,7 +8,7 @@ namespace transport {
 namespace tcp {
 
 TEST(TcpTest, ConnectTimeout) {
-  auto loop = std::make_shared<Loop>();
+  Loop loop;
 
   std::mutex m;
   std::condition_variable cv;
@@ -25,7 +25,7 @@ TEST(TcpTest, ConnectTimeout) {
     EXPECT_TRUE(e);
     EXPECT_TRUE(dynamic_cast<const TimeoutError*>(&e));
   };
-  connectLoop(*loop, remote, 0, 5, timeout, std::move(fn));
+  connectLoop(loop, remote, 0, 5, timeout, std::move(fn));
 
   std::unique_lock<std::mutex> lock(m);
   cv.wait(lock, [&] { return done; });

gloo/transport/tcp/context.cc

@@ -27,9 +27,14 @@ 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)) {}
+    : ::gloo::transport::Context(rank, size), device_(std::move(device)) {
+  connecting_.resize(size);
+}
 
 Context::~Context() {
+  // We need to shutdown the loop thread prior to freeing the pairs.
+  device_->shutdown();
+
   // Pairs refer to device by raw pointer.
   // Ensure they are destructed before the device.
   pairs_.clear();
@@ -50,8 +55,6 @@ void Context::createAndConnectAllPairs(std::shared_ptr<IStore> store) {
   // it's not super straightforward so left for folks having more bandwidth
   // later on.
 
-  int localRank = 0;
-  bool localRankSet = false;
   auto localHostName = getHostname();
   bool useRankAsSeqNum = useRankAsSeqNumber();
 
@@ -61,7 +64,7 @@ void Context::createAndConnectAllPairs(std::shared_ptr<IStore> store) {
   std::vector<ssize_t> pairIdentifiers;
   for (int i = 0; i < size; i++) {
     const auto& pair = createPair(i, useRankAsSeqNum);
-    if (!useRankAsSeqNum) {
+    if (!useRankAsSeqNum && !device_->isLazyInit()) {
       // Need to preserve the order of the pair identifiers if we are not using
       // the rank as seq number
       pairIdentifiers.emplace_back(
@@ -88,67 +91,103 @@ void Context::createAndConnectAllPairs(std::shared_ptr<IStore> store) {
       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;
+  store_ = store;
+
+  if (!device_->isLazyInit()) {
+    int localRank = 0;
+    bool localRankSet = false;
+    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++);
+        return store->wait_get(keyStr, getTimeout());
+      });
     }
-  } else {
-    std::generate_n(std::back_inserter(remoteRankInfos), size, [&] {
-      const auto& keyStr = std::to_string(key++);
-      return store->wait_get(keyStr, getTimeout());
-    });
-  }
 
-  // Connect every pair
-  for (int i = 0; i < size; i++) {
-    if (i == rank) {
-      // at this point we have enumerated all the ranks located on this host
-      // up to the current rank, so the current `localRank` number is
-      // what we'll set to the pairs.
-      localRankSet = true;
-      // We are not going to connect self.
-      continue;
-    }
+    // Connect every pair
+    for (int i = 0; i < size; i++) {
+      if (i == rank) {
+        // at this point we have enumerated all the ranks located on this host
+        // up to the current rank, so the current `localRank` number is
+        // what we'll set to the pairs.
+        localRankSet = true;
+        // We are not going to connect self.
+        continue;
+      }
 
-    Rank remoteRankInfo(remoteRankInfos[i]);
+      Rank remoteRankInfo(remoteRankInfos[i]);
 
-    if (!localRankSet && remoteRankInfo.hostname == localHostName) {
-      ++localRank;
-    }
+      if (!localRankSet && remoteRankInfo.hostname == localHostName) {
+        ++localRank;
+      }
 
-    const auto& pair = getPair(i);
-    auto remoteDeviceAddr = Address(remoteRankInfo.addressBytes).getSockaddr();
-    auto remoteAddr = Address(
-        remoteDeviceAddr,
-        useRankAsSeqNum ? (sequence_number_t)rank
-                        : remoteRankInfo.pairIdentifiers[rank]);
-    pair->connect(remoteAddr.bytes());
-  }
+      const auto& pair = pairs_[i];
+      auto remoteDeviceAddr =
+          Address(remoteRankInfo.addressBytes).getSockaddr();
+      auto remoteAddr = Address(
+          remoteDeviceAddr,
+          useRankAsSeqNum ? (sequence_number_t)rank
+                          : remoteRankInfo.pairIdentifiers[rank]);
+      pair->connect(remoteAddr.bytes());
+      connecting_[i] = true;
+    }
 
-  // Set the local rank info for all mesh pairs involving current rank
-  for (int i = 0; i < size; i++) {
-    if (i == rank) {
-      continue;
+    // Set the local rank info for all mesh pairs involving current rank
+    for (int i = 0; i < size; i++) {
+      if (i == rank) {
+        continue;
+      }
+      const auto& pair = getPair(i);
+      pair->setLocalRank(localRank);
     }
-    const auto& pair = getPair(i);
-    pair->setLocalRank(localRank);
   }
 
   printConnectivityInfo();
 }
 
+std::unique_ptr<transport::Pair>& Context::getPair(int rank) {
+  auto& pair = pairs_[rank];
+
+  if (!store_) {
+    // Manual context creation without store to bootstrap.
+    return pair;
+  }
+
+  // don't connect to self
+  if (rank == this->rank) {
+    return pair;
+  }
+
+  if (!connecting_[rank]) {
+    connecting_[rank] = true;
+
+    const auto& keyStr = std::to_string(rank);
+    auto remoteRankInfoBytes = store_->wait_get(keyStr, getTimeout());
+
+    Rank remoteRankInfo(remoteRankInfoBytes);
+
+    auto remoteDeviceAddr = Address(remoteRankInfo.addressBytes).getSockaddr();
+    auto remoteAddr = Address(remoteDeviceAddr, this->rank);
+    pair->connect(remoteAddr.bytes());
+  }
+  return pair;
+}
+
 std::unique_ptr<transport::Pair>& Context::createPair(int rank) {
   pairs_[rank] = std::unique_ptr<transport::Pair>(
       new tcp::Pair(this, device_.get(), rank, getTimeout(), false));
@@ -219,6 +258,11 @@ void Context::recvFromAny(
     size_t offset,
     size_t nbytes,
     std::vector<int> srcRanks) {
+  // Ensure all connections are established.
+  for (auto rank : srcRanks) {
+    getPair(rank);
+  }
+
   for (;;) {
     // Find rank of pair we can attempt a recv from
     auto rank = recvFromAnyFindRank(buf, slot, offset, nbytes, srcRanks);