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

from vllm.config import ParallelConfig
4
from vllm.logger import init_logger
5
from vllm.sequence import ExecuteModelRequest, IntermediateTensors
6
from vllm.utils import get_ip, is_hip, is_tpu, is_xpu
7
from vllm.worker.worker_base import WorkerWrapperBase
8
9

logger = init_logger(__name__)
10
11
12

try:
    import ray
13

14
    class RayWorkerWrapper(WorkerWrapperBase):
15
16
17
        """Ray wrapper for vllm.worker.Worker, allowing Worker to be
        lazliy initialized after Ray sets CUDA_VISIBLE_DEVICES."""

18
19
        def __init__(self, *args, **kwargs) -> None:
            super().__init__(*args, **kwargs)
20
21
22
23
24
            # Since the compiled DAG runs a main execution
            # in a different thread that calls cuda.set_device.
            # The flag indicates is set_device is called on
            # that thread.
            self.compiled_dag_cuda_device_set = False
25

26
27
28
29
30
31
32
33
        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

34
35
36
37
38
39
40
41
42
43
44
        def execute_model_spmd(
            self, req_or_tuple: Union[ExecuteModelRequest,
                                      Tuple[ExecuteModelRequest,
                                            IntermediateTensors]]):
            """Execute model in SPMD fashion: used only when SPMD worker and
            compiled DAG are both enabled.

            Args:
                req_or_tuple: The request to execute the model, or a tuple
                    containing the request and intermediate tensors.
            """
45
46
47
            # TODO(swang): This is needed right now because Ray aDAG executes
            # on a background thread, so we need to reset torch's current
            # device.
48
49
50
51
52
            import torch
            if not self.compiled_dag_cuda_device_set:
                torch.cuda.set_device(self.worker.device)
                self.compiled_dag_cuda_device_set = True

53
54
55
56
57
58
59
60
61
62
63
            if isinstance(req_or_tuple, tuple):
                execute_model_req, intermediate_tensors = req_or_tuple
            else:
                execute_model_req = req_or_tuple
                intermediate_tensors = None

            output = self.worker._execute_model_spmd(execute_model_req,
                                                     intermediate_tensors)
            if isinstance(output, IntermediateTensors):
                return execute_model_req, output
            return output
64

65
66
    ray_import_err = None

67
except ImportError as e:
68
    ray = None  # type: ignore
69
    ray_import_err = e
70
    RayWorkerWrapper = None  # type: ignore
71
72


73
74
75
76
77
78
79
80
81
82
83
84
def ray_is_available() -> bool:
    """Returns True if Ray is available."""
    return ray is not None


def assert_ray_available():
    """Raise an exception if Ray is not available."""
    if ray is None:
        raise ValueError("Failed to import Ray, please install Ray with "
                         "`pip install ray`.") from ray_import_err


85
def initialize_ray_cluster(
86
    parallel_config: ParallelConfig,
Zhuohan Li's avatar
Zhuohan Li committed
87
    ray_address: Optional[str] = None,
88
89
90
91
92
93
):
    """Initialize the distributed cluster with Ray.

    it will connect to the Ray cluster and create a placement group
    for the workers, which includes the specification of the resources
    for each distributed worker.
94
95
96

    Args:
        parallel_config: The configurations for parallel execution.
Zhuohan Li's avatar
Zhuohan Li committed
97
        ray_address: The address of the Ray cluster. If None, uses
98
99
            the default Ray cluster address.
    """
100
    assert_ray_available()
101
102

    # Connect to a ray cluster.
103
    if is_hip() or is_xpu():
104
105
106
107
108
109
110
111
112
        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)

    if parallel_config.placement_group:
        # Placement group is already set.
        return
113

114
    device_str = "GPU" if not is_tpu() else "TPU"
115
    # Create placement group for worker processes
116
117
118
119
120
    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.
121
        device_bundles = 0
122
        for bundle in bundles:
123
124
            bundle_devices = bundle.get(device_str, 0)
            if bundle_devices > 1:
125
                raise ValueError(
126
127
128
129
130
                    "Placement group bundle cannot have more than 1 "
                    f"{device_str}.")
            if bundle_devices:
                device_bundles += 1
        if parallel_config.world_size > device_bundles:
131
            raise ValueError(
132
133
134
135
                f"The number of required {device_str}s exceeds the total "
                f"number of available {device_str}s in the placement group."
                f"Required number of devices: {parallel_config.world_size}. "
                f"Total number of devices: {device_bundles}.")
136
    else:
137
138
        num_devices_in_cluster = ray.cluster_resources().get(device_str, 0)
        if parallel_config.world_size > num_devices_in_cluster:
139
            raise ValueError(
140
141
                f"The number of required {device_str}s exceeds the total "
                f"number of available {device_str}s in the placement group.")
142
        # Create a new placement group
143
144
145
        placement_group_specs = ([{
            device_str: 1
        }] * parallel_config.world_size)
146
147
        current_placement_group = ray.util.placement_group(
            placement_group_specs)
148
149
150
151
152
        # 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)

153
154
    # Set the placement group in the parallel config
    parallel_config.placement_group = current_placement_group