wehub-resource-sync
353 lines
11 KiB
Python
353 lines
11 KiB
Python
import time
|
|
from collections import defaultdict
|
|
from typing import Dict, List, Optional
|
|
|
|
import pytest
|
|
|
|
import ray
|
|
import ray._private.ray_constants as ray_constants
|
|
from ray.cluster_utils import Cluster
|
|
from ray.train._internal.worker_group import Worker, WorkerGroup, WorkerMetadata
|
|
from ray.util.state import list_actors
|
|
|
|
|
|
@pytest.fixture
|
|
def ray_start_2_cpus():
|
|
address_info = ray.init(num_cpus=2)
|
|
yield address_info
|
|
# The code after the yield will run as teardown code.
|
|
ray.shutdown()
|
|
|
|
|
|
@pytest.fixture
|
|
def ray_start_2_cpus_and_gpus():
|
|
address_info = ray.init(num_cpus=2, num_gpus=2)
|
|
yield address_info
|
|
# The code after the yield will run as teardown code.
|
|
ray.shutdown()
|
|
|
|
|
|
@pytest.fixture
|
|
def ray_start_2_cpus_and_neuron_core_accelerator():
|
|
address_info = ray.init(num_cpus=2, resources={ray_constants.NEURON_CORES: 2})
|
|
yield address_info
|
|
# The code after the yield will run as teardown code.
|
|
ray.shutdown()
|
|
|
|
|
|
@pytest.fixture
|
|
def ray_start_2_cpus_and_10kb_memory():
|
|
address_info = ray.init(num_cpus=2, _memory=10_000)
|
|
yield address_info
|
|
# The code after the yield will run as teardown code.
|
|
ray.shutdown()
|
|
|
|
|
|
@pytest.fixture
|
|
def ray_start_5_nodes_with_memory():
|
|
cluster = Cluster()
|
|
for _ in range(4):
|
|
cluster.add_node(num_cpus=4, memory=500)
|
|
cluster.add_node(num_cpus=4, memory=2_000)
|
|
|
|
ray.init(address=cluster.address)
|
|
|
|
yield
|
|
|
|
ray.shutdown()
|
|
cluster.shutdown()
|
|
|
|
|
|
def test_worker_creation(ray_start_2_cpus):
|
|
assert ray.available_resources()["CPU"] == 2
|
|
wg = WorkerGroup(num_workers=2)
|
|
assert len(wg.workers) == 2
|
|
time.sleep(1)
|
|
# Make sure both CPUs are being used by the actors.
|
|
assert "CPU" not in ray.available_resources()
|
|
wg.shutdown()
|
|
|
|
|
|
def test_worker_creation_num_cpus(ray_start_2_cpus):
|
|
assert ray.available_resources()["CPU"] == 2
|
|
wg = WorkerGroup(resources_per_worker={"CPU": 2})
|
|
time.sleep(1)
|
|
assert len(wg.workers) == 1
|
|
# Make sure both CPUs are being used by the actor.
|
|
assert "CPU" not in ray.available_resources()
|
|
wg.shutdown()
|
|
|
|
|
|
def test_worker_creation_with_memory(ray_start_5_nodes_with_memory):
|
|
resources_per_worker = {"memory": 1_000}
|
|
wg = WorkerGroup(num_workers=2, resources_per_worker=resources_per_worker)
|
|
assert len(wg.workers) == 2
|
|
|
|
nodes = ray.nodes()
|
|
large_node = [node for node in nodes if node["Resources"]["memory"] == 2_000][0]
|
|
large_node_id = large_node["NodeID"]
|
|
|
|
def validate_scheduling():
|
|
resources = ray.get_runtime_context().get_assigned_resources()
|
|
assert resources == resources_per_worker, "Resources should include memory."
|
|
|
|
node_id = ray.get_runtime_context().get_node_id()
|
|
assert (
|
|
node_id == large_node_id
|
|
), "Workers should be scheduled on the large node."
|
|
|
|
wg.execute(validate_scheduling)
|
|
|
|
|
|
def test_worker_shutdown(ray_start_2_cpus):
|
|
assert ray.available_resources()["CPU"] == 2
|
|
wg = WorkerGroup(num_workers=2)
|
|
time.sleep(1)
|
|
assert "CPU" not in ray.available_resources()
|
|
assert len(list_actors()) == 2
|
|
wg.shutdown()
|
|
time.sleep(1)
|
|
assert ray.available_resources()["CPU"] == 2
|
|
|
|
with pytest.raises(RuntimeError):
|
|
wg.execute(lambda: 1)
|
|
|
|
|
|
def test_worker_restart(ray_start_2_cpus):
|
|
wg = WorkerGroup(num_workers=2)
|
|
with pytest.raises(RuntimeError):
|
|
wg.start()
|
|
# Avoid race condition.
|
|
time.sleep(1)
|
|
wg.shutdown(0)
|
|
wg.start()
|
|
wg.execute(lambda: 1)
|
|
|
|
|
|
def test_worker_with_gpu_ids(ray_start_2_cpus_and_gpus):
|
|
num_gpus = 2
|
|
wg = WorkerGroup(num_workers=2, resources_per_worker={"GPU": 1})
|
|
assert len(wg.workers) == 2
|
|
time.sleep(1)
|
|
assert ray_constants.GPU not in ray.available_resources()
|
|
wg.execute(lambda: 1)
|
|
assert len(wg.workers) == 2
|
|
for w in wg.workers:
|
|
resource_ids = w.metadata.resource_ids
|
|
gpu_ids = resource_ids[ray_constants.GPU]
|
|
for gpu_id in gpu_ids:
|
|
assert gpu_id in [str(i) for i in range(num_gpus)]
|
|
assert len(resource_ids[ray_constants.NEURON_CORES]) == 0
|
|
|
|
|
|
def test_worker_with_neuron_core_accelerator_ids(
|
|
ray_start_2_cpus_and_neuron_core_accelerator,
|
|
):
|
|
num_nc = 2
|
|
wg = WorkerGroup(
|
|
num_workers=2, resources_per_worker={ray_constants.NEURON_CORES: 1}
|
|
)
|
|
assert len(wg.workers) == 2
|
|
time.sleep(1)
|
|
assert ray_constants.NEURON_CORES not in ray.available_resources()
|
|
wg.execute(lambda: 1)
|
|
assert len(wg.workers) == 2
|
|
for w in wg.workers:
|
|
resource_ids = w.metadata.resource_ids
|
|
assert len(resource_ids[ray_constants.GPU]) == 0
|
|
neuron_core_ids = resource_ids[ray_constants.NEURON_CORES]
|
|
for neuron_core_id in neuron_core_ids:
|
|
assert neuron_core_id in [str(i) for i in range(num_nc)]
|
|
|
|
|
|
def test_execute_async(ray_start_2_cpus):
|
|
wg = WorkerGroup(num_workers=2)
|
|
futures = wg.execute_async(lambda: 1)
|
|
assert len(futures) == 2
|
|
outputs = ray.get(futures)
|
|
assert all(o == 1 for o in outputs)
|
|
|
|
|
|
def test_execute(ray_start_2_cpus):
|
|
wg = WorkerGroup(num_workers=2)
|
|
outputs = wg.execute(lambda: 1)
|
|
assert len(outputs) == 2
|
|
assert all(o == 1 for o in outputs)
|
|
|
|
|
|
def test_execute_args(ray_start_2_cpus):
|
|
wg = WorkerGroup(num_workers=2)
|
|
outputs = wg.execute(lambda x: x, 1)
|
|
assert len(outputs) == 2
|
|
assert all(o == 1 for o in outputs)
|
|
|
|
|
|
def test_group_workers_by_node_id(ray_start_2_cpus):
|
|
def create_worker_group(node_ids):
|
|
wg = WorkerGroup(num_workers=2)
|
|
wg.workers = [
|
|
Worker(
|
|
actor=None,
|
|
metadata=WorkerMetadata(
|
|
node_id=node_id,
|
|
node_ip="dummy",
|
|
hostname="dummy",
|
|
resource_ids={},
|
|
pid=0,
|
|
),
|
|
)
|
|
for node_id in node_ids
|
|
]
|
|
return wg
|
|
|
|
wg = create_worker_group(["2", "3", "1", "4", "2", "1", "3", "3", "4", "2"])
|
|
wg.sort_workers_by_node_id_and_gpu_id()
|
|
expected = ["2", "2", "2", "3", "3", "3", "1", "1", "4", "4"]
|
|
node_ids = [w.metadata.node_id for w in wg.workers]
|
|
assert node_ids == expected, (
|
|
"Workers should be grouped by Node ID "
|
|
"and follow the same original order of IDs encountered (2, 3, 1, 4)."
|
|
)
|
|
|
|
wg = create_worker_group(["2", "3", "1", "4", "2", "1", "3", "3", "4", "2"])
|
|
wg.sort_workers_by_node_id_and_gpu_id(_first_node_id="1")
|
|
expected = ["1", "1", "2", "2", "2", "3", "3", "3", "4", "4"]
|
|
node_ids = [w.metadata.node_id for w in wg.workers]
|
|
assert (
|
|
node_ids == expected
|
|
), "Workers should be grouped by Node ID, with the first ID being 1."
|
|
|
|
|
|
def test_sort_local_workers_by_gpu_id(ray_start_2_cpus):
|
|
def create_worker_group(pids, node_ids, gpu_ids):
|
|
wg = WorkerGroup(num_workers=2)
|
|
wg.workers = [
|
|
Worker(
|
|
actor=None,
|
|
metadata=WorkerMetadata(
|
|
node_id=node_id,
|
|
node_ip="dummy",
|
|
hostname="dummy",
|
|
resource_ids={"GPU": gpu_id.split() if gpu_id else []},
|
|
pid=pid,
|
|
),
|
|
)
|
|
for pid, node_id, gpu_id in zip(pids, node_ids, gpu_ids)
|
|
]
|
|
return wg
|
|
|
|
def setup_and_check_worker_group(
|
|
pids: List[int],
|
|
node_ids: List[str],
|
|
gpu_ids: List[Optional[str]],
|
|
expected_local_ranks: Dict[int, int],
|
|
):
|
|
"""
|
|
Create a worker group, group workers by Node ID,
|
|
and check local ranks assignment.
|
|
|
|
Args:
|
|
pids: List of unique process IDs.
|
|
node_ids: List of Node IDs corresponding to each PID.
|
|
gpu_ids: List of GPU IDs or None for each PID.
|
|
expected_local_ranks: Dictionary mapping PID to the
|
|
expected local rank.
|
|
"""
|
|
wg = create_worker_group(pids=pids, node_ids=node_ids, gpu_ids=gpu_ids)
|
|
wg.sort_workers_by_node_id_and_gpu_id()
|
|
|
|
# Build local ranks according to the logics in
|
|
# `BackendExecutor._create_rank_world_size_mappings()`
|
|
node_id_dict = defaultdict(int)
|
|
local_ranks_map = defaultdict(int)
|
|
for w in wg.workers:
|
|
local_ranks_map[w.metadata.pid] = node_id_dict[w.metadata.node_id]
|
|
node_id_dict[w.metadata.node_id] += 1
|
|
|
|
local_ranks = [local_ranks_map[pid] for pid in pids]
|
|
|
|
assert (
|
|
local_ranks == expected_local_ranks
|
|
), "Incorrect local ranks allocation!\n"
|
|
f"Expect: {expected_local_ranks}\nGot: {local_ranks}"
|
|
|
|
# Define the worker configurations for different scenarios
|
|
# For workers without GPU resources, their original order will be preserved
|
|
cpu_workers_config = {
|
|
"pids": [0, 1, 2, 3, 4, 5, 6, 7],
|
|
"node_ids": ["2", "2", "1", "1", "2", "1", "1", "2"],
|
|
"gpu_ids": [None] * 8,
|
|
"expected_local_ranks": [0, 1, 0, 1, 2, 2, 3, 3],
|
|
}
|
|
|
|
gpu_workers_single_gpu_config = {
|
|
"pids": [0, 1, 2, 3, 4, 5, 6, 7],
|
|
"node_ids": ["2", "2", "1", "1", "2", "1", "1", "2"],
|
|
"gpu_ids": ["1", "0", "3", "2", "2", "0", "1", "3"],
|
|
"expected_local_ranks": [1, 0, 3, 2, 2, 0, 1, 3],
|
|
}
|
|
|
|
# For workers with multiple gpus, sort by their lowest gpu id
|
|
gpu_workers_multiple_gpus_config = {
|
|
"pids": [0, 1, 2, 3],
|
|
"node_ids": ["2", "1", "1", "2"],
|
|
"gpu_ids": ["1,3", "2,1", "0,3", "0,2"],
|
|
"expected_local_ranks": [1, 1, 0, 0],
|
|
}
|
|
|
|
# Setup and check worker groups for each configuration
|
|
setup_and_check_worker_group(**cpu_workers_config)
|
|
setup_and_check_worker_group(**gpu_workers_single_gpu_config)
|
|
setup_and_check_worker_group(**gpu_workers_multiple_gpus_config)
|
|
|
|
|
|
def test_execute_single(ray_start_2_cpus):
|
|
wg = WorkerGroup(num_workers=2)
|
|
|
|
def f():
|
|
import os
|
|
|
|
os.environ["TEST"] = "1"
|
|
|
|
wg.execute_single(1, f)
|
|
|
|
def check():
|
|
import os
|
|
|
|
return os.environ.get("TEST", "0")
|
|
|
|
assert wg.execute(check) == ["0", "1"]
|
|
|
|
|
|
def test_bad_resources(ray_start_2_cpus):
|
|
with pytest.raises(ValueError):
|
|
WorkerGroup(num_workers=-1)
|
|
|
|
with pytest.raises(ValueError):
|
|
WorkerGroup(resources_per_worker={"CPU": -1})
|
|
|
|
with pytest.raises(ValueError):
|
|
WorkerGroup(resources_per_worker={"GPU": -1})
|
|
|
|
with pytest.raises(ValueError):
|
|
WorkerGroup(resources_per_worker={"memory": -1})
|
|
|
|
|
|
def test_placement_group(ray_start_2_cpus):
|
|
"""Tests that workers can be removed and added to a placement group."""
|
|
num_workers = 2
|
|
bundle = {"CPU": 1}
|
|
bundles = [bundle.copy() for _ in range(num_workers)]
|
|
placement_group = ray.util.placement_group(bundles)
|
|
wg = WorkerGroup(num_workers=num_workers, placement_group=placement_group)
|
|
wg.remove_workers([0])
|
|
wg.add_workers(1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
import pytest
|
|
|
|
sys.exit(pytest.main(["-v", "-x", __file__]))
|