chore: import upstream snapshot with attribution

python/ray/experimental/channel/accelerator_context.py

@@ -0,0 +1,246 @@
+import importlib
+import threading
+from contextlib import nullcontext
+from typing import TYPE_CHECKING, ContextManager, List, Optional, Type
+
+import ray
+from ray._private.accelerators import get_accelerator_manager_for_resource
+from ray.experimental.channel.communicator import Communicator
+
+if TYPE_CHECKING:
+    import torch
+
+# The accelerator context singleton on this process.
+_accelerator_context_lock = threading.Lock()
+_default_accelerator_context: Optional["AcceleratorContext"] = None
+_global_custom_context: Optional["AcceleratorContext"] = None
+
+
+class AcceleratorContext:
+    """
+    Provides a unified interface for managing different accelerator backends
+    This includes stream management, event creation, device context control,
+    and communicator support for distributed communication.
+    """
+
+    def __init__(self, torch_module_name: str, communicator_cls: Type[Communicator]):
+        """
+        Initializes an accelerator context with the specified torch device module
+        and communicator class.
+
+        Args:
+            torch_module_name: Name of the torch device module (e.g., "cuda", "cpu").
+            communicator_cls: Class used to handle communication.
+        """
+
+        # The name of the torch module (e.g., 'cuda', 'npu')
+        self._torch_module_name: str = torch_module_name
+        # The Communicator class used to manage communication
+        self._communicator_cls: Type[Communicator] = communicator_cls
+
+        # Import the torch backend module (e.g., torch.cuda) if the device is not 'cpu'.
+        if torch_module_name != "cpu":
+            self._torch_mod = importlib.import_module(f"torch.{torch_module_name}")
+
+    @staticmethod
+    def get() -> "AcceleratorContext":
+        """
+        Returns the singleton instance of the accelerator context.
+
+        If a custom accelerator has been registered, initializes the context
+        based on the registration. Otherwise, selects an appropriate runtime
+        based on the available device (CUDA or CPU) and registers the
+        corresponding default communicator.
+
+        Returns:
+            AcceleratorContext: A singleton instance of the appropriate
+            runtime context.
+        """
+
+        global _default_accelerator_context, _global_custom_context
+
+        with _accelerator_context_lock:
+            if _global_custom_context is not None:
+                return _global_custom_context
+
+            if _default_accelerator_context is None:
+                if len(ray.get_gpu_ids()) > 0:
+                    from ray.experimental.channel.nccl_group import _NcclGroup
+
+                    _default_accelerator_context = AcceleratorContext(
+                        "cuda", _NcclGroup
+                    )
+                else:
+                    from ray.experimental.channel.cpu_communicator import (
+                        CPUCommunicator,
+                    )
+
+                    _default_accelerator_context = AcceleratorContext(
+                        "cpu", CPUCommunicator
+                    )
+
+            return _default_accelerator_context
+
+    @staticmethod
+    def set(accelerator_context: "AcceleratorContext") -> None:
+        """
+        Overwrites the default accelerator context.
+
+        Args:
+            accelerator_context: The context to register.
+        """
+        global _global_custom_context
+
+        # Accelerator context is registered.
+        _global_custom_context = accelerator_context
+
+    def get_accelerator_devices(self) -> List["torch.device"]:
+        """
+        Gets the torch device list configured for this process.
+
+        Returns:
+            List[torch.device]: The torch device list.
+        """
+        import torch
+
+        if self._torch_module_name == "cpu":
+            return [torch.device("cpu")]
+
+        if self._torch_module_name == "cuda":
+            accelerator_ids = [str(id) for id in ray.get_gpu_ids()]
+            accelerator_manager = get_accelerator_manager_for_resource("GPU")
+        else:
+            accelerator_ids = [
+                str(id)
+                for id in ray.get_runtime_context().get_accelerator_ids()[
+                    self._torch_module_name.upper()
+                ]
+            ]
+            accelerator_manager = get_accelerator_manager_for_resource(
+                self._torch_module_name.upper()
+            )
+
+        device_ids = []
+
+        if len(accelerator_ids) > 0:
+            accelerator_visible_list = (
+                accelerator_manager.get_current_process_visible_accelerator_ids()
+            )
+            if accelerator_visible_list is None:
+                accelerator_visible_list = []
+
+            # If there are multiple Accelerators, return a list of devices.
+            # If using fractional Accelerators, these IDs are not guaranteed
+            # to be unique across different processes.
+            for accelerator_id in accelerator_ids:
+                try:
+                    device_ids.append(accelerator_visible_list.index(accelerator_id))
+                except ValueError:
+                    raise RuntimeError(
+                        f"{accelerator_manager.get_visible_accelerator_ids_env_var()} set incorrectly. "
+                        f"expected to include {accelerator_id}. "
+                        "Did you override this environment"
+                        " variable? If not, please help file an issue on Github."
+                    )
+
+        else:
+            # If called on the driver or outside of Ray Train, return the
+            # 0th device.
+            device_ids.append(0)
+
+        return [
+            torch.device(f"{self._torch_module_name}:{device_id}")
+            for device_id in device_ids
+        ]
+
+    def get_device_context(self, device: "torch.device") -> ContextManager:
+        """
+        Retrieves the context manager for the specified accelerator device.
+        There is no device context for CPU, returning a nullcontext.
+
+        Args:
+            device: The target device for which the context manager is required.
+
+        Returns:
+            ContextManager: A context manager specific to the device type.
+        """
+        if device.type == "cpu":
+            return nullcontext()
+
+        return self._torch_mod.device(device)
+
+    def current_stream(self):
+        """
+        Retrieves the current execution stream for the accelerator device.
+        """
+        return self._torch_mod.current_stream()
+
+    def create_event(self):
+        """
+        Creates an event object for the accelerator device.
+        """
+        return self._torch_mod.Event()
+
+    def generate_communicator_id(self) -> str:
+        """
+        Generates a communication identifier for communication group.
+        """
+        return self._communicator_cls.generate_communicator_id()
+
+    def create_communicator(self, *args, **kwargs) -> Communicator:
+        """
+        Creates a communication group for collective operations.
+        """
+        return self._communicator_cls(*args, **kwargs)
+
+    @property
+    def module_name(self) -> str:
+        """
+        Gets the name of the torch module backing the accelerator.
+        """
+        return self._torch_module_name
+
+    @property
+    def communicator_cls(self) -> Optional[Type[Communicator]]:
+        """
+        Returns the communicator class.
+        """
+        return self._communicator_cls
+
+    @property
+    def accelerator_count(self) -> int:
+        """
+        Returns the number of accelerators assigned by ray.
+        """
+        if self._torch_module_name == "cuda":
+            return len(ray.get_gpu_ids())
+        else:
+            accelerator_ids = ray.get_runtime_context().get_accelerator_ids()
+            return len(accelerator_ids.get(self._torch_module_name.upper(), []))
+
+
+def register_accelerator_context(
+    torch_module_name: str, communicator_cls: Type[Communicator]
+):
+    """
+    Registers the accelerator context with the specified device type and communicator.
+
+    Args:
+        torch_module_name: The name of the device module under torch.
+        communicator_cls: The communicator class associated with the device.
+    """
+    accelerator_context = AcceleratorContext(torch_module_name, communicator_cls)
+    AcceleratorContext.set(accelerator_context)
+
+
+def is_accelerator_context_registered():
+    """
+    Checks whether a custom accelerator context has been registered.
+
+    Returns:
+        bool: True if a custom accelerator context is registered
+              (_global_custom_context is not None), False otherwise.
+    """
+    if _global_custom_context is not None:
+        return True
+    return False