wehub-resource-sync
93 lines
2.8 KiB
Python
93 lines
2.8 KiB
Python
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__]))
|