| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Make sure ray assigns GPU workers to the correct node. |
| 4 | |
| 5 | Run: |
| 6 | ```sh |
| 7 | cd $VLLM_PATH/tests |
| 8 | |
| 9 | pytest distributed/test_multi_node_assignment.py |
| 10 | ``` |
| 11 | """ |
| 12 | |
| 13 | import os |
| 14 | import time |
| 15 | |
| 16 | import pytest |
| 17 | import ray |
| 18 | from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy |
| 19 | |
| 20 | from vllm import initialize_ray_cluster |
| 21 | from vllm.config import ParallelConfig |
| 22 | from vllm.utils.network_utils import get_ip |
| 23 | from vllm.v1.executor.ray_utils import available_resources_per_node |
| 24 | |
| 25 | VLLM_MULTI_NODE = os.getenv("VLLM_MULTI_NODE", "0") == "1" |
| 26 | |
| 27 | |
| 28 | @pytest.mark.skipif( |
| 29 | not VLLM_MULTI_NODE, reason="Need at least 2 nodes to run the test." |
| 30 | ) |
| 31 | def test_multi_node_assignment() -> None: |
| 32 | # NOTE: important to keep this class definition here |
| 33 | # to let ray use cloudpickle to serialize it. |
| 34 | class Actor: |
| 35 | def get_ip(self): |
| 36 | return get_ip() |
| 37 | |
| 38 | for _ in range(10): |
| 39 | config = ParallelConfig(1, 2) |
| 40 | initialize_ray_cluster(config) |
| 41 | |
| 42 | current_ip = get_ip() |
| 43 | workers = [] |
| 44 | for bundle_id, bundle in enumerate(config.placement_group.bundle_specs): |
| 45 | if not bundle.get("GPU", 0): |
| 46 | continue |
| 47 | scheduling_strategy = PlacementGroupSchedulingStrategy( |
| 48 | placement_group=config.placement_group, |
| 49 | placement_group_capture_child_tasks=True, |
| 50 | placement_group_bundle_index=bundle_id, |
| 51 | ) |
| 52 | |
| 53 | worker = ray.remote( |
| 54 | num_cpus=0, |
| 55 | num_gpus=1, |
| 56 | scheduling_strategy=scheduling_strategy, |
| 57 | )(Actor).remote() |
| 58 | worker_ip = ray.get(worker.get_ip.remote()) |
| 59 | assert worker_ip == current_ip |
| 60 | workers.append(worker) |
| 61 | |
| 62 | for worker in workers: |
| 63 | ray.kill(worker) |
| 64 | |
| 65 | ray.util.remove_placement_group(config.placement_group) |
| 66 | |
| 67 | gpus_needed = len(workers) |
| 68 | deadline = time.monotonic() + 60 |
| 69 | node_id = ray.get_runtime_context().get_node_id() |
| 70 | while True: |
| 71 | free_gpus = available_resources_per_node().get(node_id, {}).get("GPU", 0) |
| 72 | if free_gpus >= gpus_needed: |
| 73 | break |
| 74 | |
| 75 | remaining = deadline - time.monotonic() |
| 76 | if remaining <= 0: |
| 77 | raise TimeoutError("placement group resources were not released") |
| 78 | |
| 79 | time.sleep(min(1, remaining)) |