Flushing Proxy Channels at CPU side upon reaching the Inflight Request Limit

include/mscclpp/core.hpp

@@ -33,7 +33,7 @@ std::string version();
 /// Base class for bootstraps.
 class Bootstrap {
  public:
-  Bootstrap(){};
+  Bootstrap() {};
   virtual ~Bootstrap() = default;
   virtual int getRank() = 0;
   virtual int getNranks() = 0;
@@ -388,6 +388,11 @@ class Endpoint {
   /// @return The transport used.
   Transport transport();
 
+  /// Get the max inflight requests.
+  ///
+  /// @return max inflight requests.
+  int maxInflightRequests();
+
   /// Serialize the Endpoint object to a vector of characters.
   ///
   /// @return A vector of characters representing the serialized Endpoint object.
@@ -416,6 +421,10 @@ class Endpoint {
 /// Represents a connection between two processes.
 class Connection {
  public:
+  /// Constructor.
+  /// @param maxInflightRequests The maximum number of inflight requests.
+  Connection(int maxInflightRequests) : maxInflightRequests(maxInflightRequests) {};
+
   virtual ~Connection() = default;
 
   /// Write data from a source @ref RegisteredMemory to a destination @ref RegisteredMemory.
@@ -454,10 +463,13 @@ class Connection {
   /// @return name of @ref transport() -> @ref remoteTransport()
   std::string getTransportName();
 
+  int getMaxInflightRequest();
+
  protected:
   // Internal methods for getting implementation pointers.
   static std::shared_ptr<RegisteredMemory::Impl> getImpl(RegisteredMemory& memory);
   static std::shared_ptr<Endpoint::Impl> getImpl(Endpoint& memory);
+  int maxInflightRequests;
 };
 
 /// Used to configure an endpoint.
@@ -472,14 +484,13 @@ struct EndpointConfig {
   int ibMaxCqPollNum = DefaultMaxCqPollNum;
   int ibMaxSendWr = DefaultMaxSendWr;
   int ibMaxWrPerSend = DefaultMaxWrPerSend;
-
-  /// Default constructor. Sets transport to Transport::Unknown.
-  EndpointConfig() : transport(Transport::Unknown) {}
+  int maxInflightRequests;
 
   /// Constructor that takes a transport and sets the other fields to their default values.
   ///
   /// @param transport The transport to use.
-  EndpointConfig(Transport transport) : transport(transport) {}
+  EndpointConfig(Transport transport = Transport::Unknown, int maxInflightRequests = -1)
+      : transport(transport), maxInflightRequests(maxInflightRequests) {}
 };
 
 /// Represents a context for communication. This provides a low-level interface for forming connections in use-cases

include/mscclpp/proxy_channel.hpp

@@ -72,6 +72,7 @@ class ProxyService : public BaseProxyService {
   std::vector<RegisteredMemory> memories_;
   std::shared_ptr<Proxy> proxy_;
   int deviceNumaNode;
+  std::unordered_map<std::shared_ptr<Connection>, int> inflightRequests;
 
   void bindThread();
 

src/connection.cc

@@ -36,10 +36,13 @@ std::string Connection::getTransportName() {
          TransportNames[static_cast<int>(this->remoteTransport())];
 }
 
+int Connection::getMaxInflightRequest() { return maxInflightRequests; }
+
 // CudaIpcConnection
 
 CudaIpcConnection::CudaIpcConnection(Endpoint localEndpoint, Endpoint remoteEndpoint, cudaStream_t stream)
-    : stream_(stream) {
+    : Connection(localEndpoint.maxInflightRequests() != -1 ? localEndpoint.maxInflightRequests() : INT_MAX),
+      stream_(stream) {
   if (localEndpoint.transport() != Transport::CudaIpc) {
     throw mscclpp::Error("Cuda IPC connection can only be made from a Cuda IPC endpoint", ErrorCode::InvalidUsage);
   }
@@ -119,7 +122,9 @@ void CudaIpcConnection::flush(int64_t timeoutUsec) {
 // IBConnection
 
 IBConnection::IBConnection(Endpoint localEndpoint, Endpoint remoteEndpoint, Context& context)
-    : transport_(localEndpoint.transport()),
+    : Connection(localEndpoint.maxInflightRequests() != -1 ? localEndpoint.maxInflightRequests()
+                                                           : EndpointConfig::DefaultMaxCqPollNum),
+      transport_(localEndpoint.transport()),
       remoteTransport_(remoteEndpoint.transport()),
       dummyAtomicSource_(std::make_unique<uint64_t>(0)) {
   qp = getImpl(localEndpoint)->ibQp_;
@@ -231,7 +236,10 @@ void IBConnection::flush(int64_t timeoutUsec) {
 
 EthernetConnection::EthernetConnection(Endpoint localEndpoint, Endpoint remoteEndpoint, uint64_t sendBufferSize,
                                        uint64_t recvBufferSize)
-    : abortFlag_(0), sendBufferSize_(sendBufferSize), recvBufferSize_(recvBufferSize) {
+    : Connection(localEndpoint.maxInflightRequests() != -1 ? localEndpoint.maxInflightRequests() : INT_MAX),
+      abortFlag_(0),
+      sendBufferSize_(sendBufferSize),
+      recvBufferSize_(recvBufferSize) {
   // Validating Transport Protocol
   if (localEndpoint.transport() != Transport::Ethernet || remoteEndpoint.transport() != Transport::Ethernet) {
     throw mscclpp::Error("Ethernet connection can only be made from Ethernet endpoints", ErrorCode::InvalidUsage);

src/endpoint.cc

@@ -13,7 +13,7 @@
 namespace mscclpp {
 
 Endpoint::Impl::Impl(EndpointConfig config, Context::Impl& contextImpl)
-    : transport_(config.transport), hostHash_(getHostHash()) {
+    : transport_(config.transport), hostHash_(getHostHash()), maxInflightRequests_(config.maxInflightRequests) {
   if (AllIBTransports.has(transport_)) {
     ibLocal_ = true;
     ibQp_ = contextImpl.getIbContext(transport_)
@@ -34,6 +34,8 @@ Endpoint::Impl::Impl(EndpointConfig config, Context::Impl& contextImpl)
 
 MSCCLPP_API_CPP Transport Endpoint::transport() { return pimpl_->transport_; }
 
+MSCCLPP_API_CPP int Endpoint::maxInflightRequests() { return pimpl_->maxInflightRequests_; }
+
 MSCCLPP_API_CPP std::vector<char> Endpoint::serialize() {
   std::vector<char> data;
   std::copy_n(reinterpret_cast<char*>(&pimpl_->transport_), sizeof(pimpl_->transport_), std::back_inserter(data));

src/include/endpoint.hpp

@@ -20,6 +20,7 @@ struct Endpoint::Impl {
 
   Transport transport_;
   uint64_t hostHash_;
+  int maxInflightRequests_;
 
   // The following are only used for IB and are undefined for other transports.
   bool ibLocal_;

src/proxy_channel.cc

@@ -76,15 +76,19 @@ ProxyHandlerResult ProxyService::handleTrigger(ProxyTrigger triggerRaw) {
     RegisteredMemory& src = memories_[trigger->fields.srcMemoryId];
     semaphore->connection()->write(dst, trigger->fields.dstOffset, src, trigger->fields.srcOffset,
                                    trigger->fields.size);
+    inflightRequests[semaphore->connection()]++;
   }
 
   if (trigger->fields.type & TriggerFlag) {
     semaphore->signal();
+    inflightRequests[semaphore->connection()]++;
   }
 
-  if (trigger->fields.type & TriggerSync) {
+  if (trigger->fields.type & TriggerSync ||
+      inflightRequests[semaphore->connection()] > semaphore->connection()->getMaxInflightRequest()) {
     semaphore->connection()->flush();
     result = ProxyHandlerResult::FlushFifoTailAndContinue;
+    inflightRequests[semaphore->connection()] = 0;
   }
 
   return result;