Fix #514 & minor updates

Author: chhwang

python/mscclpp/comm.py

@@ -107,7 +107,7 @@ def make_connection(
         return connections
 
     def register_tensor_with_connections(
-        self, tensor: Type[cp.ndarray] or Type[np.ndarray], connections: dict[int, Connection]
+        self, tensor: Type[cp.ndarray] | Type[np.ndarray], connections: dict[int, Connection]
     ) -> dict[int, RegisteredMemory]:
         transport_flags = TransportFlags()
         for rank in connections:
@@ -135,7 +135,7 @@ def register_tensor_with_connections(
     def make_semaphore(
         self,
         connections: dict[int, Connection],
-        semaphore_type: Type[Host2HostSemaphore] or Type[Host2DeviceSemaphore] or Type[MemoryDevice2DeviceSemaphore],
+        semaphore_type: Type[Host2HostSemaphore] | Type[Host2DeviceSemaphore] | Type[MemoryDevice2DeviceSemaphore],
     ) -> dict[int, Host2HostSemaphore]:
         semaphores = {}
         for rank in connections:

src/bootstrap/bootstrap.cc

@@ -528,6 +528,7 @@ std::shared_ptr<Socket> TcpBootstrap::Impl::getPeerRecvSocket(int peer, int tag)
     if (recvPeer == peer && recvTag == tag) {
       return sock;
     }
+    // TODO(chhwang): set an exit condition or timeout
   }
 }
 

src/communicator.cc

@@ -31,30 +31,41 @@ MSCCLPP_API_CPP RegisteredMemory Communicator::registerMemory(void* ptr, size_t
 }
 
 MSCCLPP_API_CPP void Communicator::sendMemory(RegisteredMemory memory, int remoteRank, int tag) {
-  pimpl_->bootstrap_->send(memory.serialize(), remoteRank, tag);
+  bootstrap()->send(memory.serialize(), remoteRank, tag);
 }
 
 MSCCLPP_API_CPP std::shared_future<RegisteredMemory> Communicator::recvMemory(int remoteRank, int tag) {
-  return std::async(std::launch::deferred, [this, remoteRank, tag]() {
+  size_t numRecvItemsAhead = pimpl_->recvQueues_.getNumRecvItems(remoteRank, tag);
+  auto future = std::async(std::launch::deferred, [this, remoteRank, tag, numRecvItemsAhead]() {
+    pimpl_->recvQueues_.waitN(remoteRank, tag, numRecvItemsAhead);
     std::vector<char> data;
     bootstrap()->recv(data, remoteRank, tag);
     return RegisteredMemory::deserialize(data);
   });
+  auto shared_future = std::shared_future<RegisteredMemory>(std::move(future));
+  pimpl_->recvQueues_.addRecvItem(remoteRank, tag, shared_future);
+  return shared_future;
 }
 
 MSCCLPP_API_CPP std::shared_future<std::shared_ptr<Connection>> Communicator::connect(int remoteRank, int tag,
                                                                                       EndpointConfig localConfig) {
-  auto localEndpoint = pimpl_->context_->createEndpoint(localConfig);
-  pimpl_->bootstrap_->send(localEndpoint.serialize(), remoteRank, tag);
+  auto localEndpoint = context()->createEndpoint(localConfig);
+  bootstrap()->send(localEndpoint.serialize(), remoteRank, tag);
 
-  return std::async(std::launch::deferred, [this, remoteRank, tag, localEndpoint = std::move(localEndpoint)]() mutable {
+  size_t numRecvItemsAhead = pimpl_->recvQueues_.getNumRecvItems(remoteRank, tag);
+  auto future = std::async(std::launch::deferred, [this, remoteRank, tag, numRecvItemsAhead,
+                                                   localEndpoint = std::move(localEndpoint)]() mutable {
+    pimpl_->recvQueues_.waitN(remoteRank, tag, numRecvItemsAhead);
     std::vector<char> data;
     bootstrap()->recv(data, remoteRank, tag);
     auto remoteEndpoint = Endpoint::deserialize(data);
     auto connection = context()->connect(localEndpoint, remoteEndpoint);
     pimpl_->connectionInfos_[connection.get()] = {remoteRank, tag};
     return connection;
   });
+  auto shared_future = std::shared_future<std::shared_ptr<Connection>>(std::move(future));
+  pimpl_->recvQueues_.addRecvItem(remoteRank, tag, shared_future);
+  return shared_future;
 }
 
 MSCCLPP_API_CPP int Communicator::remoteRankOf(const Connection& connection) {

src/include/communicator.hpp

@@ -9,9 +9,49 @@
 #include <unordered_map>
 #include <vector>
 
+#include "utils_internal.hpp"
+
 namespace mscclpp {
 
-class ConnectionBase;
+class BaseRecvItem {
+ public:
+  virtual ~BaseRecvItem() = default;
+  virtual void wait() = 0;
+};
+
+template <typename T>
+class RecvItem : public BaseRecvItem {
+ public:
+  RecvItem(std::shared_future<T> future) : future_(future) {}
+
+  void wait() { future_.wait(); }
+
+ private:
+  std::shared_future<T> future_;
+};
+
+class RecvQueues {
+ public:
+  RecvQueues() = default;
+
+  template <typename T>
+  void addRecvItem(int remoteRank, int tag, std::shared_future<T> future) {
+    auto& queue = queues_[std::make_pair(remoteRank, tag)];
+    queue.emplace_back(std::make_shared<RecvItem<T>>(future));
+  }
+
+  size_t getNumRecvItems(int remoteRank, int tag) { return queues_[std::make_pair(remoteRank, tag)].size(); }
+
+  void waitN(int remoteRank, int tag, size_t n) {
+    auto& queue = queues_[std::make_pair(remoteRank, tag)];
+    for (size_t i = 0; i < n; ++i) {
+      queue[i]->wait();
+    }
+  }
+
+ private:
+  std::unordered_map<std::pair<int, int>, std::vector<std::shared_ptr<BaseRecvItem>>, PairHash> queues_;
+};
 
 struct ConnectionInfo {
   int remoteRank;
@@ -22,6 +62,7 @@ struct Communicator::Impl {
   std::shared_ptr<Bootstrap> bootstrap_;
   std::shared_ptr<Context> context_;
   std::unordered_map<const Connection*, ConnectionInfo> connectionInfos_;
+  RecvQueues recvQueues_;
 
   Impl(std::shared_ptr<Bootstrap> bootstrap, std::shared_ptr<Context> context);