Asynchronous setup (#514)

Cherry-picked a part of features from #167: now `Communicator::setup()`
is unneeded. `Communicator::sendMemory()` conducts the task inline, and
`Communicator::recvMemory()` and `Communicator::connect()` conducts the
task asynchronously without explicit setup.

Author: chhwang

include/mscclpp/core.hpp

@@ -557,64 +557,19 @@ class Context {
   friend class Endpoint;
 };
 
-/// A base class for objects that can be set up during @ref Communicator::setup().
-struct Setuppable {
-  virtual ~Setuppable() = default;
-
-  /// Called inside @ref Communicator::setup() before any call to @ref endSetup() of any @ref Setuppable object that is
-  /// being set up within the same @ref Communicator::setup() call.
-  ///
-  /// @param bootstrap A shared pointer to the bootstrap implementation.
-  virtual void beginSetup(std::shared_ptr<Bootstrap> bootstrap);
-
-  /// Called inside @ref Communicator::setup() after all calls to @ref beginSetup() of all @ref Setuppable objects that
-  /// are being set up within the same @ref Communicator::setup() call.
-  ///
-  /// @param bootstrap A shared pointer to the bootstrap implementation.
-  virtual void endSetup(std::shared_ptr<Bootstrap> bootstrap);
-};
-
-/// A non-blocking future that can be used to check if a value is ready and retrieve it.
 template <typename T>
-class NonblockingFuture {
-  std::shared_future<T> future;
-
- public:
-  /// Default constructor.
-  NonblockingFuture() = default;
-
-  /// Constructor that takes a shared future and moves it into the NonblockingFuture.
-  ///
-  /// @param future The shared future to move.
-  NonblockingFuture(std::shared_future<T>&& future) : future(std::move(future)) {}
-
-  /// Check if the value is ready to be retrieved.
-  ///
-  /// @return True if the value is ready, false otherwise.
-  bool ready() const { return future.wait_for(std::chrono::seconds(0)) == std::future_status::ready; }
-
-  /// Get the value.
-  ///
-  /// @return The value.
-  ///
-  /// @throws Error if the value is not ready.
-  T get() const {
-    if (!ready()) throw Error("NonblockingFuture::get() called before ready", ErrorCode::InvalidUsage);
-    return future.get();
-  }
-};
+using NonblockingFuture [[deprecated("Use std::shared_future instead. This will be removed in a future release.")]] =
+    std::shared_future<T>;
 
 /// A class that sets up all registered memories and connections between processes.
 ///
 /// A typical way to use this class:
-///   1. Call @ref connectOnSetup() to declare connections between the calling process with other processes.
-///   2. Call @ref registerMemory() to register memory regions that will be used for communication.
-///   3. Call @ref sendMemoryOnSetup() or @ref recvMemoryOnSetup() to send/receive registered memory regions to/from
+///   1. Call connect() to declare connections between the calling process with other processes.
+///   2. Call registerMemory() to register memory regions that will be used for communication.
+///   3. Call sendMemory() or recvMemory() to send/receive registered memory regions to/from
 ///      other processes.
-///   4. Call @ref setup() to set up all registered memories and connections declared in the previous steps.
-///   5. Call @ref NonblockingFuture<RegisteredMemory>::get() to get the registered memory regions received from other
-///      processes.
-///   6. All done; use connections and registered memories to build channels.
+///   4. Call get() on all futures returned by connect() and recvMemory().
+///   5. All done; use connections and registered memories to build channels.
 ///
 class Communicator {
  public:
@@ -645,40 +600,57 @@ class Communicator {
   /// @return RegisteredMemory A handle to the buffer.
   RegisteredMemory registerMemory(void* ptr, size_t size, TransportFlags transports);
 
-  /// Send information of a registered memory to the remote side on setup.
+  /// Send information of a registered memory to the remote side.
   ///
-  /// This function registers a send to a remote process that will happen by a following call of @ref setup(). The send
-  /// will carry information about a registered memory on the local process.
+  /// The send will be performed immediately upon calling this function.
   ///
   /// @param memory The registered memory buffer to send information about.
   /// @param remoteRank The rank of the remote process.
   /// @param tag The tag to use for identifying the send.
-  void sendMemoryOnSetup(RegisteredMemory memory, int remoteRank, int tag);
+  void sendMemory(RegisteredMemory memory, int remoteRank, int tag);
 
-  /// Receive memory on setup.
+  [[deprecated("Use sendMemory() instead. This will be removed in a future release.")]] void sendMemoryOnSetup(
+      RegisteredMemory memory, int remoteRank, int tag) {
+    sendMemory(memory, remoteRank, tag);
+  }
+
+  /// Receive memory information from a corresponding sendMemory call on the remote side.
   ///
-  /// This function registers a receive from a remote process that will happen by a following call of @ref setup(). The
-  /// receive will carry information about a registered memory on the remote process.
+  /// This function returns a future immediately. The actual receive will be performed upon calling
+  /// the first get() on the future.
   ///
   /// @param remoteRank The rank of the remote process.
   /// @param tag The tag to use for identifying the receive.
-  /// @return NonblockingFuture<RegisteredMemory> A non-blocking future of registered memory.
-  NonblockingFuture<RegisteredMemory> recvMemoryOnSetup(int remoteRank, int tag);
+  /// @return std::shared_future<RegisteredMemory> A non-blocking future of registered memory.
+  std::shared_future<RegisteredMemory> recvMemory(int remoteRank, int tag);
 
-  /// Connect to a remote rank on setup.
+  [[deprecated(
+      "Use recvMemory() instead. This will be removed in a future release.")]] NonblockingFuture<RegisteredMemory>
+  recvMemoryOnSetup(int remoteRank, int tag) {
+    return recvMemory(remoteRank, tag);
+  }
+
+  /// Connect to a remote rank.
   ///
-  /// This function only prepares metadata for connection. The actual connection is made by a following call of
-  /// @ref setup(). Note that this function is two-way and a connection from rank `i` to remote rank `j` needs
+  /// This function will immediately send metadata about the local endpoint to the remote rank, and return a future
+  /// without waiting for the remote rank to respond. The connection will be established when the remote rank
+  /// responds with its own endpoint and the local rank calls the first get() on the future.
+  /// Note that this function is two-way and a connection from rank `i` to remote rank `j` needs
   /// to have a counterpart from rank `j` to rank `i`. Note that with IB, buffers are registered at a page level and if
   /// a buffer is spread through multiple pages and do not fully utilize all of them, IB's QP has to register for all
   /// involved pages. This potentially has security risks if the connection's accesses are given to a malicious process.
   ///
   /// @param remoteRank The rank of the remote process.
   /// @param tag The tag of the connection for identifying it.
   /// @param config The configuration for the local endpoint.
-  /// @return NonblockingFuture<NonblockingFuture<std::shared_ptr<Connection>>> A non-blocking future of shared pointer
-  /// to the connection.
-  NonblockingFuture<std::shared_ptr<Connection>> connectOnSetup(int remoteRank, int tag, EndpointConfig localConfig);
+  /// @return std::shared_future<std::shared_ptr<Connection>> A non-blocking future of shared pointer to the connection.
+  std::shared_future<std::shared_ptr<Connection>> connect(int remoteRank, int tag, EndpointConfig localConfig);
+
+  [[deprecated("Use connect() instead. This will be removed in a future release.")]] NonblockingFuture<
+      std::shared_ptr<Connection>>
+  connectOnSetup(int remoteRank, int tag, EndpointConfig localConfig) {
+    return connect(remoteRank, tag, localConfig);
+  }
 
   /// Get the remote rank a connection is connected to.
   ///
@@ -692,17 +664,7 @@ class Communicator {
   /// @return The tag the connection was made with.
   int tagOf(const Connection& connection);
 
-  /// Add a custom Setuppable object to a list of objects to be setup later, when @ref setup() is called.
-  ///
-  /// @param setuppable A shared pointer to the Setuppable object.
-  void onSetup(std::shared_ptr<Setuppable> setuppable);
-
-  /// Setup all objects that have registered for setup.
-  ///
-  /// This includes previous calls of @ref sendMemoryOnSetup(), @ref recvMemoryOnSetup(), @ref connectOnSetup(), and
-  /// @ref onSetup(). It is allowed to call this function multiple times, where the n-th call will only setup objects
-  /// that have been registered after the (n-1)-th call.
-  void setup();
+  [[deprecated("setup() is now no-op and no longer needed. This will be removed in a future release.")]] void setup() {}
 
  private:
   // The interal implementation.

include/mscclpp/semaphore.hpp

@@ -30,7 +30,7 @@ template <template <typename> typename InboundDeleter, template <typename> typen
 class BaseSemaphore {
  protected:
   /// The registered memory for the remote peer's inbound semaphore ID.
-  NonblockingFuture<RegisteredMemory> remoteInboundSemaphoreIdsRegMem_;
+  std::shared_future<RegisteredMemory> remoteInboundSemaphoreIdsRegMem_;
 
   /// The inbound semaphore ID that is incremented by the remote peer and waited on by the local peer.
   ///

python/mscclpp/core_py.cpp

@@ -29,9 +29,7 @@ extern void register_gpu_utils(nb::module_& m);
 template <typename T>
 void def_nonblocking_future(nb::handle& m, const std::string& typestr) {
   std::string pyclass_name = std::string("NonblockingFuture") + typestr;
-  nb::class_<NonblockingFuture<T>>(m, pyclass_name.c_str())
-      .def("ready", &NonblockingFuture<T>::ready)
-      .def("get", &NonblockingFuture<T>::get);
+  nb::class_<NonblockingFuture<T>>(m, pyclass_name.c_str()).def("get", &NonblockingFuture<T>::get);
 }
 
 void register_core(nb::module_& m) {

src/communicator.cc

@@ -30,79 +30,31 @@ MSCCLPP_API_CPP RegisteredMemory Communicator::registerMemory(void* ptr, size_t
   return context()->registerMemory(ptr, size, transports);
 }
 
-struct MemorySender : public Setuppable {
-  MemorySender(RegisteredMemory memory, int remoteRank, int tag)
-      : memory_(memory), remoteRank_(remoteRank), tag_(tag) {}
-
-  void beginSetup(std::shared_ptr<Bootstrap> bootstrap) override {
-    bootstrap->send(memory_.serialize(), remoteRank_, tag_);
-  }
-
-  RegisteredMemory memory_;
-  int remoteRank_;
-  int tag_;
-};
-
-MSCCLPP_API_CPP void Communicator::sendMemoryOnSetup(RegisteredMemory memory, int remoteRank, int tag) {
-  onSetup(std::make_shared<MemorySender>(memory, remoteRank, tag));
+MSCCLPP_API_CPP void Communicator::sendMemory(RegisteredMemory memory, int remoteRank, int tag) {
+  pimpl_->bootstrap_->send(memory.serialize(), remoteRank, tag);
 }
 
-struct MemoryReceiver : public Setuppable {
-  MemoryReceiver(int remoteRank, int tag) : remoteRank_(remoteRank), tag_(tag) {}
-
-  void endSetup(std::shared_ptr<Bootstrap> bootstrap) override {
+MSCCLPP_API_CPP std::shared_future<RegisteredMemory> Communicator::recvMemory(int remoteRank, int tag) {
+  return std::async(std::launch::deferred, [this, remoteRank, tag]() {
     std::vector<char> data;
-    bootstrap->recv(data, remoteRank_, tag_);
-    memoryPromise_.set_value(RegisteredMemory::deserialize(data));
-  }
-
-  std::promise<RegisteredMemory> memoryPromise_;
-  int remoteRank_;
-  int tag_;
-};
-
-MSCCLPP_API_CPP NonblockingFuture<RegisteredMemory> Communicator::recvMemoryOnSetup(int remoteRank, int tag) {
-  auto memoryReceiver = std::make_shared<MemoryReceiver>(remoteRank, tag);
-  onSetup(memoryReceiver);
-  return NonblockingFuture<RegisteredMemory>(memoryReceiver->memoryPromise_.get_future());
+    bootstrap()->recv(data, remoteRank, tag);
+    return RegisteredMemory::deserialize(data);
+  });
 }
 
-struct Communicator::Impl::Connector : public Setuppable {
-  Connector(Communicator& comm, Communicator::Impl& commImpl_, int remoteRank, int tag, EndpointConfig localConfig)
-      : comm_(comm),
-        commImpl_(commImpl_),
-        remoteRank_(remoteRank),
-        tag_(tag),
-        localEndpoint_(comm.context()->createEndpoint(localConfig)) {}
-
-  void beginSetup(std::shared_ptr<Bootstrap> bootstrap) override {
-    bootstrap->send(localEndpoint_.serialize(), remoteRank_, tag_);
-  }
+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);
 
-  void endSetup(std::shared_ptr<Bootstrap> bootstrap) override {
+  return std::async(std::launch::deferred, [this, remoteRank, tag, localEndpoint = std::move(localEndpoint)]() mutable {
     std::vector<char> data;
-    bootstrap->recv(data, remoteRank_, tag_);
+    bootstrap()->recv(data, remoteRank, tag);
     auto remoteEndpoint = Endpoint::deserialize(data);
-    auto connection = comm_.context()->connect(localEndpoint_, remoteEndpoint);
-    commImpl_.connectionInfos_[connection.get()] = {remoteRank_, tag_};
-    connectionPromise_.set_value(connection);
-    INFO(MSCCLPP_INIT, "Connection %d -> %d created (%s)", comm_.bootstrap()->getRank(), remoteRank_,
-         connection->getTransportName().c_str());
-  }
-
-  std::promise<std::shared_ptr<Connection>> connectionPromise_;
-  Communicator& comm_;
-  Communicator::Impl& commImpl_;
-  int remoteRank_;
-  int tag_;
-  Endpoint localEndpoint_;
-};
-
-MSCCLPP_API_CPP NonblockingFuture<std::shared_ptr<Connection>> Communicator::connectOnSetup(
-    int remoteRank, int tag, EndpointConfig localConfig) {
-  auto connector = std::make_shared<Communicator::Impl::Connector>(*this, *pimpl_, remoteRank, tag, localConfig);
-  onSetup(connector);
-  return NonblockingFuture<std::shared_ptr<Connection>>(connector->connectionPromise_.get_future());
+    auto connection = context()->connect(localEndpoint, remoteEndpoint);
+    pimpl_->connectionInfos_[connection.get()] = {remoteRank, tag};
+    return connection;
+  });
 }
 
 MSCCLPP_API_CPP int Communicator::remoteRankOf(const Connection& connection) {
@@ -113,18 +65,4 @@ MSCCLPP_API_CPP int Communicator::tagOf(const Connection& connection) {
   return pimpl_->connectionInfos_.at(&connection).tag;
 }
 
-MSCCLPP_API_CPP void Communicator::onSetup(std::shared_ptr<Setuppable> setuppable) {
-  pimpl_->toSetup_.push_back(setuppable);
-}
-
-MSCCLPP_API_CPP void Communicator::setup() {
-  for (auto& setuppable : pimpl_->toSetup_) {
-    setuppable->beginSetup(pimpl_->bootstrap_);
-  }
-  for (auto& setuppable : pimpl_->toSetup_) {
-    setuppable->endSetup(pimpl_->bootstrap_);
-  }
-  pimpl_->toSetup_.clear();
-}
-
 }  // namespace mscclpp

src/core.cc

@@ -89,10 +89,6 @@ const TransportFlags AllIBTransports = Transport::IB0 | Transport::IB1 | Transpo
 
 const TransportFlags AllTransports = AllIBTransports | Transport::CudaIpc | Transport::Ethernet;
 
-void Setuppable::beginSetup(std::shared_ptr<Bootstrap>) {}
-
-void Setuppable::endSetup(std::shared_ptr<Bootstrap>) {}
-
 }  // namespace mscclpp
 
 namespace std {

src/include/communicator.hpp

@@ -22,7 +22,6 @@ struct Communicator::Impl {
   std::shared_ptr<Bootstrap> bootstrap_;
   std::shared_ptr<Context> context_;
   std::unordered_map<const Connection*, ConnectionInfo> connectionInfos_;
-  std::vector<std::shared_ptr<Setuppable>> toSetup_;
 
   Impl(std::shared_ptr<Bootstrap> bootstrap, std::shared_ptr<Context> context);
 

test/unit/core_tests.cc

@@ -18,20 +18,6 @@ class LocalCommunicatorTest : public ::testing::Test {
   std::shared_ptr<mscclpp::Communicator> comm;
 };
 
-class MockSetuppable : public mscclpp::Setuppable {
- public:
-  MOCK_METHOD(void, beginSetup, (std::shared_ptr<mscclpp::Bootstrap> bootstrap), (override));
-  MOCK_METHOD(void, endSetup, (std::shared_ptr<mscclpp::Bootstrap> bootstrap), (override));
-};
-
-TEST_F(LocalCommunicatorTest, OnSetup) {
-  auto mockSetuppable = std::make_shared<MockSetuppable>();
-  comm->onSetup(mockSetuppable);
-  EXPECT_CALL(*mockSetuppable, beginSetup(std::dynamic_pointer_cast<mscclpp::Bootstrap>(bootstrap)));
-  EXPECT_CALL(*mockSetuppable, endSetup(std::dynamic_pointer_cast<mscclpp::Bootstrap>(bootstrap)));
-  comm->setup();
-}
-
 TEST_F(LocalCommunicatorTest, RegisterMemory) {
   int dummy[42];
   auto memory = comm->registerMemory(&dummy, sizeof(dummy), mscclpp::NoTransports);