Merge branch 'main' into caiorocha/fixing_bug_constoffset_executor

microsoft · Nov 11, 2024 · f0288f9 · f0288f9
2 parents 496ac4e + 85fdde7
commit f0288f9
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/src/context.cc b/src/context.cc
@@ -11,7 +11,7 @@
 
 namespace mscclpp {
 
-Context::Impl::Impl() : ipcStream_(cudaStreamNonBlocking) {}
+Context::Impl::Impl() {}
 
 IbCtx* Context::Impl::getIbContext(Transport ibTransport) {
   // Find IB context or create it
@@ -43,7 +43,10 @@ MSCCLPP_API_CPP std::shared_ptr<Connection> Context::connect(Endpoint localEndpo
     if (remoteEndpoint.transport() != Transport::CudaIpc) {
       throw mscclpp::Error("Local transport is CudaIpc but remote is not", ErrorCode::InvalidUsage);
     }
-    conn = std::make_shared<CudaIpcConnection>(localEndpoint, remoteEndpoint, pimpl_->ipcStream_);
+    if (!(pimpl_->ipcStream_)) {
+      pimpl_->ipcStream_ = std::make_shared<CudaStreamWithFlags>(cudaStreamNonBlocking);
+    }
+    conn = std::make_shared<CudaIpcConnection>(localEndpoint, remoteEndpoint, cudaStream_t(*(pimpl_->ipcStream_)));
   } else if (AllIBTransports.has(localEndpoint.transport())) {
     if (!AllIBTransports.has(remoteEndpoint.transport())) {
       throw mscclpp::Error("Local transport is IB but remote is not", ErrorCode::InvalidUsage);

diff --git a/src/include/context.hpp b/src/include/context.hpp
@@ -16,7 +16,7 @@ namespace mscclpp {
 struct Context::Impl {
   std::vector<std::shared_ptr<Connection>> connections_;
   std::unordered_map<Transport, std::unique_ptr<IbCtx>> ibContexts_;
-  CudaStreamWithFlags ipcStream_;
+  std::shared_ptr<CudaStreamWithFlags> ipcStream_;
   CUmemGenericAllocationHandle mcHandle_;
 
   Impl();