add other platform and refactory

Signed-off-by: noemotiovon <757486878@qq.com>

Author: noemotiovon

roll/distributed/executor/cluster.py

@@ -22,7 +22,6 @@
 from roll.distributed.scheduler.resource_manager import ResourceManager
 from roll.utils.import_utils import safe_import_class
 from roll.utils.logging import get_logger
-from roll.utils.ray_utils import RayUtils
 
 logger = get_logger()
 
@@ -114,9 +113,7 @@ def _create_workers(self):
                 env_vars["MASTER_ADDR"] = self.master_addr
                 env_vars["MASTER_PORT"] = str(self.master_port)
             if deploy_pg["gpu_rank"] is not None:
-                RayUtils.update_env_vars_for_visible_devices(
-                    env_vars=env_vars,
-                    gpu_ranks=pg_zero_gpu_ranks)
+                current_platform.update_env_vars_for_visible_devices(env_vars=env_vars, gpu_ranks=pg_zero_gpu_ranks)
             if "ROLL_LOG_DIR" in os.environ:
                 env_vars["ROLL_LOG_DIR"] = os.environ["ROLL_LOG_DIR"]
             env_vars.update(self.worker_config.system_envs)
@@ -135,10 +132,14 @@ def _create_workers(self):
             if current_platform.ray_device_key == "GPU":
                 worker_options.update({"num_gpus": 0.01 if self.worker_config.device_mapping else 0})
             elif current_platform.ray_device_key == "NPU":
-                worker_options.update({
-                    "num_cpus": 0,
-                    "resources": {current_platform.ray_device_key: 0.01 if self.worker_config.device_mapping else 0},
-                })
+                worker_options.update(
+                    {
+                        "num_gpus": 0,
+                        "resources": {
+                            current_platform.ray_device_key: 0.01 if self.worker_config.device_mapping else 0
+                        },
+                    }
+                )
 
             worker = self.worker_cls.options(**worker_options).remote(worker_config=self.worker_config)
             self.workers.append(worker)

roll/distributed/executor/worker.py

@@ -17,8 +17,6 @@
 from roll.utils.logging import get_logger
 from roll.utils.offload_states import OffloadStateType
 from roll.platforms import current_platform
-from roll.utils.ray_utils import RayUtils
-from roll.platforms import current_platform
 
 
 @dataclass
@@ -121,7 +119,7 @@ def get_master_addr_and_port(self):
 
     @staticmethod
     def get_visible_gpus():
-        return RayUtils.get_visible_gpus()
+        return current_platform.get_visible_gpus()
 
     def get_devices_info(self):
         devices_info = [

roll/distributed/scheduler/driver_utils.py

@@ -15,27 +15,27 @@ def is_driver():
 
 
 def get_driver_rank():
-    assert is_driver(), "this function should only be ran on a driver"
+    assert is_driver(), "this function should only be run on a driver"
     return int(os.getenv("RANK", "0"))
 
 
 def get_driver_world_size():
-    assert is_driver(), "this function should only be ran on a driver"
+    assert is_driver(), "this function should only be run on a driver"
     return int(os.getenv("WORLD_SIZE", "1"))
 
 
 def get_driver_master_addr():
-    assert is_driver(), "this function should only be ran on a driver"
+    assert is_driver(), "this function should only be run on a driver"
     return os.getenv("MASTER_ADDR", "127.0.0.1")
 
 
 def get_driver_master_port():
-    assert is_driver(), "this function should only be ran on a driver"
+    assert is_driver(), "this function should only be run on a driver"
     return os.getenv("MASTER_PORT", "6379")
 
 
 def get_driver_node_name():
-    assert is_driver(), "this function should only be ran on a driver"
+    assert is_driver(), "this function should only be run on a driver"
     return os.getenv("WORKER_ID", f"{get_driver_master_addr()}:{get_driver_rank()}")
 
 def is_multi_tenant():

roll/distributed/scheduler/initialize.py

@@ -18,7 +18,7 @@
 from roll.distributed.scheduler.log_monitor import LogMonitorListener
 from roll.utils.constants import RAY_NAMESPACE
 from roll.utils.logging import get_logger
-from roll.utils.ray_utils import RayUtils
+from roll.platforms import current_platform
 
 logger = get_logger()
 
@@ -60,7 +60,7 @@ def init():
     manual_start = start_ray_cluster()
 
     runtime_env = {
-        "env_vars": RayUtils.get_custom_env_env_vars(),
+        "env_vars": current_platform.get_custom_env_vars(),
     }
 
     if not ray.is_initialized():

roll/distributed/scheduler/resource_manager.py

@@ -55,7 +55,7 @@ def __init__(self, num_gpus_per_node, num_nodes):
                         if current_platform.ray_device_key == "GPU"
                         else {"resources": {current_platform.ray_device_key: self.gpu_per_node}}
                     )
-                ).remote(current_platform.ray_device_key)
+                ).remote(current_platform.device_control_env_var)
                 for pg in self.placement_groups
             ])
             print(f"gpu ranks: {gpu_ranks}")

roll/platforms/__init__.py

@@ -1,29 +1,56 @@
-import logging
+import torch
+
 from .platform import Platform
 from .cuda import CudaPlatform
-from .npu import NPUPlatform
+from .npu import NpuPlatform
+from .rocm import RocmPlatform
+from .unknown import UnknownPlatform
+from .cpu import CpuPlatform
+
+from roll.utils.logging import get_logger
+
+
+logger = get_logger()
 
-logger = logging.getLogger(__name__)
 
 def _init_platform() -> Platform:
-    """Initialize and return the current platform instance.
+    """
+    Detect and initialize the appropriate platform based on available devices.
+
+    Priority:
+    1. CUDA (NVIDIA / AMD ROCm)
+    2. NPU (if torch_npu is installed)
+    3. CPU (fallback)
 
-    Automatically selects the platform based on environment and availability:
-    - If torch_npu is installed, use NPUPlatform.
-    - Otherwise, fall back to CudaPlatform.
+    Returns:
+        An instance of a subclass of Platform corresponding to the detected hardware.
     """
-    try:
-        import torch_npu  # noqa: F401
-        logger.info("Detected torch_npu. Initializing NPU platform.")
-        return NPUPlatform()
-    except ImportError:
-        logger.info("Initializing ROLL default device backend: Cuda platform.")
-        return CudaPlatform()
+    if torch.cuda.is_available():
+        device_name = torch.cuda.get_device_name().upper()
+        logger.info(f"Detected CUDA device: {device_name}")
+        if "NVIDIA" in device_name:
+            logger.info("Initializing CUDA platform (NVIDIA).")
+            return CudaPlatform()
+        elif "AMD" in device_name:
+            logger.info("Initializing ROCm platform (AMD).")
+            return RocmPlatform()
+        logger.warning("Unrecognized CUDA device. Falling back to UnknownPlatform.")
+        return UnknownPlatform()
+    else:
+        try:
+            import torch_npu  # noqa: F401
+
+            logger.info("Detected torch_npu. Initializing NPU platform.")
+            return NpuPlatform()
+        except ImportError:
+            logger.info("No supported accelerator detected. Initializing CPU platform.")
+            return CpuPlatform()
+
 
 # Global singleton representing the current platform in use.
 current_platform: Platform = _init_platform()
 
 __all__ = [
     "Platform",
     "current_platform",
-]
+]

roll/platforms/cpu.py

@@ -0,0 +1,21 @@
+from .platform import Platform
+from roll.utils.logging import get_logger
+
+
+logger = get_logger()
+
+
+class CpuPlatform(Platform):
+    device_name: str = "CPU"
+    device_type: str = "cpu"
+    dispatch_key: str = "CPU"
+    ray_device_key: str = "CPU"
+    communication_backend: str = "gloo"
+
+    @classmethod
+    def clear_cublas_workspaces(cls):
+        pass
+
+    @classmethod
+    def get_custom_env_vars(cls) -> dict:
+        return {}

roll/platforms/cuda.py

@@ -1,9 +1,13 @@
 from .platform import Platform
-from roll.utils.logging import logger
+from roll.utils.logging import get_logger
 
 import torch
 
+logger = get_logger()
+
+
 class CudaPlatform(Platform):
+    device_name: str = "NVIDIA"
     device_type: str = "cuda"
     dispatch_key: str = "CUDA"
     ray_device_key: str = "GPU"
@@ -14,24 +18,37 @@ class CudaPlatform(Platform):
     @classmethod
     def clear_cublas_workspaces(cls):
         torch._C._cuda_clearCublasWorkspaces()
-    
+
     @classmethod
     def get_vllm_worker_class(clas):
         try:
             from vllm import envs
+
             if envs.VLLM_USE_V1:
                 from vllm.v1.worker.gpu_worker import Worker
+
                 logger.info("Successfully imported vLLM V1 Worker.")
                 return Worker
             else:
                 from vllm.worker.worker import Worker
+
                 logger.info("Successfully imported vLLM V0 Worker.")
                 return Worker
         except ImportError as e:
-            logger.error("Failed to import vLLM Worker. "
-                         "Make sure vLLM is installed correctly: %s", e)
+            logger.error("Failed to import vLLM Worker. " "Make sure vLLM is installed correctly: %s", e)
             raise RuntimeError("vLLM is not installed or not properly configured.") from e
-        
+
     @classmethod
     def set_allocator_settings(cls):
         torch.cuda.memory._set_allocator_settings("expandable_segments:False")
+
+    @classmethod
+    def get_custom_env_vars(cls) -> dict:
+        env_vars = {
+            # "RAY_DEBUG": "legacy"
+            "TORCHINDUCTOR_COMPILE_THREADS": "2",
+            "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True",
+            "NCCL_CUMEM_ENABLE": "0",  # https://github.com/NVIDIA/nccl/issues/1234
+            "NCCL_NVLS_ENABLE": "0",
+        }
+        return env_vars

roll/platforms/npu.py

@@ -1,9 +1,11 @@
 from .platform import Platform
-from roll.utils.logging import logger
+from roll.utils.logging import get_logger
 
-import torch
+logger = get_logger()
 
-class NPUPlatform(Platform):
+
+class NpuPlatform(Platform):
+    device_name: str = "ASCEND"
     device_type: str = "npu"
     dispatch_key: str = "PrivateUse1"
     ray_device_key: str = "NPU"
@@ -14,24 +16,33 @@ class NPUPlatform(Platform):
     @classmethod
     def clear_cublas_workspaces(cls):
         pass
-    
+
     @classmethod
     def get_vllm_worker_class(clas):
         try:
             from vllm import envs
+
             if envs.VLLM_USE_V1:
                 from vllm_ascend.worker.worker_v1 import NPUWorker as Worker
+
                 logger.info("Successfully imported vLLM V1 Worker.")
                 return Worker
             else:
                 from vllm_ascend.worker.worker import NPUWorker as Worker
+
                 logger.info("Successfully imported vLLM V0 Worker.")
                 return Worker
         except ImportError as e:
-            logger.error("Failed to import vLLM Worker. "
-                         "Make sure vLLM is installed correctly: %s", e)
+            logger.error("Failed to import vLLM Worker. " "Make sure vLLM is installed correctly: %s", e)
             raise RuntimeError("vLLM is not installed or not properly configured.") from e
-        
+
     @classmethod
     def set_allocator_settings(cls):
         pass
+
+    @classmethod
+    def get_custom_env_vars(cls) -> dict:
+        env_vars = {
+            "TORCHINDUCTOR_COMPILE_THREADS": "2",
+        }
+        return env_vars

roll/platforms/platform.py

@@ -1,5 +1,9 @@
 import torch
-from roll.utils.logging import logger
+import os
+
+from roll.utils.logging import get_logger
+
+logger = get_logger()
 
 
 class Platform:
@@ -23,7 +27,12 @@ class Platform:
        - `get_vllm_worker_class`: to specify the vLLM Ray worker class.
        - `set_allocator_settings`: to configure platform-specific memory allocators.
     """
-    # Corresponding torch module name, e.g., "cuda", "npu"
+    # High-level platform name, used for readability and logging.
+    # Examples: "NVIDIA", "AMD", "ASCEND"
+    device_name: str
+
+    # Corresponding torch module name
+    # Examples: "cuda", "npu"
     device_type: str
 
     # available dispatch keys:
@@ -92,3 +101,21 @@ def get_vllm_worker_class(cls):
     def set_allocator_settings(cls):
         """Configure memory allocator settings based on the device type."""
         raise NotImplementedError
+
+    @classmethod
+    def get_custom_env_vars(cls) -> dict:
+        raise NotImplementedError
+
+    @classmethod
+    def update_env_vars_for_visible_devices(cls, env_vars: dict, gpu_ranks: list):
+        visible_devices_env_vars = {
+            cls.device_control_env_var: ",".join(map(str, gpu_ranks)),
+            cls.ray_experimental_noset: "1",
+        }
+        env_vars.update(visible_devices_env_vars)
+
+    @classmethod
+    def get_visible_gpus(cls) -> list:
+        if cls.device_control_env_var is not None:
+            return os.environ.get(cls.device_control_env_var, "").split(",")
+        return []