chore: import upstream snapshot with attribution

python/ray/train/tests/test_worker_group.py

@@ -0,0 +1,352 @@
+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__]))