wehub-resource-sync
972 lines
33 KiB
Python
972 lines
33 KiB
Python
import sys
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
import ray
|
|
from ray._private.accelerators import TPUAcceleratorManager, tpu
|
|
from ray.util.tpu import SlicePlacementGroup
|
|
|
|
|
|
def test_get_current_pod_name_smoke():
|
|
with patch(
|
|
"ray._private.accelerators.tpu.TPUAcceleratorManager.get_current_node_tpu_name",
|
|
return_value="my-tpu",
|
|
):
|
|
name = ray.util.tpu.get_current_pod_name()
|
|
assert name == "my-tpu"
|
|
|
|
|
|
def test_empty_get_current_pod_name_returns_none():
|
|
with patch(
|
|
"ray._private.accelerators.tpu.TPUAcceleratorManager.get_current_node_tpu_name",
|
|
return_value="",
|
|
):
|
|
name = ray.util.tpu.get_current_pod_name()
|
|
assert name is None
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"test_case",
|
|
[
|
|
# (number_chips_per_host, parsed accl_type, expected_worker_count)
|
|
(4, "v2-4", 1),
|
|
(4, "v3-32", 4),
|
|
(4, "v4-8", 1),
|
|
(4, "v4-16", 2),
|
|
(8, "v5litepod-4", 1),
|
|
(8, "v5litepod-8", 1),
|
|
(8, "v5litepod-16", 2),
|
|
(8, "v5litepod-32", 4),
|
|
(4, "v5p-4", 1),
|
|
(4, "v5p-8", 1),
|
|
(4, "v5p-16", 2),
|
|
(4, "v6e-4", 1),
|
|
(8, "v6e-8", 1),
|
|
(4, "v6e-8", 2),
|
|
(8, "v6e-16", 2),
|
|
(4, "v7x-8", 1),
|
|
(4, "v7x-16", 2),
|
|
],
|
|
)
|
|
@patch("glob.glob")
|
|
def test_worker_count(mock_glob, test_case):
|
|
num_devices, accelerator_type, expected_worker_count = test_case
|
|
mock_glob.return_value = ["/dev/accel" + str(x) for x in range(num_devices)]
|
|
TPUAcceleratorManager.get_current_node_num_accelerators.cache_clear()
|
|
|
|
with patch(
|
|
"ray._private.accelerators.tpu.TPUAcceleratorManager."
|
|
"get_current_node_tpu_pod_type",
|
|
return_value=accelerator_type,
|
|
):
|
|
worker_count = ray.util.tpu.get_current_pod_worker_count()
|
|
|
|
assert worker_count == expected_worker_count
|
|
|
|
|
|
@patch("glob.glob")
|
|
def test_num_tpu_chips(mock_glob):
|
|
mock_glob.return_value = [
|
|
"/dev/accel0",
|
|
"/dev/accel1",
|
|
"/dev/accel2",
|
|
"/dev/accel3",
|
|
]
|
|
TPUAcceleratorManager.get_current_node_num_accelerators.cache_clear()
|
|
num_tpu_chips = ray.util.tpu.get_num_tpu_chips_on_node()
|
|
assert num_tpu_chips == 4
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"test_case",
|
|
[
|
|
# (accelerator_type, accelerator_topology, expected_result)
|
|
("v2-16", "4x4", True),
|
|
("v2-256", "16x16", True),
|
|
("v2-4", "2x2", False),
|
|
("v3-16", "4x4", True),
|
|
("v3-1024", "32x32", True),
|
|
("v3-4", "4x16", False),
|
|
("v4-4", "2x2x1", True),
|
|
("v4-32", "2x4x4", True),
|
|
("v4-2048", "8x8x16", True),
|
|
("v4-4", "16x16x16", False),
|
|
("v5p-128", "4x4x4", True),
|
|
("v5p-4096", "16x16x16", True),
|
|
("v5p-12288", "16x16x24", True),
|
|
("v5p-4", "24x24x24", False),
|
|
("v5litepod-16", "2x8", True),
|
|
("v5litepod-256", "16x16", True),
|
|
("v5litepod-4", "2x2", True),
|
|
("v6e-8", "2x4", True),
|
|
("v6e-16", "4x4", True),
|
|
("v6e-64", "8x8", True),
|
|
("v6e-4", "4x16", False),
|
|
("tpu7x-16", "2x2x2", True),
|
|
("tpu7x-64", "2x4x4", True),
|
|
("v7x-8", "4x4", False),
|
|
],
|
|
)
|
|
@patch("glob.glob")
|
|
def test_is_valid_tpu_accelerator_topology(_mock_glob, test_case):
|
|
"""Test valid TPU accelerator topologies."""
|
|
accelerator_type, accelerator_topology, expected_result = test_case
|
|
actual_result = TPUAcceleratorManager.is_valid_tpu_accelerator_topology(
|
|
accelerator_type, accelerator_topology
|
|
)
|
|
|
|
assert actual_result == expected_result
|
|
|
|
|
|
def test_get_current_node_labels_env_only(monkeypatch):
|
|
# Simulate GKE TPU environment variables
|
|
monkeypatch.setenv("TPU_NAME", "tpu-worker-group-2")
|
|
monkeypatch.setenv("TPU_WORKER_ID", "0")
|
|
monkeypatch.setenv("TPU_ACCELERATOR_TYPE", "v6e-16")
|
|
monkeypatch.setenv("TPU_TOPOLOGY", "4x4")
|
|
|
|
tpu_labels = TPUAcceleratorManager.get_current_node_accelerator_labels()
|
|
|
|
assert tpu_labels["ray.io/tpu-slice-name"] == "tpu-worker-group-2"
|
|
assert tpu_labels["ray.io/tpu-worker-id"] == "0"
|
|
assert tpu_labels["ray.io/tpu-topology"] == "4x4"
|
|
assert tpu_labels["ray.io/tpu-pod-type"] == "v6e-16"
|
|
|
|
|
|
def test_get_current_node_tpu_topology_from_metadata():
|
|
tpu_env_string = "TPU_ACCELERATOR:v6e.\nTOPOLOGY: '2x2x4'\nTPU_HOST_BOUNDS:0,1,1,2"
|
|
|
|
with patch(
|
|
"ray._private.accelerators.tpu._get_tpu_metadata", return_value=tpu_env_string
|
|
):
|
|
topology = TPUAcceleratorManager.get_current_node_tpu_topology()
|
|
assert topology == "2x2x4"
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"topology, accelerator_type, expected_pod_type, should_raise",
|
|
[
|
|
("2x4", "TPU-V6E", "v6e-8", False),
|
|
("2x2x2", "TPU-V4", "v4-16", False),
|
|
("4x8", "TPU-V3", "v3-64", False),
|
|
("2x2x1", "TPU-V5P", "v5p-8", False),
|
|
("4x4", "TPU-V5P", "v5p-32", False),
|
|
("8x16", "TPU-V6E", "v6e-128", False),
|
|
("", "TPU-V3", None, False),
|
|
("4x", "TPU-V3", None, True),
|
|
("2x2x2", "TPU-V7X", "v7x-16", False),
|
|
],
|
|
)
|
|
def test_infer_tpu_pod_type_from_topology(
|
|
topology, accelerator_type, expected_pod_type, should_raise
|
|
):
|
|
if should_raise:
|
|
with pytest.raises(ValueError):
|
|
tpu.infer_tpu_pod_type_from_topology(topology, accelerator_type)
|
|
else:
|
|
actual_result = tpu.infer_tpu_pod_type_from_topology(topology, accelerator_type)
|
|
assert actual_result == expected_pod_type
|
|
|
|
|
|
@pytest.fixture
|
|
def ray_start_cpu():
|
|
address_info = ray.init(num_cpus=1)
|
|
yield address_info
|
|
ray.shutdown()
|
|
|
|
|
|
@pytest.fixture
|
|
def ray_tpu_cluster(ray_start_cluster):
|
|
"""
|
|
Simulates a Ray cluster with two multi-host TPU v4-16 slices.
|
|
"""
|
|
pod_type = "v4-16"
|
|
topology = "2x2x2"
|
|
|
|
cluster = ray_start_cluster
|
|
slice_0_env_common = {
|
|
"TPU_NAME": "test-slice-0",
|
|
"TPU_ACCELERATOR_TYPE": pod_type,
|
|
"TPU_TOPOLOGY": topology,
|
|
}
|
|
slice_0_head_labels = {
|
|
"ray.io/tpu-slice-name": "test-slice-0",
|
|
"ray.io/tpu-worker-id": "0",
|
|
"ray.io/tpu-pod-type": pod_type,
|
|
"ray.io/tpu-topology": topology,
|
|
}
|
|
slice_0_worker_labels = {
|
|
"ray.io/tpu-slice-name": "test-slice-0",
|
|
"ray.io/tpu-worker-id": "1",
|
|
"ray.io/tpu-pod-type": pod_type,
|
|
"ray.io/tpu-topology": topology,
|
|
}
|
|
cluster.add_node(
|
|
num_cpus=2,
|
|
resources={"TPU": 4, f"TPU-{pod_type}-head": 1},
|
|
env_vars={**slice_0_env_common, "TPU_WORKER_ID": "0"},
|
|
labels=slice_0_head_labels,
|
|
)
|
|
cluster.add_node(
|
|
num_cpus=2,
|
|
resources={"TPU": 4},
|
|
env_vars={**slice_0_env_common, "TPU_WORKER_ID": "1"},
|
|
labels=slice_0_worker_labels,
|
|
)
|
|
|
|
slice_1_env_common = {
|
|
"TPU_NAME": "test-slice-1",
|
|
"TPU_ACCELERATOR_TYPE": pod_type,
|
|
"TPU_TOPOLOGY": topology,
|
|
}
|
|
slice_1_head_labels = {
|
|
"ray.io/tpu-slice-name": "test-slice-1",
|
|
"ray.io/tpu-worker-id": "0",
|
|
"ray.io/tpu-pod-type": pod_type,
|
|
"ray.io/tpu-topology": topology,
|
|
}
|
|
slice_1_worker_labels = {
|
|
"ray.io/tpu-slice-name": "test-slice-1",
|
|
"ray.io/tpu-worker-id": "1",
|
|
"ray.io/tpu-pod-type": pod_type,
|
|
"ray.io/tpu-topology": topology,
|
|
}
|
|
cluster.add_node(
|
|
num_cpus=2,
|
|
resources={"TPU": 4, f"TPU-{pod_type}-head": 1},
|
|
env_vars={**slice_1_env_common, "TPU_WORKER_ID": "0"},
|
|
labels=slice_1_head_labels,
|
|
)
|
|
cluster.add_node(
|
|
num_cpus=2,
|
|
resources={"TPU": 4},
|
|
env_vars={**slice_1_env_common, "TPU_WORKER_ID": "1"},
|
|
labels=slice_1_worker_labels,
|
|
)
|
|
|
|
ray.init(address=cluster.address)
|
|
yield cluster
|
|
ray.shutdown()
|
|
|
|
|
|
@pytest.fixture
|
|
def ray_v6e_tpu_cluster(ray_start_cluster):
|
|
"""
|
|
Simulates a Ray cluster with two v6e-8 slices (2x4 topology).
|
|
|
|
"""
|
|
pod_type = "v6e-8"
|
|
topology = "2x4"
|
|
cluster = ray_start_cluster
|
|
|
|
for i in range(2):
|
|
env_common = {
|
|
"TPU_NAME": f"test-v6e-slice-{i}",
|
|
"TPU_ACCELERATOR_TYPE": pod_type,
|
|
"TPU_TOPOLOGY": topology,
|
|
}
|
|
head_labels = {
|
|
"ray.io/tpu-slice-name": f"test-v6e-slice-{i}",
|
|
"ray.io/tpu-worker-id": "0",
|
|
"ray.io/tpu-pod-type": pod_type,
|
|
"ray.io/tpu-topology": topology,
|
|
}
|
|
# A single-host v6e-8 has 8 chips on one node
|
|
cluster.add_node(
|
|
num_cpus=4,
|
|
resources={"TPU": 8, f"TPU-{pod_type}-head": 1},
|
|
env_vars={**env_common, "TPU_WORKER_ID": "0"},
|
|
labels=head_labels,
|
|
)
|
|
|
|
ray.init(address=cluster.address)
|
|
yield cluster
|
|
ray.shutdown()
|
|
|
|
|
|
def test_fetch_tpu_slice_name_from_pg(ray_tpu_cluster):
|
|
"""Tests that the slice name can be fetched from a PG."""
|
|
tpu_head_pg = ray.util.placement_group(bundles=[{"TPU-v4-16-head": 1}])
|
|
ray.get(tpu_head_pg.ready())
|
|
|
|
expected_unique_slice_names = {"test-slice-0", "test-slice-1"}
|
|
slice_name = tpu.fetch_tpu_slice_name_from_pg(tpu_head_pg)
|
|
assert slice_name in expected_unique_slice_names
|
|
|
|
ray.util.remove_placement_group(tpu_head_pg)
|
|
|
|
|
|
def test_reserve_tpu_slice(ray_tpu_cluster):
|
|
"""Tests that a TPU slice can be successfully reserved."""
|
|
reserved_name_0, hg_pg_0 = tpu.reserve_tpu_slice(
|
|
topology="2x2x2", accelerator_type="TPU-V4"
|
|
)
|
|
reserved_name_1, hg_pg_1 = tpu.reserve_tpu_slice(
|
|
topology="2x2x2", accelerator_type="TPU-V4"
|
|
)
|
|
|
|
# Ensure the placement groups reserving the TPU slice using the head worker are valid.
|
|
assert hg_pg_0 is not None, "Expected placement group for slice 0, got None"
|
|
assert hg_pg_1 is not None, "Expected placement group for slice 1, got None"
|
|
|
|
assert (
|
|
reserved_name_0 != reserved_name_1
|
|
), f"Expected to reserve two different slices, but got the same name: {reserved_name_0}"
|
|
expected_unique_slice_names = {"test-slice-0", "test-slice-1"}
|
|
actual_reserved_names = {reserved_name_0, reserved_name_1}
|
|
assert actual_reserved_names == expected_unique_slice_names, (
|
|
f"Got unexpected slice names. Expected {expected_unique_slice_names}, "
|
|
f"but got {actual_reserved_names}"
|
|
)
|
|
|
|
|
|
def test_slice_placement_group(ray_tpu_cluster):
|
|
"""Test that single TPU slice can be successfully reserved."""
|
|
slice_placement_group = ray.util.tpu.slice_placement_group(
|
|
topology="2x2x2",
|
|
accelerator_version="v4",
|
|
)
|
|
assert slice_placement_group.chips_per_host == 4
|
|
assert slice_placement_group.num_hosts == 2
|
|
assert slice_placement_group.placement_group.bundle_count == 2
|
|
assert slice_placement_group.placement_group.bundle_specs == [
|
|
{"TPU": 4, "CPU": 1.0},
|
|
{"TPU": 4, "CPU": 1.0},
|
|
]
|
|
|
|
|
|
def test_multi_slice_placement_group(ray_tpu_cluster):
|
|
"""Test that multiple whole TPU slices can be successfully reserved"""
|
|
multi_slice_placement_group = ray.util.tpu.slice_placement_group(
|
|
topology="2x2x2",
|
|
accelerator_version="v4",
|
|
num_slices=2,
|
|
)
|
|
assert multi_slice_placement_group.placement_group.bundle_count == 4
|
|
assert multi_slice_placement_group.num_hosts == 4
|
|
assert multi_slice_placement_group.placement_group.bundle_specs == [
|
|
{"TPU": 4, "CPU": 1.0}, # slice 1, host 1
|
|
{"TPU": 4, "CPU": 1.0}, # slice 1, host 2
|
|
{"TPU": 4, "CPU": 1.0}, # slice 2, host 1
|
|
{"TPU": 4, "CPU": 1.0}, # slice 2, host 2
|
|
]
|
|
|
|
|
|
@patch("ray.util.tpu.placement_group")
|
|
@patch("ray.util.tpu.remove_placement_group")
|
|
@patch("ray.util.tpu.reserve_tpu_slice")
|
|
def test_slice_placement_group_partial_failure_cleanup(
|
|
mock_reserve, mock_remove_pg, mock_create_pg
|
|
):
|
|
"""
|
|
Verifies that if a multi-slice request fails halfway through,
|
|
the TPU head placement groups are cleaned up to prevent leaks.
|
|
"""
|
|
fake_head_pg_1 = MagicMock(name="head_pg_1")
|
|
mock_reserve.side_effect = [("slice_1", fake_head_pg_1), None]
|
|
|
|
with pytest.raises(RuntimeError, match="Failed to reserve TPU slice"):
|
|
SlicePlacementGroup(topology="2x2x2", accelerator_version="v4", num_slices=2)
|
|
|
|
# Validate that 2 TPU util attempted to reserve two slices, failed, and
|
|
# correctly cleaned up the hanging TPU head placement groups.
|
|
assert mock_reserve.call_count == 2
|
|
mock_remove_pg.assert_called_once_with(fake_head_pg_1)
|
|
mock_create_pg.assert_not_called()
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"accelerator_type, expected_version",
|
|
[
|
|
# type with "TPU-" prefix
|
|
("TPU-V4", "v4"),
|
|
("TPU-v4", "v4"),
|
|
("TPU-V6E", "v6e"),
|
|
("TPU-v5p", "v5p"),
|
|
("TPU-V7X", "v7x"),
|
|
# Only the TPU version - no parsing necessary.
|
|
("v4", "v4"),
|
|
("v3", "v3"),
|
|
("v6e", "v6e"),
|
|
("v5litepod", "v5litepod"),
|
|
("v7x", "v7x"),
|
|
],
|
|
)
|
|
def test_get_tpu_version_valid(accelerator_type, expected_version):
|
|
assert ray.util.tpu.get_tpu_version_from_type(accelerator_type) == expected_version
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"invalid_type",
|
|
[
|
|
"A100", # GPU type
|
|
"random-invalid-type", # Random string
|
|
"TPU-invalid", # TPU prefix
|
|
"", # Empty string
|
|
],
|
|
)
|
|
def test_get_tpu_version_invalid(invalid_type):
|
|
with pytest.raises(ValueError, match="Invalid accelerator_type"):
|
|
ray.util.tpu.get_tpu_version_from_type(invalid_type)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"topology, accelerator_type, num_workers, resources_per_worker, expected_slices",
|
|
[
|
|
# "2x2x1" has 4 chips, for 4 workers with TPU: 1 each we expect num_slices=1.
|
|
("2x2x1", "TPU-V4", 4, {"TPU": 1}, 1),
|
|
# "2x2x1" has 4 chips, for 8 workers with TPU: 1 each we expect num_slices=2.
|
|
("2x2x1", "v4", 8, {"TPU": 1}, 2),
|
|
# "2x2x2" has 8 chips and 2 hosts, defaulting to 1 TPU worker per host
|
|
# and requesting 4 workers, we expect num_slices=2.
|
|
("2x2x2", "TPU-V4", 4, None, 2),
|
|
# "2x2x4" has 16 chips and 4 hosts, defaulting to 1 TPU worker per host
|
|
# and requesting 4 workers, we expect num_slices=1.
|
|
("2x2x4", "TPU-V4", 4, None, 1),
|
|
# 0 workers requested -> fallback to 1 slice.
|
|
("2x2x1", "v4", 0, None, 1),
|
|
# Invalid topology -> fallback to 1 slice.
|
|
("", "v4", 4, {"TPU": 1}, 1),
|
|
("2x2x1", "", 4, {"TPU": 1}, 1),
|
|
],
|
|
)
|
|
def test_get_tpu_num_slices_for_workers(
|
|
topology, accelerator_type, num_workers, resources_per_worker, expected_slices
|
|
):
|
|
num_slices = ray.util.tpu.get_tpu_num_slices_for_workers(
|
|
topology=topology,
|
|
accelerator_type=accelerator_type,
|
|
num_workers=num_workers,
|
|
resources_per_worker=resources_per_worker,
|
|
)
|
|
assert num_slices == expected_slices
|
|
|
|
|
|
def _make_mock_tpu_node(
|
|
alive, pod_type, slice_name, worker_id, tpu_chips=4, node_id=None
|
|
):
|
|
"""Helper to mock a Ray Node dictionary returned by ray.nodes()."""
|
|
if node_id is None:
|
|
node_id = f"node_{slice_name}_{worker_id}"
|
|
return {
|
|
"NodeID": node_id,
|
|
"Alive": alive,
|
|
"Labels": {
|
|
"ray.io/tpu-pod-type": pod_type,
|
|
"ray.io/tpu-slice-name": slice_name,
|
|
"ray.io/tpu-worker-id": str(worker_id),
|
|
},
|
|
"Resources": {"TPU": tpu_chips},
|
|
}
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"topology, accelerator_type, mock_nodes, mock_avail_resources, expected_ready",
|
|
[
|
|
# 1 fully intact and available v4 slice (2 physical hosts).
|
|
(
|
|
"2x2x2",
|
|
"v4",
|
|
[
|
|
_make_mock_tpu_node(True, "v4-16", "slice-1", 0, node_id="A"),
|
|
_make_mock_tpu_node(True, "v4-16", "slice-1", 1, node_id="B"),
|
|
],
|
|
{
|
|
"A": {"TPU": 4},
|
|
"B": {"TPU": 4},
|
|
},
|
|
1,
|
|
),
|
|
# 1 fully intact slice, but one node is using 2 TPUs (unavailable) -> 0 ready slices.
|
|
(
|
|
"2x2x2",
|
|
"v4",
|
|
[
|
|
_make_mock_tpu_node(True, "v4-16", "slice-1", 0, node_id="A"),
|
|
_make_mock_tpu_node(True, "v4-16", "slice-1", 1, node_id="B"),
|
|
],
|
|
{
|
|
"A": {"TPU": 2},
|
|
"B": {"TPU": 4},
|
|
},
|
|
0,
|
|
),
|
|
# Fractured slice (missing a physical host) -> 0 ready slices.
|
|
(
|
|
"2x2x2",
|
|
"v4",
|
|
[
|
|
_make_mock_tpu_node(True, "v4-16", "slice-1", 0, node_id="A"),
|
|
# Worker 1 is missing
|
|
],
|
|
{
|
|
"A": {"TPU": 4},
|
|
},
|
|
0,
|
|
),
|
|
# Correct number of hosts, but missing the head node (rank 0) -> 0 ready slices.
|
|
(
|
|
"2x2x2",
|
|
"v4",
|
|
[
|
|
_make_mock_tpu_node(True, "v4-16", "slice-1", 1, node_id="A"),
|
|
_make_mock_tpu_node(True, "v4-16", "slice-1", 2, node_id="B"),
|
|
],
|
|
{
|
|
"A": {"TPU": 4},
|
|
"B": {"TPU": 4},
|
|
},
|
|
0,
|
|
),
|
|
# Fractured slice (one physical host is dead) -> 0 ready slices.
|
|
(
|
|
"2x2x2",
|
|
"v4",
|
|
[
|
|
_make_mock_tpu_node(True, "v4-16", "slice-1", 0, node_id="A"),
|
|
_make_mock_tpu_node(False, "v4-16", "slice-1", 1, node_id="B"),
|
|
],
|
|
{
|
|
"A": {"TPU": 4},
|
|
},
|
|
0,
|
|
),
|
|
# 2 slices total: one intact & available, one fractured -> 1 ready slice.
|
|
(
|
|
"2x2x2",
|
|
"v4",
|
|
[
|
|
_make_mock_tpu_node(True, "v4-16", "slice-A", 0, node_id="A0"),
|
|
_make_mock_tpu_node(True, "v4-16", "slice-A", 1, node_id="A1"),
|
|
_make_mock_tpu_node(True, "v4-16", "slice-B", 0, node_id="B0"),
|
|
# slice-B worker 1 is missing
|
|
],
|
|
{
|
|
"A0": {"TPU": 4},
|
|
"A1": {"TPU": 4},
|
|
"B0": {"TPU": 4},
|
|
},
|
|
1,
|
|
),
|
|
# 1 fully intact and available v6e 2x4 slice (single-host).
|
|
(
|
|
"2x4",
|
|
"v6e",
|
|
[
|
|
_make_mock_tpu_node(
|
|
True, "v6e-8", "slice-1", 0, tpu_chips=8, node_id="A"
|
|
),
|
|
],
|
|
{
|
|
"A": {"TPU": 8},
|
|
},
|
|
1,
|
|
),
|
|
# 1 fully intact and available v6e 2x4 slice (2 physical hosts).
|
|
(
|
|
"2x4",
|
|
"v6e",
|
|
[
|
|
_make_mock_tpu_node(
|
|
True, "v6e-8", "slice-1", 0, tpu_chips=4, node_id="A"
|
|
),
|
|
_make_mock_tpu_node(
|
|
True, "v6e-8", "slice-1", 1, tpu_chips=4, node_id="B"
|
|
),
|
|
],
|
|
{
|
|
"A": {"TPU": 4},
|
|
"B": {"TPU": 4},
|
|
},
|
|
1,
|
|
),
|
|
# 2 fully intact v6e slices.
|
|
(
|
|
"4x4",
|
|
"v6e",
|
|
[
|
|
_make_mock_tpu_node(True, "v6e-16", "slice-1", 0, node_id="S1_0"),
|
|
_make_mock_tpu_node(True, "v6e-16", "slice-1", 1, node_id="S1_1"),
|
|
_make_mock_tpu_node(True, "v6e-16", "slice-1", 2, node_id="S1_2"),
|
|
_make_mock_tpu_node(True, "v6e-16", "slice-1", 3, node_id="S1_3"),
|
|
_make_mock_tpu_node(True, "v6e-16", "slice-2", 0, node_id="S2_0"),
|
|
_make_mock_tpu_node(True, "v6e-16", "slice-2", 1, node_id="S2_1"),
|
|
_make_mock_tpu_node(True, "v6e-16", "slice-2", 2, node_id="S2_2"),
|
|
_make_mock_tpu_node(True, "v6e-16", "slice-2", 3, node_id="S2_3"),
|
|
],
|
|
{
|
|
"S1_0": {"TPU": 4},
|
|
"S1_1": {"TPU": 4},
|
|
"S1_2": {"TPU": 4},
|
|
"S1_3": {"TPU": 4},
|
|
"S2_0": {"TPU": 4},
|
|
"S2_1": {"TPU": 4},
|
|
"S2_2": {"TPU": 4},
|
|
"S2_3": {"TPU": 4},
|
|
},
|
|
2,
|
|
),
|
|
],
|
|
)
|
|
@patch("ray.is_initialized", return_value=True)
|
|
@patch("ray._private.state.available_resources_per_node")
|
|
@patch("ray.nodes")
|
|
def test_get_num_ready_tpu_slices_calculation(
|
|
mock_nodes_call,
|
|
mock_avail_resources_call,
|
|
mock_is_initialized,
|
|
topology,
|
|
accelerator_type,
|
|
mock_nodes,
|
|
mock_avail_resources,
|
|
expected_ready,
|
|
):
|
|
"""Test that the TPU slice readiness utility correctly calculates the number of ready
|
|
slices in different mocked scenarios, including idle resource verification."""
|
|
mock_nodes_call.return_value = mock_nodes
|
|
mock_avail_resources_call.return_value = mock_avail_resources
|
|
|
|
actual_ready = ray.util.tpu.get_num_ready_tpu_slices(
|
|
topology=topology,
|
|
accelerator_type=accelerator_type,
|
|
)
|
|
assert actual_ready == expected_ready
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"topology, accelerator_type, mock_nodes, expected_intact",
|
|
[
|
|
# 1 fully intact v4 slice (2 physical hosts).
|
|
(
|
|
"2x2x2",
|
|
"v4",
|
|
[
|
|
_make_mock_tpu_node(True, "v4-16", "slice-1", 0, node_id="A"),
|
|
_make_mock_tpu_node(True, "v4-16", "slice-1", 1, node_id="B"),
|
|
],
|
|
1,
|
|
),
|
|
# Fractured slice (missing a physical host) -> 0 intact slices.
|
|
(
|
|
"2x2x2",
|
|
"v4",
|
|
[
|
|
_make_mock_tpu_node(True, "v4-16", "slice-1", 0, node_id="A"),
|
|
],
|
|
0,
|
|
),
|
|
# Missing head node (rank 0) -> 0 intact slices.
|
|
(
|
|
"2x2x2",
|
|
"v4",
|
|
[
|
|
_make_mock_tpu_node(True, "v4-16", "slice-1", 1, node_id="A"),
|
|
_make_mock_tpu_node(True, "v4-16", "slice-1", 2, node_id="B"),
|
|
],
|
|
0,
|
|
),
|
|
# One physical host is dead -> 0 intact slices.
|
|
(
|
|
"2x2x2",
|
|
"v4",
|
|
[
|
|
_make_mock_tpu_node(True, "v4-16", "slice-1", 0, node_id="A"),
|
|
_make_mock_tpu_node(False, "v4-16", "slice-1", 1, node_id="B"),
|
|
],
|
|
0,
|
|
),
|
|
# 2 slices: one intact, one fractured -> 1 intact slice.
|
|
(
|
|
"2x2x2",
|
|
"v4",
|
|
[
|
|
_make_mock_tpu_node(True, "v4-16", "slice-A", 0, node_id="A0"),
|
|
_make_mock_tpu_node(True, "v4-16", "slice-A", 1, node_id="A1"),
|
|
_make_mock_tpu_node(True, "v4-16", "slice-B", 0, node_id="B0"),
|
|
],
|
|
1,
|
|
),
|
|
# 2 fully intact v6e slices.
|
|
(
|
|
"4x4",
|
|
"v6e",
|
|
[
|
|
_make_mock_tpu_node(True, "v6e-16", "slice-1", 0, node_id="S1_0"),
|
|
_make_mock_tpu_node(True, "v6e-16", "slice-1", 1, node_id="S1_1"),
|
|
_make_mock_tpu_node(True, "v6e-16", "slice-1", 2, node_id="S1_2"),
|
|
_make_mock_tpu_node(True, "v6e-16", "slice-1", 3, node_id="S1_3"),
|
|
_make_mock_tpu_node(True, "v6e-16", "slice-2", 0, node_id="S2_0"),
|
|
_make_mock_tpu_node(True, "v6e-16", "slice-2", 1, node_id="S2_1"),
|
|
_make_mock_tpu_node(True, "v6e-16", "slice-2", 2, node_id="S2_2"),
|
|
_make_mock_tpu_node(True, "v6e-16", "slice-2", 3, node_id="S2_3"),
|
|
],
|
|
2,
|
|
),
|
|
],
|
|
)
|
|
@patch("ray.is_initialized", return_value=True)
|
|
@patch("ray.nodes")
|
|
def test_get_num_tpu_slices_calculation(
|
|
mock_nodes_call,
|
|
mock_is_initialized,
|
|
topology,
|
|
accelerator_type,
|
|
mock_nodes,
|
|
expected_intact,
|
|
):
|
|
"""Test that the intact TPU slice utility counts slices based purely on
|
|
physical integrity (all hosts alive, correct chip count) without checking
|
|
whether they are idle."""
|
|
mock_nodes_call.return_value = mock_nodes
|
|
|
|
actual_intact = ray.util.tpu.get_num_tpu_slices(
|
|
topology=topology,
|
|
accelerator_type=accelerator_type,
|
|
)
|
|
assert actual_intact == expected_intact
|
|
|
|
|
|
@patch("ray.is_initialized", return_value=False)
|
|
def test_get_num_tpu_slices_uninitialized(mock_is_initialized):
|
|
"""Test that the utility gracefully handles an uninitialized Ray context."""
|
|
assert ray.util.tpu.get_num_tpu_slices("2x2x2", "v4") == 0
|
|
|
|
|
|
def test_get_num_ready_tpu_slices(ray_tpu_cluster):
|
|
"""
|
|
Tests the get_num_ready_tpu_slices utility against a real Ray cluster.
|
|
The ray_tpu_cluster fixture provisions two v4-16 slices (2x2x2 topology).
|
|
"""
|
|
ready_slices = ray.util.tpu.get_num_ready_tpu_slices(
|
|
topology="2x2x2", accelerator_type="v4"
|
|
)
|
|
assert ready_slices == 2
|
|
|
|
|
|
@patch("ray.is_initialized", return_value=False)
|
|
def test_get_num_ready_tpu_slices_uninitialized(mock_is_initialized):
|
|
"""Test that the utility gracefully handles an uninitialized Ray context."""
|
|
assert ray.util.tpu.get_num_ready_tpu_slices("2x2x2", "v4") == 0
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"node_dict, expected_slice_name",
|
|
[
|
|
(_make_mock_tpu_node(True, "v4-16", "slice-1", 0), "slice-1"),
|
|
(_make_mock_tpu_node(True, "v6e-8", "slice-A", 1), "slice-A"),
|
|
(_make_mock_tpu_node(True, "v4-16", "", 0), ""),
|
|
({"Alive": True, "Labels": {}}, None), # Missing TPU slice name
|
|
({"Alive": True}, None), # Missing Node labels dict
|
|
],
|
|
)
|
|
def test_get_tpu_slice_name_from_node(node_dict, expected_slice_name):
|
|
"""Tests that the utility correctly extracts the TPU slice name from a node dictionary."""
|
|
assert ray.util.tpu.get_tpu_slice_name_from_node(node_dict) == expected_slice_name
|
|
|
|
|
|
@patch("ray.is_initialized", return_value=True)
|
|
@patch("ray.nodes")
|
|
def test_get_tpu_nodes_for_slice(mock_nodes_call, mock_is_initialized):
|
|
"""Tests that the utility correctly filters alive nodes for a specific slice."""
|
|
mock_nodes = [
|
|
_make_mock_tpu_node(True, "v4-16", "slice-A", 0),
|
|
_make_mock_tpu_node(True, "v4-16", "slice-A", 1),
|
|
_make_mock_tpu_node(False, "v4-16", "slice-A", 2), # Dead node
|
|
_make_mock_tpu_node(True, "v4-16", "slice-B", 0), # Wrong slice
|
|
]
|
|
mock_nodes_call.return_value = mock_nodes
|
|
|
|
# Call ray.nodes() to fetch from GCS
|
|
nodes_a = ray.util.tpu.get_tpu_nodes_for_slice("slice-A")
|
|
assert len(nodes_a) == 2
|
|
assert nodes_a[0]["Labels"][ray._raylet.RAY_NODE_TPU_WORKER_ID_KEY] == "0"
|
|
assert nodes_a[1]["Labels"][ray._raylet.RAY_NODE_TPU_WORKER_ID_KEY] == "1"
|
|
assert mock_nodes_call.call_count == 1
|
|
|
|
# Pass cached nodes directly
|
|
nodes_b = ray.util.tpu.get_tpu_nodes_for_slice("slice-B", nodes=mock_nodes)
|
|
assert len(nodes_b) == 1
|
|
assert nodes_b[0]["Labels"][ray._raylet.RAY_NODE_TPU_SLICE_NAME_KEY] == "slice-B"
|
|
assert mock_nodes_call.call_count == 1 # Call count remains 1
|
|
|
|
# Use non-existent slice
|
|
nodes_c = ray.util.tpu.get_tpu_nodes_for_slice("slice-C", nodes=mock_nodes)
|
|
assert len(nodes_c) == 0
|
|
|
|
|
|
@patch("ray.is_initialized", return_value=False)
|
|
def test_get_tpu_nodes_for_slice_uninitialized(mock_is_initialized):
|
|
"""Test that the utility gracefully handles an uninitialized Ray context."""
|
|
assert ray.util.tpu.get_tpu_nodes_for_slice("slice-A") == []
|
|
|
|
|
|
def test_get_tpu_worker_resources_chips_per_vm_override():
|
|
"""Test that chips_per_vm correctly overrides the default resource calculations."""
|
|
|
|
# Default behavior: v6e 2x4 defaults to a single 8-chip host
|
|
num_workers, resources = ray.util.tpu.get_tpu_worker_resources(
|
|
topology="2x4", accelerator_type="v6e"
|
|
)
|
|
assert num_workers == 1
|
|
assert resources["TPU"] == 8
|
|
|
|
# Override behavior: v6e 2x4 forced to 4 chips per VM (2 hosts)
|
|
num_workers_override, resources_override = ray.util.tpu.get_tpu_worker_resources(
|
|
topology="2x4", accelerator_type="v6e", chips_per_vm=4
|
|
)
|
|
assert num_workers_override == 2
|
|
assert resources_override["TPU"] == 4
|
|
|
|
|
|
def test_slice_placement_group_chips_per_vm_override(ray_v6e_tpu_cluster):
|
|
"""Test that SlicePlacementGroup respects chips_per_vm for host calculation."""
|
|
|
|
# Default behavior (1 VM with 8 chips)
|
|
default_pg = SlicePlacementGroup(topology="2x4", accelerator_version="v6e")
|
|
assert default_pg.chips_per_host == 8
|
|
assert default_pg.num_hosts == 1
|
|
assert default_pg.num_bundles == 1
|
|
assert default_pg.bundle_resources["TPU"] == 8
|
|
|
|
# User-specified override behavior (2 VMs with 4 chips each)
|
|
override_pg = SlicePlacementGroup(
|
|
topology="2x4", accelerator_version="v6e", chips_per_vm=4
|
|
)
|
|
assert override_pg.chips_per_host == 4
|
|
assert override_pg.num_hosts == 2
|
|
assert override_pg.num_bundles == 2
|
|
assert override_pg.bundle_resources["TPU"] == 4
|
|
|
|
|
|
def test_user_bundle_label_selector_merged(ray_tpu_cluster):
|
|
"""Verifies that user-passed bundle_label_selector is merged with dynamic TPU labels."""
|
|
user_selectors = [{"env": "prod"}, {"env": "test"}]
|
|
|
|
# 2x2x2 v4 = 2 hosts = 2 bundles
|
|
slice_pg = SlicePlacementGroup(
|
|
topology="2x2x2", accelerator_version="v4", bundle_label_selector=user_selectors
|
|
)
|
|
|
|
assert len(slice_pg._bundle_label_selector) == 2
|
|
|
|
# Verify slice 0
|
|
assert slice_pg._bundle_label_selector[0]["env"] == "prod"
|
|
assert ray._raylet.RAY_NODE_TPU_SLICE_NAME_KEY in slice_pg._bundle_label_selector[0]
|
|
|
|
# Verify slice 1
|
|
assert slice_pg._bundle_label_selector[1]["env"] == "test"
|
|
assert ray._raylet.RAY_NODE_TPU_SLICE_NAME_KEY in slice_pg._bundle_label_selector[1]
|
|
|
|
|
|
def test_user_bundle_label_selector_collision_dynamic_wins(ray_v6e_tpu_cluster):
|
|
"""Verifies that dynamic TPU labels take precedence on collision."""
|
|
user_selectors = [{ray._raylet.RAY_NODE_TPU_SLICE_NAME_KEY: "user-requested-slice"}]
|
|
|
|
# v6e-8 is single host (1 bundle)
|
|
slice_pg = SlicePlacementGroup(
|
|
topology="2x4", accelerator_version="v6e", bundle_label_selector=user_selectors
|
|
)
|
|
|
|
assert len(slice_pg._bundle_label_selector) == 1
|
|
# The dynamic value should win (it generates test-v6e-slice-N)
|
|
actual_val = slice_pg._bundle_label_selector[0][
|
|
ray._raylet.RAY_NODE_TPU_SLICE_NAME_KEY
|
|
]
|
|
assert actual_val != "user-requested-slice"
|
|
assert "test-v6e-slice-" in actual_val
|
|
|
|
|
|
def test_user_bundle_label_selector_length_mismatch_raises():
|
|
"""Verifies that providing wrong length of selector list raises ValueError."""
|
|
user_selectors = [{"env": "prod"}] # Only 1 provided but 2x2x2 v4 has 2 hosts
|
|
|
|
with pytest.raises(ValueError, match="bundle_label_selector length"):
|
|
SlicePlacementGroup(
|
|
topology="2x2x2",
|
|
accelerator_version="v4",
|
|
bundle_label_selector=user_selectors,
|
|
)
|
|
|
|
|
|
def test_release_head_pgs_idempotent(ray_tpu_cluster):
|
|
"""Verifies that release_head_pgs() is idempotent."""
|
|
slice_pg = SlicePlacementGroup(topology="2x2x2", accelerator_version="v4")
|
|
|
|
assert len(slice_pg.head_placement_groups) == 1
|
|
|
|
slice_pg.release_head_pgs()
|
|
assert len(slice_pg.head_placement_groups) == 0
|
|
|
|
# Call again, should not raise
|
|
slice_pg.release_head_pgs()
|
|
assert len(slice_pg.head_placement_groups) == 0
|
|
|
|
|
|
def test_shutdown_idempotent(ray_tpu_cluster):
|
|
"""Verifies that shutdown() is idempotent."""
|
|
slice_pg = SlicePlacementGroup(topology="2x2x2", accelerator_version="v4")
|
|
|
|
slice_pg.shutdown()
|
|
assert slice_pg.placement_group is None
|
|
assert len(slice_pg.head_placement_groups) == 0
|
|
|
|
# Call again, should not raise
|
|
slice_pg.shutdown()
|
|
|
|
|
|
def test_shutdown_safe_after_construction_failure():
|
|
"""Verifies that shutdown() is safe to call on a partially-constructed instance."""
|
|
with patch(
|
|
"ray.util.tpu.SlicePlacementGroup._reserve_slice",
|
|
side_effect=RuntimeError("Test failure"),
|
|
):
|
|
with pytest.raises(RuntimeError, match="Test failure"):
|
|
SlicePlacementGroup(topology="2x2x2", accelerator_version="v4")
|
|
|
|
# If the above didn't crash or leak resources, we are good.
|
|
# We can also manually construct a partial instance and call shutdown.
|
|
partial_pg = SlicePlacementGroup.__new__(SlicePlacementGroup)
|
|
partial_pg._head_pgs = []
|
|
partial_pg._placement_group = None
|
|
|
|
# Should not raise even though it's missing attributes
|
|
partial_pg.shutdown()
|
|
|
|
|
|
def test_release_head_pgs_after_ready_then_shutdown(ray_tpu_cluster):
|
|
"""Validates Slice PG lifecycle: wait until ready, release head PGs, then shutdown."""
|
|
slice_pg = SlicePlacementGroup(topology="2x2x2", accelerator_version="v4")
|
|
|
|
# Wait for ready
|
|
ray.get(slice_pg.placement_group.ready())
|
|
|
|
slice_pg.release_head_pgs()
|
|
assert len(slice_pg.head_placement_groups) == 0
|
|
|
|
slice_pg.shutdown()
|
|
assert slice_pg.placement_group is None
|
|
|
|
|
|
def test_chips_per_vm_zero_raises_value_error():
|
|
"""Verifies that passing chips_per_vm=0 explicitly raises a ValueError instead of silently using the topology default."""
|
|
with pytest.raises(ValueError):
|
|
SlicePlacementGroup(
|
|
topology="2x2x2",
|
|
accelerator_version="v4",
|
|
chips_per_vm=0,
|
|
)
|
|
|
|
# Also verify when custom resources already include a "TPU" key
|
|
with pytest.raises(ValueError):
|
|
SlicePlacementGroup(
|
|
topology="2x2x2",
|
|
accelerator_version="v4",
|
|
resources_per_bundle={"TPU": 4},
|
|
chips_per_vm=0,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-sv", __file__]))
|