chore: import upstream snapshot with attribution

python/ray/tests/accelerators/test_rbln.py

@@ -0,0 +1,82 @@
+import os
+import sys
+
+import pytest
+
+from ray._private.accelerators.rbln import (
+    NOSET_RBLN_RT_VISIBLE_DEVICES_ENV_VAR,
+    RBLN_RT_VISIBLE_DEVICES_ENV_VAR,
+    RBLNAcceleratorManager,
+)
+
+
+@pytest.fixture(autouse=True)
+def mock_rebel_module(monkeypatch):
+    from ray.tests.accelerators import mock_rebel
+
+    monkeypatch.setitem(sys.modules, "rebel", mock_rebel)
+
+
+@pytest.fixture
+def clear_rbln_environment():
+    original_env = os.environ.get(RBLN_RT_VISIBLE_DEVICES_ENV_VAR)
+    original_no_set_env = os.environ.get(NOSET_RBLN_RT_VISIBLE_DEVICES_ENV_VAR)
+
+    os.environ.pop(RBLN_RT_VISIBLE_DEVICES_ENV_VAR, None)
+    os.environ.pop(NOSET_RBLN_RT_VISIBLE_DEVICES_ENV_VAR, None)
+
+    yield
+
+    if original_env is not None:
+        os.environ[RBLN_RT_VISIBLE_DEVICES_ENV_VAR] = original_env
+    if original_no_set_env is not None:
+        os.environ[NOSET_RBLN_RT_VISIBLE_DEVICES_ENV_VAR] = original_no_set_env
+
+
+@pytest.mark.usefixtures("clear_rbln_environment")
+class TestRBLNAcceleratorManager:
+    def test_get_resource_name(self):
+        assert RBLNAcceleratorManager.get_resource_name() == "RBLN"
+
+    def test_get_visible_accelerator_ids_env_var(self):
+        assert (
+            RBLNAcceleratorManager.get_visible_accelerator_ids_env_var()
+            == RBLN_RT_VISIBLE_DEVICES_ENV_VAR
+        )
+
+    def test_get_current_process_visible_accelerator_ids(self):
+        os.environ[RBLN_RT_VISIBLE_DEVICES_ENV_VAR] = "0,1,2,3"
+        assert RBLNAcceleratorManager.get_current_process_visible_accelerator_ids() == [
+            "0",
+            "1",
+            "2",
+            "3",
+        ]
+
+        os.environ[RBLN_RT_VISIBLE_DEVICES_ENV_VAR] = ""
+        assert (
+            RBLNAcceleratorManager.get_current_process_visible_accelerator_ids() == []
+        )
+
+        os.environ.pop(RBLN_RT_VISIBLE_DEVICES_ENV_VAR)
+        assert (
+            RBLNAcceleratorManager.get_current_process_visible_accelerator_ids() is None
+        )
+
+    def test_get_current_node_num_accelerators(self):
+        assert RBLNAcceleratorManager.get_current_node_num_accelerators() == 4
+
+    def test_get_current_node_accelerator_type(self):
+        assert RBLNAcceleratorManager.get_current_node_accelerator_type() == "RBLN-CA02"
+
+    def test_set_current_process_visible_accelerator_ids(self):
+        RBLNAcceleratorManager.set_current_process_visible_accelerator_ids(["0", "1"])
+        assert os.environ[RBLN_RT_VISIBLE_DEVICES_ENV_VAR] == "0,1"
+
+        os.environ[NOSET_RBLN_RT_VISIBLE_DEVICES_ENV_VAR] = "1"
+        RBLNAcceleratorManager.set_current_process_visible_accelerator_ids(["2", "3"])
+        assert os.environ[RBLN_RT_VISIBLE_DEVICES_ENV_VAR] == "0,1"
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))