wehub-resource-sync
159 lines
4.7 KiB
Python
159 lines
4.7 KiB
Python
"""Tests for get_bundle_indices_sorted_by_node."""
|
|
|
|
import sys
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
import ray
|
|
from ray._private.test_utils import placement_group_assert_no_leak
|
|
from ray.llm._internal.serve.utils.pg_utils import get_bundle_indices_sorted_by_node
|
|
from ray.serve._private.utils import get_head_node_id
|
|
|
|
# Realistic 56-char hex node IDs
|
|
NODE_A = "ab7c1d2e3f4a5b6c7d8e9f0a1b2c3d4e5f6a7b8c9d0e1f2a3b4c" # driver node
|
|
NODE_B = "1234567890abcdef1234567890abcdef1234567890abcdef12345678"
|
|
NODE_C = "fedcba0987654321fedcba0987654321fedcba0987654321fedcba09"
|
|
|
|
|
|
def _assert_grouped_by_node(result, bundles_to_node_id):
|
|
seen_nodes = set()
|
|
prev_node = None
|
|
for idx in result:
|
|
node = bundles_to_node_id[idx]
|
|
if node != prev_node:
|
|
assert (
|
|
node not in seen_nodes
|
|
), f"Node {node} appeared non-contiguously in {result}"
|
|
seen_nodes.add(node)
|
|
prev_node = node
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"bundles_to_node_id,expected_sorted_bundle_indices",
|
|
[
|
|
pytest.param(
|
|
{0: NODE_C, 1: NODE_A, 2: NODE_B, 3: NODE_C, 4: NODE_A, 5: NODE_B},
|
|
[1, 4, 2, 5, 0, 3],
|
|
),
|
|
pytest.param(
|
|
{0: NODE_B, 1: NODE_B, 2: NODE_A, 3: NODE_A},
|
|
[2, 3, 0, 1],
|
|
),
|
|
pytest.param(
|
|
# Driver has no bundles
|
|
{0: NODE_C, 1: NODE_B, 2: NODE_C, 3: NODE_B},
|
|
[1, 3, 0, 2],
|
|
),
|
|
pytest.param(
|
|
{0: NODE_A, 1: NODE_A, 2: NODE_A},
|
|
[0, 1, 2],
|
|
),
|
|
pytest.param(
|
|
{0: NODE_A},
|
|
[0],
|
|
),
|
|
pytest.param(
|
|
{},
|
|
[],
|
|
),
|
|
],
|
|
)
|
|
def test_sort_bundle_indices_by_node(
|
|
bundles_to_node_id, expected_sorted_bundle_indices
|
|
):
|
|
mock_pg = MagicMock()
|
|
|
|
with (
|
|
patch(
|
|
"ray.llm._internal.serve.utils.pg_utils.placement_group_table",
|
|
return_value={"bundles_to_node_id": bundles_to_node_id},
|
|
),
|
|
patch(
|
|
"ray.llm._internal.serve.utils.pg_utils.get_head_node_id",
|
|
return_value=NODE_A,
|
|
),
|
|
):
|
|
result = get_bundle_indices_sorted_by_node(mock_pg)
|
|
|
|
assert result == expected_sorted_bundle_indices
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"bundles_to_node_id,driver_node_id,expected_sorted_bundle_indices",
|
|
[
|
|
# Explicit driver node orders that node's bundles first, even though the
|
|
# head node (NODE_A) also owns bundles.
|
|
pytest.param(
|
|
{0: NODE_A, 1: NODE_B, 2: NODE_A, 3: NODE_B},
|
|
NODE_B,
|
|
[1, 3, 0, 2],
|
|
),
|
|
# Head node (NODE_A) owns no bundles; explicit driver still wins over
|
|
# lexicographic fallback (which would otherwise pick NODE_B first).
|
|
pytest.param(
|
|
{0: NODE_B, 1: NODE_C, 2: NODE_B, 3: NODE_C},
|
|
NODE_C,
|
|
[1, 3, 0, 2],
|
|
),
|
|
],
|
|
)
|
|
def test_explicit_driver_node_ordered_first(
|
|
bundles_to_node_id, driver_node_id, expected_sorted_bundle_indices
|
|
):
|
|
mock_pg = MagicMock()
|
|
|
|
with (
|
|
patch(
|
|
"ray.llm._internal.serve.utils.pg_utils.placement_group_table",
|
|
return_value={"bundles_to_node_id": bundles_to_node_id},
|
|
),
|
|
patch(
|
|
"ray.llm._internal.serve.utils.pg_utils.get_head_node_id",
|
|
return_value=NODE_A,
|
|
),
|
|
):
|
|
result = get_bundle_indices_sorted_by_node(
|
|
mock_pg, driver_node_id=driver_node_id
|
|
)
|
|
|
|
assert result == expected_sorted_bundle_indices
|
|
|
|
|
|
# Integration test: verifies the full path through a real cluster, though
|
|
# bundle ordering across nodes is non-deterministic and may already be sorted.
|
|
@pytest.mark.parametrize("strategy", ["SPREAD", "PACK"])
|
|
def test_sort_bundles(ray_start_cluster, strategy):
|
|
cluster = ray_start_cluster
|
|
for _ in range(2):
|
|
cluster.add_node(num_cpus=2)
|
|
ray.init(address=cluster.address)
|
|
cluster.wait_for_nodes()
|
|
|
|
driver_node_id = get_head_node_id()
|
|
|
|
pg = ray.util.placement_group(
|
|
name=f"test_sort_bundles_{strategy}",
|
|
strategy=strategy,
|
|
bundles=[{"CPU": 1}] * 4,
|
|
)
|
|
ray.get(pg.ready())
|
|
|
|
result = get_bundle_indices_sorted_by_node(pg)
|
|
|
|
table = ray.util.placement_group_table(pg)
|
|
bundles_to_node_id = table["bundles_to_node_id"]
|
|
|
|
assert sorted(result) == list(range(4))
|
|
_assert_grouped_by_node(result, bundles_to_node_id)
|
|
|
|
driver_bundles = [i for i in result if bundles_to_node_id[i] == driver_node_id]
|
|
if driver_bundles:
|
|
assert result[: len(driver_bundles)] == driver_bundles
|
|
|
|
placement_group_assert_no_leak([pg])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-v", __file__]))
|