ray_utils.py 4.52 KB
Newer Older
1
from typing import Optional, List, Tuple, TYPE_CHECKING
2
3

from vllm.config import ParallelConfig
4
from vllm.logger import init_logger
5
from vllm.utils import is_hip, set_cuda_visible_devices, get_ip
6
7

logger = init_logger(__name__)
8
9
10

try:
    import ray
11

12
    class RayWorkerVllm:
13
14
15
        """Ray wrapper for vllm.worker.Worker, allowing Worker to be
        lazliy initialized after Ray sets CUDA_VISIBLE_DEVICES."""

16
17
18
19
        def __init__(self, init_cached_hf_modules=False) -> None:
            if init_cached_hf_modules:
                from transformers.dynamic_module_utils import init_hf_modules
                init_hf_modules()
20
21
22
23
24
25
26
27
28
29
30
31
            self.worker = None

        def init_worker(self, worker_init_fn):
            self.worker = worker_init_fn()

        def __getattr__(self, name):
            return getattr(self.worker, name)

        def execute_method(self, method, *args, **kwargs):
            executor = getattr(self, method)
            return executor(*args, **kwargs)

32
33
34
35
36
37
38
39
40
41
42
        def get_node_ip(self) -> str:
            return get_ip()

        def get_node_and_gpu_ids(self) -> Tuple[str, List[int]]:
            node_id = ray.get_runtime_context().get_node_id()
            gpu_ids = ray.get_gpu_ids()
            return node_id, gpu_ids

        def set_cuda_visible_devices(self, device_ids) -> None:
            set_cuda_visible_devices(device_ids)

43
44
45
46
except ImportError as e:
    logger.warning(f"Failed to import Ray with {e!r}. "
                   "For distributed inference, please install Ray with "
                   "`pip install ray pandas pyarrow`.")
47
    ray = None
48
    RayWorkerVllm = None
49

50
51
if TYPE_CHECKING:
    from ray.util.placement_group import PlacementGroup
52
53
54
55


def initialize_cluster(
    parallel_config: ParallelConfig,
Zhuohan Li's avatar
Zhuohan Li committed
56
57
    engine_use_ray: bool = False,
    ray_address: Optional[str] = None,
ljss's avatar
ljss committed
58
) -> Optional["PlacementGroup"]:
59
60
61
62
    """Initialize the distributed cluster probably with Ray.

    Args:
        parallel_config: The configurations for parallel execution.
Zhuohan Li's avatar
Zhuohan Li committed
63
64
        engine_use_ray: Whether to use Ray for async engine.
        ray_address: The address of the Ray cluster. If None, uses
65
66
67
            the default Ray cluster address.

    Returns:
68
69
70
        An optional `PlacementGroup`. It includes the specification
        of the resources for each distributed worker. None if Ray is
        not used.
71
    """
Zhuohan Li's avatar
Zhuohan Li committed
72
    if parallel_config.worker_use_ray or engine_use_ray:
73
74
75
76
77
        if ray is None:
            raise ImportError(
                "Ray is not installed. Please install Ray to use distributed "
                "serving.")
        # Connect to a ray cluster.
78
79
80
81
82
83
        if is_hip():
            ray.init(address=ray_address,
                     ignore_reinit_error=True,
                     num_gpus=parallel_config.world_size)
        else:
            ray.init(address=ray_address, ignore_reinit_error=True)
84
85

    if not parallel_config.worker_use_ray:
86
87
88
        assert parallel_config.world_size == 1, (
            "Ray is required if parallel_config.world_size > 1.")
        return None
89

90
    # Create placement group for worker processes
91
92
93
94
95
96
97
    current_placement_group = ray.util.get_current_placement_group()
    if current_placement_group:
        # We are in a placement group
        bundles = current_placement_group.bundle_specs
        # Verify that we can use the placement group.
        gpu_bundles = 0
        for bundle in bundles:
98
99
100
101
102
            bundle_gpus = bundle.get("GPU", 0)
            if bundle_gpus > 1:
                raise ValueError(
                    "Placement group bundle cannot have more than 1 GPU.")
            if bundle_gpus:
103
104
                gpu_bundles += 1
        if parallel_config.world_size > gpu_bundles:
105
            raise ValueError(
106
107
                "The number of required GPUs exceeds the total number of "
                "available GPUs in the placement group.")
108
    else:
109
110
        num_gpus_in_cluster = ray.cluster_resources().get("GPU", 0)
        if parallel_config.world_size > num_gpus_in_cluster:
111
            raise ValueError(
112
113
114
                "The number of required GPUs exceeds the total number of "
                "available GPUs in the cluster.")
        # Create a new placement group
115
116
117
        placement_group_specs = ([{"GPU": 1}] * parallel_config.world_size)
        current_placement_group = ray.util.placement_group(
            placement_group_specs)
118
119
120
121
122
        # Wait until PG is ready - this will block until all
        # requested resources are available, and will timeout
        # if they cannot be provisioned.
        ray.get(current_placement_group.ready(), timeout=1800)

123
    return current_placement_group