forked from EngineX-Cambricon/enginex-mlu370-vllm
93 lines
3.6 KiB
Python
Executable File
93 lines
3.6 KiB
Python
Executable File
import os
|
|
import sys
|
|
import pytest
|
|
from unittest.mock import patch
|
|
|
|
import ray
|
|
from ray._private.accelerators import MLUAcceleratorManager as Accelerator
|
|
|
|
|
|
@patch("glob.glob")
|
|
@patch("os.listdir")
|
|
def test_autodetect_num_mlus(mock_list, mock_glob):
|
|
mock_glob.return_value = [f"/dev/davinci{i}" for i in range(4)]
|
|
# mock_list.return_value = []
|
|
assert Accelerator.get_current_node_num_accelerators() == 4
|
|
|
|
|
|
@patch("glob.glob")
|
|
@patch("os.listdir")
|
|
def test_autodetect_num_mlus_without_devices(mock_list, mock_glob):
|
|
mock_glob.side_effect = Exception
|
|
# mock_list.return_value = []
|
|
assert Accelerator.get_current_node_num_accelerators() == 0
|
|
|
|
|
|
def test_mlu_accelerator_manager_api():
|
|
assert Accelerator.get_resource_name() == "MLU"
|
|
assert Accelerator.get_visible_accelerator_ids_env_var() == "MLU_VISIBLE_DEVICES"
|
|
assert Accelerator.validate_resource_request_quantity(0.5) == (True, None)
|
|
assert Accelerator.validate_resource_request_quantity(1) == (True, None)
|
|
|
|
|
|
def test_visible_mlu_type(monkeypatch, shutdown_only):
|
|
with patch.object(
|
|
Accelerator, "get_current_node_num_accelerators", return_value=4
|
|
), patch.object(
|
|
Accelerator, "get_current_node_accelerator_type", return_value="MLU370"
|
|
):
|
|
monkeypatch.setenv("MLU_VISIBLE_DEVICES", "0,1,2")
|
|
manager = ray._private.accelerators.get_accelerator_manager_for_resource("MLU")
|
|
assert manager.get_current_node_accelerator_type() == "MLU370"
|
|
|
|
@pytest.mark.skipif(sys.platform == "win32", reason="Not supported mock on Windows")
|
|
def test_visible_mlu_ids(monkeypatch, shutdown_only):
|
|
monkeypatch.setenv("MLU_VISIBLE_DEVICES", "0,1,2")
|
|
with patch.object(Accelerator, "get_current_node_num_accelerators", return_value=4):
|
|
|
|
ray.init()
|
|
manager = ray._private.accelerators.get_accelerator_manager_for_resource("MLU")
|
|
assert manager.get_current_node_num_accelerators() == 4
|
|
assert manager.__name__ == "MLUAcceleratorManager"
|
|
assert ray.available_resources()["MLU"] == 3
|
|
|
|
def test_get_current_process_visible_accelerator_ids(monkeypatch, shutdown_only):
|
|
monkeypatch.setenv("MLU_VISIBLE_DEVICES", "0,1,2")
|
|
assert Accelerator.get_current_process_visible_accelerator_ids() == ["0", "1", "2"]
|
|
|
|
monkeypatch.delenv("MLU_VISIBLE_DEVICES")
|
|
assert Accelerator.get_current_process_visible_accelerator_ids() is None
|
|
|
|
monkeypatch.setenv("MLU_VISIBLE_DEVICES", "")
|
|
assert Accelerator.get_current_process_visible_accelerator_ids() == []
|
|
|
|
monkeypatch.setenv("MLU_VISIBLE_DEVICES", "NoDevFiles")
|
|
assert Accelerator.get_current_process_visible_accelerator_ids() == []
|
|
|
|
|
|
def test_set_current_process_visible_accelerator_ids(shutdown_only):
|
|
Accelerator.set_current_process_visible_accelerator_ids(["0"])
|
|
assert os.environ["MLU_VISIBLE_DEVICES"] == "0"
|
|
|
|
Accelerator.set_current_process_visible_accelerator_ids(["0", "1"])
|
|
assert os.environ["MLU_VISIBLE_DEVICES"] == "0,1"
|
|
|
|
Accelerator.set_current_process_visible_accelerator_ids(["0", "1", "2"])
|
|
assert os.environ["MLU_VISIBLE_DEVICES"] == "0,1,2"
|
|
|
|
|
|
@pytest.mark.skipif(sys.platform == "win32", reason="Not supported mock on Windows")
|
|
def test_auto_detected_more_than_visible(monkeypatch, shutdown_only):
|
|
with patch.object(Accelerator, "get_current_node_num_accelerators", return_value=4):
|
|
# If more MLUs are detected than visible.
|
|
monkeypatch.setenv("MLU_VISIBLE_DEVICES", "0,1,2")
|
|
|
|
ray.init()
|
|
assert ray.available_resources()["MLU"] == 3
|
|
|
|
if __name__ == "__main__":
|
|
if os.environ.get("PARALLEL_CI"):
|
|
sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
|
|
else:
|
|
sys.exit(pytest.main(["-sv", __file__]))
|