chore: import upstream snapshot with attribution

python/ray/tests/accelerators/test_amd_gpu.py

@@ -0,0 +1,92 @@
+import os
+import sys
+from unittest.mock import patch
+
+import pytest
+
+import ray
+from ray._private.accelerators import (
+    AMDGPUAcceleratorManager,
+    get_accelerator_manager_for_resource,
+)
+
+
+@patch(
+    "ray._private.accelerators.AMDGPUAcceleratorManager.get_current_node_num_accelerators",  # noqa: E501
+    return_value=4,
+)
+def test_visible_amd_gpu_ids(mock_get_num_accelerators, monkeypatch, shutdown_only):
+    monkeypatch.setenv("HIP_VISIBLE_DEVICES", "0,1,2")
+    # Delete the cache so it can be re-populated the next time
+    # we call get_accelerator_manager_for_resource
+    del get_accelerator_manager_for_resource._resource_name_to_accelerator_manager
+    ray.init()
+    _ = mock_get_num_accelerators.called
+    assert ray.available_resources()["GPU"] == 3
+
+
+@patch(
+    "ray._private.accelerators.AMDGPUAcceleratorManager._get_amd_device_ids",
+    return_value=["0x74a1", "0x74a1", "0x74a1", "0x74a1"],
+)
+def test_visible_amd_gpu_type(mock_get_amd_device_ids, shutdown_only):
+    ray.init()
+    _ = mock_get_amd_device_ids.called
+    assert (
+        AMDGPUAcceleratorManager.get_current_node_accelerator_type()
+        == "AMD-Instinct-MI300X-OAM"
+    )
+
+
+@patch(
+    "ray._private.accelerators.AMDGPUAcceleratorManager._get_amd_device_ids",
+    return_value=["0x640f", "0x640f", "0x640f", "0x640f"],
+)
+def test_visible_amd_gpu_type_bad_device_id(mock_get_num_accelerators, shutdown_only):
+    ray.init()
+    _ = mock_get_num_accelerators.called
+    assert AMDGPUAcceleratorManager.get_current_node_accelerator_type() is None
+
+
+def test_get_current_process_visible_accelerator_ids(monkeypatch):
+    monkeypatch.setenv("HIP_VISIBLE_DEVICES", "0,1,2")
+    assert AMDGPUAcceleratorManager.get_current_process_visible_accelerator_ids() == [
+        "0",
+        "1",
+        "2",
+    ]
+
+    monkeypatch.setenv("HIP_VISIBLE_DEVICES", "0,2,7")
+    assert AMDGPUAcceleratorManager.get_current_process_visible_accelerator_ids() == [
+        "0",
+        "2",
+        "7",
+    ]
+
+    monkeypatch.setenv("HIP_VISIBLE_DEVICES", "")
+    assert AMDGPUAcceleratorManager.get_current_process_visible_accelerator_ids() == []
+
+    del os.environ["HIP_VISIBLE_DEVICES"]
+    assert (
+        AMDGPUAcceleratorManager.get_current_process_visible_accelerator_ids() is None
+    )
+
+
+def test_set_current_process_visible_accelerator_ids():
+    AMDGPUAcceleratorManager.set_current_process_visible_accelerator_ids(["0"])
+    env_var = AMDGPUAcceleratorManager.get_visible_accelerator_ids_env_var()
+    assert os.environ[env_var] == "0"
+
+    AMDGPUAcceleratorManager.set_current_process_visible_accelerator_ids(["0", "1"])
+    assert os.environ[env_var] == "0,1"
+
+    AMDGPUAcceleratorManager.set_current_process_visible_accelerator_ids(
+        ["0", "1", "7"]
+    )
+    assert os.environ[env_var] == "0,1,7"
+
+    del os.environ[env_var]
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))