[core][refactor] Move accelerator-specific environment variables to ray_constants.py to avoid redefining them

python/ray/_private/ray_constants.py

@@ -422,15 +422,6 @@ def env_set_by_user(key):
 
 LANGUAGE_WORKER_TYPES = ["python", "java", "cpp"]
 
-# Accelerator constants
-NOSET_CUDA_VISIBLE_DEVICES_ENV_VAR = "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES"
-
-CUDA_VISIBLE_DEVICES_ENV_VAR = "CUDA_VISIBLE_DEVICES"
-ROCR_VISIBLE_DEVICES_ENV_VAR = "ROCR_VISIBLE_DEVICES"
-NEURON_RT_VISIBLE_CORES_ENV_VAR = "NEURON_RT_VISIBLE_CORES"
-TPU_VISIBLE_CHIPS_ENV_VAR = "TPU_VISIBLE_CHIPS"
-NPU_RT_VISIBLE_DEVICES_ENV_VAR = "ASCEND_RT_VISIBLE_DEVICES"
-
 NEURON_CORES = "neuron_cores"
 GPU = "GPU"
 TPU = "TPU"

python/ray/air/_internal/device_manager/npu.py

@@ -7,6 +7,7 @@
 import ray
 import ray._private.ray_constants as ray_constants
 from ray.air._internal.device_manager.torch_device_manager import TorchDeviceManager
+from ray._private.accelerators.npu import ASCEND_RT_VISIBLE_DEVICES_ENV_VAR
 
 
 def is_package_present(package_name: str) -> bool:
@@ -54,9 +55,7 @@ def get_devices(self) -> List[torch.device]:
             device_ids = []
 
             if len(npu_ids) > 0:
-                npu_visible_str = os.environ.get(
-                    ray_constants.NPU_RT_VISIBLE_DEVICES_ENV_VAR, ""
-                )
+                npu_visible_str = os.environ.get(ASCEND_RT_VISIBLE_DEVICES_ENV_VAR, "")
                 if npu_visible_str and npu_visible_str != "NoDevFiles":
                     npu_visible_list = npu_visible_str.split(",")
                 else:

python/ray/dashboard/modules/job/job_manager.py

@@ -13,6 +13,7 @@
 from ray._private.event.event_logger import get_event_logger
 from ray._private.gcs_utils import GcsAioClient
 from ray._private.utils import run_background_task
+from ray._private.accelerators.nvidia_gpu import NOSET_CUDA_VISIBLE_DEVICES_ENV_VAR
 from ray.actor import ActorHandle
 from ray.core.generated.event_pb2 import Event
 from ray.dashboard.consts import (
@@ -359,7 +360,7 @@ def _get_supervisor_runtime_env(
             # Don't set CUDA_VISIBLE_DEVICES for the supervisor actor so the
             # driver can use GPUs if it wants to. This will be removed from
             # the driver's runtime_env so it isn't inherited by tasks & actors.
-            env_vars[ray_constants.NOSET_CUDA_VISIBLE_DEVICES_ENV_VAR] = "1"
+            env_vars[NOSET_CUDA_VISIBLE_DEVICES_ENV_VAR] = "1"
         runtime_env["env_vars"] = env_vars
 
         if os.getenv(RAY_STREAM_RUNTIME_ENV_LOG_TO_JOB_DRIVER_LOG_ENV_VAR, "0") == "1":

python/ray/dashboard/modules/job/job_supervisor.py

@@ -16,6 +16,7 @@
 from ray._private.ray_logging.formatters import JSONFormatter, TextFormatter
 from ray._private.runtime_env.constants import RAY_JOB_CONFIG_JSON_ENV_VAR
 from ray._private.utils import remove_ray_internal_flags_from_env
+from ray._private.accelerators.nvidia_gpu import NOSET_CUDA_VISIBLE_DEVICES_ENV_VAR
 from ray.actor import ActorHandle
 from ray.dashboard.modules.job.common import (
     JOB_ID_METADATA_KEY,
@@ -139,7 +140,7 @@ def _get_driver_runtime_env(
         # Allow CUDA_VISIBLE_DEVICES to be set normally for the driver's tasks
         # & actors.
         env_vars = curr_runtime_env.get("env_vars", {})
-        env_vars.pop(ray_constants.NOSET_CUDA_VISIBLE_DEVICES_ENV_VAR)
+        env_vars.pop(NOSET_CUDA_VISIBLE_DEVICES_ENV_VAR)
         env_vars.pop(ray_constants.RAY_WORKER_NICENESS)
         curr_runtime_env["env_vars"] = env_vars
         return curr_runtime_env

python/ray/dashboard/modules/job/tests/subprocess_driver_scripts/check_cuda_devices.py

@@ -2,7 +2,7 @@
 
 import ray
 
-cuda_env = ray._private.ray_constants.NOSET_CUDA_VISIBLE_DEVICES_ENV_VAR
+cuda_env = ray._private.accelerators.nvidia_gpu.NOSET_CUDA_VISIBLE_DEVICES_ENV_VAR
 if os.environ.get("RAY_TEST_RESOURCES_SPECIFIED") == "1":
     assert cuda_env not in os.environ
     assert "CUDA_VISIBLE_DEVICES" in os.environ

python/ray/train/_internal/backend_executor.py

@@ -7,6 +7,10 @@
 
 import ray
 import ray._private.ray_constants as ray_constants
+from ray._private.accelerators.neuron import NEURON_RT_VISIBLE_CORES_ENV_VAR
+from ray._private.accelerators.npu import ASCEND_RT_VISIBLE_DEVICES_ENV_VAR
+from ray._private.accelerators.amd_gpu import ROCR_VISIBLE_DEVICES_ENV_VAR
+from ray._private.accelerators.nvidia_gpu import CUDA_VISIBLE_DEVICES_ENV_VAR
 from ray._private.ray_constants import env_integer
 from ray.data import Dataset
 from ray.exceptions import RayActorError
@@ -118,18 +122,18 @@ def __init__(
             ResourceConfig(
                 ray_constants.NEURON_CORES,
                 ENABLE_SHARE_NEURON_CORES_ACCELERATOR_ENV,
-                ray_constants.NEURON_RT_VISIBLE_CORES_ENV_VAR,
+                NEURON_RT_VISIBLE_CORES_ENV_VAR,
             ),
             ResourceConfig(
                 ray_constants.NPU,
                 ENABLE_SHARE_NPU_RT_VISIBLE_DEVICES_ENV,
-                ray_constants.NPU_RT_VISIBLE_DEVICES_ENV_VAR,
+                ASCEND_RT_VISIBLE_DEVICES_ENV_VAR,
             ),
             # For AMD GPUs, they are using ROCR_VISIBLE_DEVICES env var.
             ResourceConfig(
                 ray_constants.GPU,
                 ENABLE_SHARE_ROCR_VISIBLE_DEVICES_ENV,
-                ray_constants.ROCR_VISIBLE_DEVICES_ENV_VAR,
+                ROCR_VISIBLE_DEVICES_ENV_VAR,
             ),
         ]
 
@@ -299,9 +303,7 @@ def _share_cuda_visible_devices(self):
             - Worker3: "0,1"
 
         """
-        self._share_resource_ids(
-            ray_constants.GPU, ray_constants.CUDA_VISIBLE_DEVICES_ENV_VAR
-        )
+        self._share_resource_ids(ray_constants.GPU, CUDA_VISIBLE_DEVICES_ENV_VAR)
 
     def _share_resource_ids(self, resource: str, env_var: str):
         """Sets the given env_var on all workers.

python/ray/train/tests/test_backend.py

@@ -8,7 +8,7 @@
 import pytest
 
 import ray
-import ray._private.ray_constants as ray_constants
+from ray._private.accelerators.neuron import NEURON_RT_VISIBLE_CORES_ENV_VAR
 from ray import train
 from ray.air._internal.util import StartTraceback
 
@@ -493,7 +493,7 @@ def test_neuron_core_accelerator_ids(ray_2_node_2_neuron_cores, worker_results):
     config = TestConfig()
 
     def get_resources():
-        neuron_resource_ids = os.environ[ray_constants.NEURON_RT_VISIBLE_CORES_ENV_VAR]
+        neuron_resource_ids = os.environ[NEURON_RT_VISIBLE_CORES_ENV_VAR]
         # Sort the runtime ids to have exact match with expected result.
         sorted_devices = [
             int(device) for device in sorted(neuron_resource_ids.split(","))
@@ -530,7 +530,7 @@ def test_neuron_core_accelerator_ids_sharing_disabled(
     config = TestConfig()
 
     def get_resources():
-        neuron_resource_ids = os.environ[ray_constants.NEURON_RT_VISIBLE_CORES_ENV_VAR]
+        neuron_resource_ids = os.environ[NEURON_RT_VISIBLE_CORES_ENV_VAR]
         # Sort the runtime ids to have exact match with expected result.
         sorted_devices = [
             int(device) for device in sorted(neuron_resource_ids.split(","))

python/ray/train/v2/_internal/callbacks/accelerators.py

@@ -5,6 +5,7 @@
 
 import ray._private.ray_constants as ray_constants
 from ray._private.ray_constants import env_bool
+from ray._private.accelerators.nvidia_gpu import CUDA_VISIBLE_DEVICES_ENV_VAR
 from ray.train import BackendConfig
 from ray.train.constants import ENABLE_SHARE_CUDA_VISIBLE_DEVICES_ENV
 from ray.train.v2._internal.execution.callback import WorkerGroupCallback
@@ -63,7 +64,7 @@ def _share_cuda_visible_devices(worker_group: WorkerGroup):
         - Worker2: "0,1"
     """
     _share_accelerator_ids(
-        worker_group, ray_constants.GPU, ray_constants.CUDA_VISIBLE_DEVICES_ENV_VAR
+        worker_group, ray_constants.GPU, CUDA_VISIBLE_DEVICES_ENV_VAR
     )