wehub-resource-sync
347 lines
12 KiB
Python
347 lines
12 KiB
Python
import os
|
|
import sys
|
|
from unittest import mock
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
import requests
|
|
|
|
from ray._private.accelerators import TPUAcceleratorManager, tpu
|
|
|
|
|
|
@patch("glob.glob")
|
|
def test_autodetect_num_tpus_accel(mock_glob):
|
|
mock_glob.return_value = [
|
|
"/dev/accel0",
|
|
"/dev/accel1",
|
|
"/dev/accel2",
|
|
"/dev/accel3",
|
|
]
|
|
TPUAcceleratorManager.get_current_node_num_accelerators.cache_clear()
|
|
assert TPUAcceleratorManager.get_current_node_num_accelerators() == 4
|
|
|
|
|
|
@patch("os.path.isdir")
|
|
@patch("glob.glob")
|
|
@patch("os.listdir")
|
|
def test_autodetect_num_tpus_accel_ignores_blackwell_directory(
|
|
mock_list, mock_glob, mock_isdir
|
|
):
|
|
# NVIDIA drivers 570.x (Blackwell-class GPUs, e.g. RTX 5090) create
|
|
# /dev/accel as a directory containing /dev/accel/accel0. The non-recursive
|
|
# glob matches the directory entry; filtering directories out keeps real
|
|
# TPU chips (character devices at /dev/accel0..N) while rejecting the
|
|
# NVIDIA false positive.
|
|
mock_glob.return_value = ["/dev/accel"]
|
|
mock_isdir.side_effect = lambda p: p == "/dev/accel"
|
|
mock_list.side_effect = FileNotFoundError
|
|
TPUAcceleratorManager.get_current_node_num_accelerators.cache_clear()
|
|
assert TPUAcceleratorManager.get_current_node_num_accelerators() == 0
|
|
|
|
|
|
@patch("glob.glob")
|
|
@patch("os.listdir")
|
|
def test_autodetect_num_tpus_vfio(mock_list, mock_glob):
|
|
mock_glob.return_value = []
|
|
mock_list.return_value = [f"{i}" for i in range(4)]
|
|
TPUAcceleratorManager.get_current_node_num_accelerators.cache_clear()
|
|
assert TPUAcceleratorManager.get_current_node_num_accelerators() == 4
|
|
|
|
|
|
@patch("glob.glob")
|
|
@patch("os.listdir")
|
|
def test_autodetect_num_tpus_without_devices(mock_list, mock_glob):
|
|
mock_list.side_effect = FileNotFoundError
|
|
mock_glob.return_value = []
|
|
TPUAcceleratorManager.get_current_node_num_accelerators.cache_clear()
|
|
assert TPUAcceleratorManager.get_current_node_num_accelerators() == 0
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"accelerator_type_version_tuple",
|
|
[
|
|
("gce", "v2-8", "TPU-V2"),
|
|
("gce", "v2-32", "TPU-V2"),
|
|
("gce", "v3-8", "TPU-V3"),
|
|
("gce", "v3-128", "TPU-V3"),
|
|
("gce", "v4-8", "TPU-V4"),
|
|
("gce", "v4-2048", "TPU-V4"),
|
|
("gce", "v5p-8", "TPU-V5P"),
|
|
("gce", "v5litepod-8", "TPU-V5LITEPOD"),
|
|
("gce", "v6e-8", "TPU-V6E"),
|
|
("gke", "v2-8", "TPU-V2"),
|
|
("gke", "v2-32", "TPU-V2"),
|
|
("gke", "v3-8", "TPU-V3"),
|
|
("gke", "v3-128", "TPU-V3"),
|
|
("gke", "v4-8", "TPU-V4"),
|
|
("gke", "v4-2048", "TPU-V4"),
|
|
("gke", "v5p-8", "TPU-V5P"),
|
|
("gke", "v5litepod-8", "TPU-V5LITEPOD"),
|
|
("gke", "v6e-8", "TPU-V6E"),
|
|
("gke", "tpu7x-16", "TPU-V7X"),
|
|
],
|
|
)
|
|
@patch("requests.get")
|
|
@patch("os.getenv")
|
|
def test_autodetect_tpu_accelerator_type(
|
|
mock_os, mock_request, accelerator_type_version_tuple
|
|
):
|
|
gce_or_gke, accelerator_type, expected_version = accelerator_type_version_tuple
|
|
if gce_or_gke == "gce":
|
|
mock_response = mock.MagicMock()
|
|
mock_response.status_code = 200
|
|
mock_response.text = accelerator_type
|
|
mock_request.return_value = mock_response
|
|
mock_os.return_value = None
|
|
else:
|
|
mock_os.return_value = accelerator_type
|
|
assert TPUAcceleratorManager.get_current_node_accelerator_type() == expected_version
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"test_case",
|
|
[
|
|
("gce", "0", 0),
|
|
("gke", "0", 0),
|
|
],
|
|
)
|
|
@patch("requests.get")
|
|
@patch("os.getenv")
|
|
def test_get_current_node_tpu_worker_id(mock_os, mock_request, test_case):
|
|
gce_or_gke, worker_id, expected_value = test_case
|
|
if gce_or_gke == "gce":
|
|
mock_response = mock.MagicMock()
|
|
mock_response.status_code = 200
|
|
mock_response.text = worker_id
|
|
mock_request.return_value = mock_response
|
|
mock_os.return_value = None
|
|
else:
|
|
mock_os.return_value = worker_id
|
|
assert TPUAcceleratorManager.get_current_node_tpu_worker_id() == expected_value
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"test_case",
|
|
[
|
|
("gce", "my-tpu"),
|
|
("gke", "my-tpu"),
|
|
],
|
|
)
|
|
@patch("requests.get")
|
|
@patch("os.getenv")
|
|
def test_get_tpu_unique_id(mock_os, mock_request, test_case):
|
|
gce_or_gke, worker_id = test_case
|
|
if gce_or_gke == "gce":
|
|
mock_response = mock.MagicMock()
|
|
mock_response.status_code = 200
|
|
mock_response.text = worker_id
|
|
mock_request.return_value = mock_response
|
|
mock_os.return_value = None
|
|
else:
|
|
mock_os.return_value = worker_id
|
|
assert TPUAcceleratorManager.get_current_node_tpu_name() == worker_id
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"test_case",
|
|
[
|
|
("gce", "not-a-valid-version"),
|
|
("gce", "vNOTVALID-8"),
|
|
("gce", "230498230948230948"),
|
|
# From issue #39913
|
|
("gce", ""),
|
|
("gke", "not-a-valid-version"),
|
|
("gke", "vNOTVALID-8"),
|
|
("gke", "230498230948230948"),
|
|
],
|
|
)
|
|
@patch("requests.get")
|
|
@patch("os.getenv")
|
|
def test_autodetect_invalid_type(mock_os, mock_request, test_case):
|
|
gce_or_gke, accelerator_type = test_case
|
|
if gce_or_gke == "gce":
|
|
mock_response = mock.MagicMock()
|
|
mock_response.status_code = 200
|
|
mock_response.text = accelerator_type
|
|
mock_request.return_value = mock_response
|
|
mock_os.return_value = None
|
|
else:
|
|
mock_os.return_value = accelerator_type
|
|
assert TPUAcceleratorManager.get_current_node_accelerator_type() is None
|
|
|
|
|
|
def test_autodetect_tpu_accelerator_type_fails_gracefully():
|
|
with patch("requests.get") as mock_get:
|
|
mock_get.side_effect = requests.exceptions.RequestException
|
|
assert TPUAcceleratorManager.get_current_node_accelerator_type() is None
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"test_config",
|
|
[
|
|
(1, False),
|
|
(0.5, True),
|
|
(3, True),
|
|
],
|
|
)
|
|
def test_validate_resource_request_quantity(test_config):
|
|
num_tpus, expect_error = test_config
|
|
|
|
if expect_error:
|
|
assert (
|
|
TPUAcceleratorManager.validate_resource_request_quantity(num_tpus)[0]
|
|
is False
|
|
)
|
|
assert (
|
|
TPUAcceleratorManager.validate_resource_request_quantity(num_tpus)[1]
|
|
is not None
|
|
)
|
|
else:
|
|
assert (
|
|
TPUAcceleratorManager.validate_resource_request_quantity(num_tpus)[0]
|
|
is True
|
|
)
|
|
assert (
|
|
TPUAcceleratorManager.validate_resource_request_quantity(num_tpus)[1]
|
|
is None
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"test_case",
|
|
[
|
|
(4, ["0"]),
|
|
(4, ["0", "1"]),
|
|
(4, ["0", "1", "2", "3"]),
|
|
(8, ["0", "1", "2", "3", "4", "5", "6", "7"]),
|
|
],
|
|
)
|
|
@patch("glob.glob")
|
|
def test_set_tpu_visible_ids_and_bounds(mock_glob, test_case):
|
|
num_devices, tpu_chips = test_case
|
|
mock_glob.return_value = ["/dev/accel" + str(x) for x in range(num_devices)]
|
|
with patch.dict("os.environ", {}, clear=True):
|
|
TPUAcceleratorManager.get_current_node_num_accelerators.cache_clear()
|
|
TPUAcceleratorManager.set_current_process_visible_accelerator_ids(tpu_chips)
|
|
if len(tpu_chips) == 1:
|
|
assert (
|
|
os.environ[tpu.TPU_CHIPS_PER_HOST_BOUNDS_ENV_VAR]
|
|
== tpu.TPU_CHIPS_PER_HOST_BOUNDS_1_CHIP_CONFIG
|
|
)
|
|
assert os.environ[tpu.TPU_HOST_BOUNDS_ENV_VAR] == tpu.TPU_SINGLE_HOST_BOUNDS
|
|
assert os.environ[tpu.TPU_VISIBLE_CHIPS_ENV_VAR] == ",".join(tpu_chips)
|
|
elif len(tpu_chips) == 2:
|
|
assert (
|
|
os.environ[tpu.TPU_CHIPS_PER_HOST_BOUNDS_ENV_VAR]
|
|
== tpu.TPU_CHIPS_PER_HOST_BOUNDS_2_CHIP_CONFIG
|
|
)
|
|
assert os.environ[tpu.TPU_HOST_BOUNDS_ENV_VAR] == tpu.TPU_SINGLE_HOST_BOUNDS
|
|
assert os.environ[tpu.TPU_VISIBLE_CHIPS_ENV_VAR] == ",".join(tpu_chips)
|
|
elif len(tpu_chips) == 4:
|
|
# Check that nothing is set, let the ML framework use the defaults.
|
|
assert os.environ.get(tpu.TPU_CHIPS_PER_HOST_BOUNDS_ENV_VAR, None) is None
|
|
assert os.environ.get(tpu.TPU_SINGLE_HOST_BOUNDS, None) is None
|
|
assert os.environ.get(tpu.TPU_VISIBLE_CHIPS_ENV_VAR, None) is None
|
|
else: # len(tpu_chips) == 8
|
|
assert os.environ.get(tpu.TPU_CHIPS_PER_HOST_BOUNDS_ENV_VAR, None) is None
|
|
assert os.environ.get(tpu.TPU_SINGLE_HOST_BOUNDS, None) is None
|
|
assert os.environ.get(tpu.TPU_VISIBLE_CHIPS_ENV_VAR, None) is None
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"test_config",
|
|
[
|
|
(0, "v4-16", {"TPU-v4-16-head": 1, "my-tpu": 1}),
|
|
(1, "v4-16", {"my-tpu": 1}),
|
|
(0, "tpu7x-16", {"TPU-v7x-16-head": 1, "my-tpu": 1}),
|
|
],
|
|
)
|
|
def test_tpu_pod_detect_and_configure_worker(test_config):
|
|
worker_id, pod_type, expected_value = test_config
|
|
final_resources = {}
|
|
with patch(
|
|
"ray._private.accelerators.tpu.TPUAcceleratorManager.get_current_node_tpu_name",
|
|
return_value="my-tpu",
|
|
):
|
|
with patch(
|
|
"ray._private.accelerators.tpu.TPUAcceleratorManager.get_current_node_tpu_worker_id",
|
|
return_value=worker_id,
|
|
):
|
|
with patch.dict(os.environ, {"TPU_ACCELERATOR_TYPE": pod_type}):
|
|
final_resources = (
|
|
TPUAcceleratorManager.get_current_node_additional_resources()
|
|
)
|
|
|
|
assert final_resources == expected_value
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"accelerator_type, expected",
|
|
[
|
|
("v2-8", True),
|
|
("v3-32", True),
|
|
("v4-8", True),
|
|
("v5p-8", True),
|
|
("v5litepod-8", True),
|
|
("v6e-8", True),
|
|
("tpu7x-16", True),
|
|
("v7x-16", True),
|
|
("v-8", False),
|
|
("8", False),
|
|
("tpu-8", False),
|
|
("v2", False),
|
|
("v2-", False),
|
|
("random-string", False),
|
|
],
|
|
)
|
|
def test_is_valid_tpu_accelerator_type(accelerator_type, expected):
|
|
assert (
|
|
TPUAcceleratorManager.is_valid_tpu_accelerator_type(accelerator_type)
|
|
== expected
|
|
)
|
|
|
|
|
|
def test_get_total_chips_from_accelerator_type():
|
|
assert tpu.get_total_chips_from_accelerator_type("v6e-16") == 16
|
|
assert tpu.get_total_chips_from_accelerator_type("v6e-8") == 8
|
|
assert (
|
|
tpu.get_total_chips_from_accelerator_type("v7x-16") == 8
|
|
) # v7x has 2 cores per chip
|
|
assert (
|
|
tpu.get_total_chips_from_accelerator_type("v4-8") == 4
|
|
) # v4 has 2 cores per chip
|
|
|
|
# Test invalid cases
|
|
with pytest.raises(ValueError, match="Accelerator type must include size"):
|
|
tpu.get_total_chips_from_accelerator_type("v6e")
|
|
|
|
with pytest.raises(ValueError, match="Invalid accelerator type"):
|
|
tpu.get_total_chips_from_accelerator_type("invalid-8")
|
|
|
|
|
|
def test_get_num_tpu_visible_chips_per_host():
|
|
# v6e multi-host (4 chips per VM)
|
|
assert tpu.get_num_tpu_visible_chips_per_host("v6e-16") == 4
|
|
assert tpu.get_num_tpu_visible_chips_per_host("v6e-32") == 4
|
|
|
|
# v6e single-host/sub-host (exact chip count)
|
|
assert tpu.get_num_tpu_visible_chips_per_host("v6e-8") == 8
|
|
assert tpu.get_num_tpu_visible_chips_per_host("v6e-4") == 4
|
|
assert tpu.get_num_tpu_visible_chips_per_host("v6e-1") == 1
|
|
|
|
# v5litepod multi-host defaults to 4, single-host is 8 chips
|
|
assert tpu.get_num_tpu_visible_chips_per_host("v5litepod-16") == 4
|
|
assert tpu.get_num_tpu_visible_chips_per_host("v5litepod-8") == 8
|
|
|
|
# v5litepod sub-host
|
|
assert tpu.get_num_tpu_visible_chips_per_host("v5litepod-4") == 4
|
|
assert tpu.get_num_tpu_visible_chips_per_host("v5litepod-1") == 1
|
|
|
|
# Other TPU generations default to 4
|
|
assert tpu.get_num_tpu_visible_chips_per_host("v4-8") == 4
|
|
assert tpu.get_num_tpu_visible_chips_per_host("v5p-8") == 4
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-sv", __file__]))
|