chore: import upstream snapshot with attribution

python/ray/util/collective/collective_group/cuda_stream.py

@@ -0,0 +1,95 @@
+import logging
+import threading
+
+import cupy
+
+from ray.util.collective.collective_group import nccl_util
+from ray.util.collective.const import ENV
+
+NCCL_STREAM_POOL_SIZE = 32
+MAX_GPU_PER_ACTOR = 16
+
+logger = logging.getLogger(__name__)
+
+
+class StreamPool:
+    """The class that represents a stream pool associated with a GPU.
+
+    When multistream is enabled, we will allocate a pool of streams for each
+    GPU, and get available stream from this pool when a collective kernel is
+    initialized. This enables overlapping computation/communication kernels
+    using multiple CUDA streams, given that the streams a appropriately
+    synchronized. The class is thread-safe.
+
+
+    Args:
+        device_idx: the absolute index of the device for this pool.
+    """
+
+    def __init__(self, device_idx: int):
+        self.device_idx = device_idx
+
+        self._initialized = False
+        self._initialized_lock = threading.Lock()
+
+        self._pool = [None] * NCCL_STREAM_POOL_SIZE
+        self._counter = 0
+        self._pool_lock = threading.Lock()
+
+    def get_stream(self):
+        """Get an available stream from the pool.
+
+        The function locks the stream pool and releases the lock before
+        returning.
+
+        Returns:
+            stream (cupy.cuda.Stream): the returned stream from pool.
+        """
+
+        # check the flag
+        self._initialized_lock.acquire()
+        if not self._initialized:
+            self._init_once()
+        self._initialized_lock.release()
+
+        # Get the stream from the pool.
+        self._pool_lock.acquire()
+        stream = self._pool[self._counter]
+        self._counter = (self._counter + 1) % NCCL_STREAM_POOL_SIZE
+        self._pool_lock.release()
+        return stream
+
+    def _init_once(self):
+        """Initialize the stream pool only for once."""
+        with nccl_util.Device(self.device_idx):
+            for i in range(NCCL_STREAM_POOL_SIZE):
+                # this is the only place where self._pool will be written.
+                if ENV.NCCL_USE_MULTISTREAM.val:
+                    logger.debug("NCCL multistream enabled.")
+                    self._pool[i] = cupy.cuda.Stream(null=False, non_blocking=False)
+                else:
+                    logger.debug("NCCL multistream disabled.")
+                    self._pool[i] = cupy.cuda.Stream.null
+        self._init_flag = True
+
+
+# This is a map from GPU index to its stream pool.
+# It is supposed to be READ-ONLY out of this file
+_device_stream_pool_map = dict()
+
+
+def _init_stream_pool():
+    global _device_stream_pool_map
+    for i in range(MAX_GPU_PER_ACTOR):
+        _device_stream_pool_map[i] = StreamPool(i)
+
+
+def get_stream_pool(device_idx):
+    """Get the CUDA stream pool of a GPU device."""
+    # In case there will be multiple threads writing to the pool.
+    lock = threading.Lock()
+    lock.acquire()
+    if not _device_stream_pool_map:
+        _init_stream_pool()
+    lock.release()
+    return _device_stream_pool_map[device_idx]