How to handle binding from c10::cuda::CUDAStream to torch.cuda.Stream

username: shadowshadow
cooked: <p>I want to subclass ProcessGroupNCCL to enable communication using a user-specified CUDA stream.</p>
<pre data-code-wrap="cpp"><code class="lang-cpp">class CustomProcessGroupNCCL: public ProcessGroupNCCL {
public:
using ProcessGroupNCCL::ProcessGroupNCCL;
  
// Set NCCL stream by device index
void setNcclStreamByDevice(int device_index, const c10::cuda::CUDAStream&amp; stream) {
  std::string key = std::to_string(device_index);
  std::lock_guard&lt;std::mutex&gt; guard(mutex_);
  if (ncclStreams_.find(key) != ncclStreams_.end()) {
    ncclStreams_.at(key) = stream;
  } else {
    TORCH_CHECK(false, "Key not found in ncclStreams_: ", key);
  }
}
  
// Get NCCL stream by device index
c10::cuda::CUDAStream getNcclStreamByDevice(int device_index) {
  std::string key = std::to_string(device_index);
  std::lock_guard&lt;std::mutex&gt; guard(mutex_);
  auto it = ncclStreams_.find(key);
  TORCH_CHECK(it != ncclStreams_.end(), "Key not found in ncclStreams_: ", key);
  return it-&gt;second;
}
  
// List all available NCCL stream keys
std::vector&lt;std::string&gt; getAvailableStreamKeys() {
  std::lock_guard&lt;std::mutex&gt; guard(mutex_);
  std::vector&lt;std::string&gt; keys;
  for (const auto&amp; pair : ncclStreams_) {
    keys.push_back(pair.first);
  }
  return keys;
}

// List all available NCCL streams 
std::unordered_map&lt;std::string, c10::cuda::CUDAStream&gt; getAvailableStreams(){
  return ncclStreams_;
} 
  
protected:
std::mutex mutex_;
};

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.doc() = "Custom ProcessGroupNCCL with stream management capabilities";
  
// Expose the factory function
m.def("create_custom_nccl_backend", &amp;c10d::createCustomNCCLBackend,
      "Create a custom NCCL backend with stream management",
      py::arg("store"), py::arg("rank"), py::arg("size"), py::arg("options"));
  
// Expose the CustomProcessGroupNCCL class
py::class_&lt;c10d::CustomProcessGroupNCCL, c10d::ProcessGroupNCCL, c10::intrusive_ptr&lt;c10d::CustomProcessGroupNCCL&gt;&gt;(m, "CustomProcessGroupNCCL")
  .def("set_nccl_stream_by_device", &amp;c10d::CustomProcessGroupNCCL::setNcclStreamByDevice,
       "Set NCCL stream by device index",
       py::arg("device_index"), py::arg("stream"))
  .def("get_nccl_stream_by_device", &amp;c10d::CustomProcessGroupNCCL::getNcclStreamByDevice,
       "Get NCCL stream by device index",
       py::arg("device_index"))
  .def("get_available_stream_keys", &amp;c10d::CustomProcessGroupNCCL::getAvailableStreamKeys,
       "List all available NCCL stream keys")
  .def("get_available_streams", &amp;c10d::CustomProcessGroupNCCL::getAvailableStreams);
}

</code></pre>
<p>I write a simple python test</p>
<pre data-code-wrap="python"><code class="lang-python">import custom_nccl_backend

def setup_custom_nccl_backend():
  """
  Register the custom NCCL backend with PyTorch distributed.
  """
  def custom_nccl_backend_creator(store, rank, size, *args):
      # Create default options for ProcessGroupNCCL
      options = torch.distributed.ProcessGroupNCCL.Options()
      # You can customize options here if needed
      # options.timeout = timedelta(seconds=300)
      # options.is_high_priority_stream = False
      
      # Use our custom factory function
      print("register args:", store, rank, size, *args)
      return custom_nccl_backend.create_custom_nccl_backend(
          store, rank, size, options
      )
  
  # Register the custom backend
  dist.Backend.register_backend("CUSTOM_NCCL", custom_nccl_backend_creator
# main
dist.init_process_group(
      backend="cuda:custom_nccl", 
      rank=rank, 
      world_size=world_size
  )
  
# Get the process group
pg = dist.group.WORLD
backend = pg._get_backend(torch.device(f"cuda:{rank}"))
# Create a custom CUDA stream
device = torch.device(f"cuda:{rank}")
torch.cuda.set_device(device)
custom_stream = torch.cuda.Stream(device=device)
print(f"[Rank {rank}] Created custom stream: {custom_stream}")
  
# Set the custom stream for this device
backend.set_nccl_stream_by_device(rank, custom_stream)

</code></pre>
<p>It seems that <code>c10::cuda::CUDAStream</code> cannot be converted to and from Python’s <code>torch.cuda.Stream</code> . How should I handle this?</p>
<pre data-code-wrap="python"><code class="lang-python">-- Process 0 terminated with the following error:
Traceback (most recent call last):
  File "/usr/local/lib/python3.12/dist-packages/torch/multiprocessing/spawn.py", line 90, in _wrap
    fn(i, *args)
  File "/mnt/code/wangtianxing.wtx/code/streamer-ttv-7b/csrc/ProcessGroupNCCLCustom/example_usage.py", line 84, in example_with_custom_streams
    backend.set_nccl_stream_by_device(rank, custom_stream)
TypeError: set_nccl_stream_by_device(): incompatible function arguments. The following argument types are supported:
    1. (self: custom_nccl_backend.CustomProcessGroupNCCL, device_index: int, stream: c10::cuda::CUDAStream) -&gt; None

Invoked with: &lt;custom_nccl_backend.CustomProcessGroupNCCL object at 0x7f1f528a39b0&gt;, 0, &lt;torch.cuda.Stream device=cuda:0 cuda_stream=0xa43a250&gt;
</code></pre>

username: albanD
cooked: <p><a class="mention" href="/u/aidyn-a">@Aidyn-A</a> answer on github:</p>
<p>That should be pretty easy to do.</p>
<p>Write a helper function in python side, in which you will need to pass <code>stream_id</code>, <code>device_index</code> and <code>device_type</code> of <code>torch.cuda.Stream</code> object, as it was done here:</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/80bf883d21ba48898e5fd40a6cbf418fa3ff3f33/torch/cuda/__init__.py#L729-L733">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/80bf883d21ba48898e5fd40a6cbf418fa3ff3f33/torch/cuda/__init__.py#L729-L733" target="_blank" rel="noopener">github.com/pytorch/pytorch</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/80bf883d21ba48898e5fd40a6cbf418fa3ff3f33/torch/cuda/__init__.py#L729-L733" target="_blank" rel="noopener">torch/cuda/__init__.py</a></h4>

<div class="git-blob-info">
  <a href="https://github.com/pytorch/pytorch/blob/80bf883d21ba48898e5fd40a6cbf418fa3ff3f33/torch/cuda/__init__.py#L729-L733" rel="noopener"><code>80bf883d2</code></a>
</div>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="729" style="counter-reset: li-counter 728 ;">
          <li>_set_stream_by_id(</li>
          <li>    stream_id=stream.stream_id,</li>
          <li>    device_index=stream.device_index,</li>
          <li>    device_type=stream.device_type,</li>
          <li>)</li>
      </ol>
    </code></pre>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>And on C++ side utilize the converter to get an object of type <code>c10::cuda::CUDAStream</code> like it is done here:</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/fca2601c9d628e1bd2d75c7318cd22c4e8c832aa/torch/csrc/cuda/Module.cpp#L242-L245">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/fca2601c9d628e1bd2d75c7318cd22c4e8c832aa/torch/csrc/cuda/Module.cpp#L242-L245" target="_blank" rel="noopener">github.com/pytorch/pytorch</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/fca2601c9d628e1bd2d75c7318cd22c4e8c832aa/torch/csrc/cuda/Module.cpp#L242-L245" target="_blank" rel="noopener">torch/csrc/cuda/Module.cpp</a></h4>

<div class="git-blob-info">
  <a href="https://github.com/pytorch/pytorch/blob/fca2601c9d628e1bd2d75c7318cd22c4e8c832aa/torch/csrc/cuda/Module.cpp#L242-L245" rel="noopener"><code>fca2601c9</code></a>
</div>



    <pre class="onebox"><code class="lang-cpp">
      <ol class="start lines" start="242" style="counter-reset: li-counter 241 ;">
          <li>auto stream = at::cuda::CUDAStream::unpack3(</li>
          <li>    stream_id,</li>
          <li>    static_cast&lt;c10::DeviceIndex&gt;(device_index),</li>
          <li>    static_cast&lt;c10::DeviceType&gt;(device_type));</li>
      </ol>
    </code></pre>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>