| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import copy |
| 4 | import os |
| 5 | import threading |
| 6 | import weakref |
| 7 | from collections import defaultdict, deque |
| 8 | from dataclasses import dataclass |
| 9 | from typing import Any |
| 10 | |
| 11 | from torch.distributed import TCPStore |
| 12 | |
| 13 | import vllm.envs as envs |
| 14 | from vllm.config import VllmConfig |
| 15 | from vllm.distributed.device_communicators.shm_broadcast import ( |
| 16 | Handle, |
| 17 | MessageQueue, |
| 18 | ) |
| 19 | from vllm.logger import configure_logging, init_logger |
| 20 | from vllm.platforms import current_platform |
| 21 | from vllm.utils.network_utils import ( |
| 22 | get_distributed_init_method, |
| 23 | ) |
| 24 | from vllm.v1.executor.multiproc_executor import ( |
| 25 | FutureWrapper, |
| 26 | MultiprocExecutor, |
| 27 | WorkerProc, |
| 28 | ) |
| 29 | from vllm.v1.executor.ray_env_utils import ( |
| 30 | get_driver_env_vars, |
| 31 | update_runtime_env_for_worker_import, |
| 32 | ) |
| 33 | from vllm.v1.executor.ray_utils import ( |
| 34 | WORKER_SPECIFIC_ENV_VARS, |
| 35 | build_actor_name, |
| 36 | get_bundles_for_indices, |
| 37 | get_bundles_sorted_by_node, |
| 38 | initialize_ray_cluster, |
| 39 | ray, |
| 40 | ) |
| 41 | |
| 42 | if ray is not None: |
| 43 | from ray.actor import ActorHandle |
| 44 | from ray.types import ObjectRef |
| 45 | from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy |
| 46 | else: |
| 47 | ActorHandle = None |
| 48 | |
| 49 | logger = init_logger(__name__) |
| 50 | |
| 51 | |
| 52 | @dataclass |
| 53 | class RayWorkerHandle: |
| 54 | """Handle for a Ray worker actor, compatible with MultiprocExecutor.""" |
| 55 | |
| 56 | actor: ActorHandle |
| 57 | """Ray worker actor""" |
| 58 | |
| 59 | rank: int |
| 60 | """Rank of the worker""" |
| 61 | |
| 62 | local_rank: int |
| 63 | """Local rank of the worker""" |
| 64 | |
| 65 | node_id: str |
| 66 | """Node ID of the worker""" |
| 67 | |
| 68 | bundle_id_idx: int = -1 |
| 69 | """Placement group bundle index for the worker""" |
| 70 | |
| 71 | run_ref: ObjectRef | None = None |
| 72 | """run() ObjectRef used as a sentinel for health monitoring""" |
| 73 | |
| 74 | def run(self): |
| 75 | """Start the worker's busy loop.""" |
| 76 | self.run_ref = self.actor.run.remote() |
| 77 | |
| 78 | |
| 79 | class RayWorkerProc(WorkerProc): |
| 80 | """Worker process that runs inside a Ray actor. |
| 81 | |
| 82 | Initialization is split into two phases: |
| 83 | 1. __init__: lightweight setup, stores init args (no device/model init) |
| 84 | 2. initialize_worker: called after GPU IDs are discovered, completes |
| 85 | the full WorkerProc initialization with the correct local_rank and |
| 86 | logical-to-physical GPU mapping. |
| 87 | |
| 88 | GPU assignment flow: |
| 89 | |
| 90 | 1. RayExecutorV2 enables RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES so Ray does |
| 91 | not set CUDA_VISIBLE_DEVICES on RayWorkerProc actors at creation time. |
| 92 | 2. Each actor is scheduled with a placement group and bundle index; Ray resolves |
| 93 | the physical GPU ID for that bundle at placement time. |
| 94 | 3. After placement, the executor discovers each worker's GPU ID and passes the |
| 95 | node's logical-to-physical mapping (assigned_physical_gpu_ids) to |
| 96 | initialize_worker(); CUDA_VISIBLE_DEVICES is never modified. |
| 97 | |
| 98 | Scheduling must complete before the mapping is known when the placement |
| 99 | group is externally managed: only then is the GPU tied to the worker's |
| 100 | bundle resolved. |
| 101 | |
| 102 | This sequence allows multiple vLLM instances to coexist on the same node: |
| 103 | each instance is unaware which physical devices others hold, and the |
| 104 | externally managed placement group avoids device assignment conflicts |
| 105 | by binding workers to specific placement group bundles. |
| 106 | """ |
| 107 | |
| 108 | def __init__( |
| 109 | self, |
| 110 | vllm_config: VllmConfig, |
| 111 | rank: int, |
| 112 | input_shm_handle: Handle, |
| 113 | is_driver_worker: bool, |
| 114 | is_driver_node: bool = False, |
| 115 | ): |
| 116 | if logging_config := getattr(vllm_config, "logging_config", None): |
| 117 | configure_logging(logging_config) |
| 118 | |
| 119 | # Defer WorkerProc.__init__ until GPU IDs are known. |
| 120 | self._is_driver_node = is_driver_node |
| 121 | self._parallel_config = vllm_config.parallel_config |
| 122 | self._init_kwargs = dict( |
| 123 | vllm_config=vllm_config, |
| 124 | rank=rank, |
| 125 | input_shm_handle=input_shm_handle, |
| 126 | shared_worker_lock=None, |
| 127 | is_driver_worker=is_driver_worker, |
| 128 | ) |
| 129 | |
| 130 | def create_dist_init_method(self) -> str: |
| 131 | host = ray.util.get_node_ip_address() |
| 132 | store = TCPStore( |
| 133 | host_name=host, |
| 134 | port=0, |
| 135 | world_size=self._parallel_config.world_size, |
| 136 | is_master=True, |
| 137 | wait_for_workers=False, |
| 138 | multi_tenant=True, |
| 139 | ) |
| 140 | # Keep the bound server alive until init_process_group reuses it. |
| 141 | self._dist_init_store = store |
| 142 | return get_distributed_init_method(host, store.port) |
| 143 | |
| 144 | def get_node_and_physical_gpu_ids(self) -> tuple[str, list[int]]: |
| 145 | """Return (node_id, physical_gpu_ids) assigned to this actor by Ray.""" |
| 146 | node_id = ray.get_runtime_context().get_node_id() |
| 147 | device_key = current_platform.ray_device_key |
| 148 | if not device_key: |
| 149 | raise RuntimeError( |
| 150 | f"current platform {current_platform.device_name} does not support ray." |
| 151 | ) |
| 152 | physical_gpu_ids = ray.get_runtime_context().get_accelerator_ids()[device_key] |
| 153 | return node_id, [ |
| 154 | current_platform.device_control_id_to_physical_device_id(str(x)) |
| 155 | for x in physical_gpu_ids |
| 156 | ] |
| 157 | |
| 158 | def initialize_worker( |
| 159 | self, |
| 160 | local_rank: int, |
| 161 | env_vars: dict[str, str], |
| 162 | distributed_init_method: str, |
| 163 | driver_env_vars: dict[str, str] | None = None, |
| 164 | assigned_physical_gpu_ids: list[int] | None = None, |
| 165 | ) -> None: |
| 166 | """Complete initialization after GPU assignment is known. |
| 167 | |
| 168 | *driver_env_vars* are applied with ``setdefault`` — they fill |
| 169 | in missing vars but never overwrite node-local values. |
| 170 | *env_vars* always overwrite. |
| 171 | *assigned_physical_gpu_ids* maps local_rank to physical CUDA device ID. |
| 172 | """ |
| 173 | if driver_env_vars: |
| 174 | for key, value in driver_env_vars.items(): |
| 175 | os.environ.setdefault(key, value) |
| 176 | for key, value in env_vars.items(): |
| 177 | os.environ[key] = value |
| 178 | |
| 179 | if assigned_physical_gpu_ids is not None: |
| 180 | vllm_config = self._init_kwargs["vllm_config"] |
| 181 | assert isinstance(vllm_config, VllmConfig) |
| 182 | vllm_config.parallel_config.assigned_physical_gpu_ids = ( |
| 183 | assigned_physical_gpu_ids |
| 184 | ) |
| 185 | |
| 186 | self._init_kwargs["distributed_init_method"] = distributed_init_method |
| 187 | self.local_rank = local_rank |
| 188 | super().__init__( |
| 189 | local_rank=local_rank, |
| 190 | **self._init_kwargs, |
| 191 | ) |
| 192 | |
| 193 | def _init_message_queues( |
| 194 | self, input_shm_handle: Handle, vllm_config: VllmConfig |
| 195 | ) -> None: |
| 196 | """Workers on the same node as the executor use shared memory for |
| 197 | both the broadcast (input) MQ and the response MQ. Workers on |
| 198 | different nodes use TCP (n_local_reader=0). |
| 199 | """ |
| 200 | self.rpc_broadcast_mq = MessageQueue.create_from_handle( |
| 201 | input_shm_handle, self.worker.rank |
| 202 | ) |
| 203 | |
| 204 | n_local = 1 if self._is_driver_node else 0 |
| 205 | # Use ray.util.get_node_ip_address() to get Ray's internal IP. |
| 206 | # get_ip() returns host's external IP which is typically not |
| 207 | # routable between nodes within the cluster. |
| 208 | self.worker_response_mq = MessageQueue( |
| 209 | n_reader=1, |
| 210 | n_local_reader=n_local, |
| 211 | connect_ip=ray.util.get_node_ip_address(), |
| 212 | ) |
| 213 | self.peer_response_handles: list[dict] = [] |
| 214 | |
| 215 | def wait_for_init(self) -> dict: |
| 216 | """Respond to the driver's wait_until_ready() barrier.""" |
| 217 | assert self.worker_response_mq is not None |
| 218 | return { |
| 219 | "status": self.READY_STR, |
| 220 | "handle": self.worker_response_mq.export_handle(), |
| 221 | } |
| 222 | |
| 223 | def run(self) -> None: |
| 224 | """Main entry point called via actor.run.remote().""" |
| 225 | try: |
| 226 | assert self.rpc_broadcast_mq is not None |
| 227 | self.rpc_broadcast_mq.wait_until_ready() |
| 228 | assert self.worker_response_mq is not None |
| 229 | self.worker_response_mq.wait_until_ready() |
| 230 | |
| 231 | self.worker_busy_loop() |
| 232 | except Exception as e: |
| 233 | logger.exception("RayWorkerProc failed: %s", e) |
| 234 | raise |
| 235 | finally: |
| 236 | self.shutdown() |
| 237 | |
| 238 | |
| 239 | class RayExecutorV2(MultiprocExecutor): |
| 240 | """Ray-based distributed executor using MessageQueue communication. |
| 241 | |
| 242 | Inherits from MultiprocExecutor to reuse the MQ-based control plane |
| 243 | and NCCL data plane. Workers are Ray actors. |
| 244 | |
| 245 | Async scheduling is enabled, inherited from MultiprocExecutor. |
| 246 | This is cricitcal for RayExecutorV2 to be performant. |
| 247 | """ |
| 248 | |
| 249 | uses_ray: bool = True |
| 250 | supports_pp: bool = True |
| 251 | |
| 252 | def __init__(self, vllm_config: VllmConfig): |
| 253 | super().__init__(vllm_config) |
| 254 | |
| 255 | def _build_runtime_env(self) -> dict: |
| 256 | """Build a runtime_env dict for RayWorkerProc actors. |
| 257 | |
| 258 | Driver env vars are applied separately via initialize_worker |
| 259 | with setdefault semantics. |
| 260 | """ |
| 261 | base = self.parallel_config.ray_runtime_env |
| 262 | runtime_env: dict = copy.deepcopy(dict(base)) if base else {} |
| 263 | |
| 264 | env_vars = runtime_env.setdefault("env_vars", {}) |
| 265 | env_vars.update({v: "1" for v in current_platform.ray_noset_device_env_vars}) |
| 266 | update_runtime_env_for_worker_import(runtime_env) |
| 267 | if self.parallel_config.ray_workers_use_nsight: |
| 268 | runtime_env["nsight"] = { |
| 269 | "t": "cuda,cudnn,cublas", |
| 270 | "o": "'worker_process_%p'", |
| 271 | "cuda-graph-trace": "node", |
| 272 | } |
| 273 | return runtime_env |
| 274 | |
| 275 | @staticmethod |
| 276 | def _get_actor_resource_kwargs() -> dict[str, Any]: |
| 277 | """Return Ray actor resource kwargs for the current platform.""" |
| 278 | num_devices = envs.VLLM_RAY_PER_WORKER_GPUS |
| 279 | device_key = current_platform.ray_device_key |
| 280 | if device_key == "GPU": |
| 281 | return {"num_gpus": num_devices} |
| 282 | return {"num_gpus": 0, "resources": {device_key: num_devices}} |
| 283 | |
| 284 | def _init_executor(self) -> None: |
| 285 | """Initialize the RayExecutorV2 executor.""" |
| 286 | self._finalizer = weakref.finalize(self, self.shutdown) |
| 287 | self.is_failed = False |
| 288 | self.failure_callback = None |
| 289 | self.shutting_down = False |
| 290 | self.shutdown_lock = threading.Lock() |
| 291 | |
| 292 | # Step 1: Initialize Ray cluster and retrieve placement group |
| 293 | if ray is None: |
| 294 | raise ImportError("Using Ray backend requires installation of ray.") |
| 295 | initialize_ray_cluster(self.parallel_config, require_gpu_on_driver=False) |
| 296 | placement_group = self.parallel_config.placement_group |
| 297 | |
| 298 | tp_size, pp_size, pcp_size = self._get_parallel_sizes() |
| 299 | assert self.world_size == tp_size * pp_size * pcp_size, ( |
| 300 | f"world_size ({self.world_size}) must be equal to the " |
| 301 | f"tensor_parallel_size ({tp_size}) x pipeline" |
| 302 | f"_parallel_size ({pp_size}) x prefill_context" |
| 303 | f"_parallel_size ({pcp_size}). " |
| 304 | ) |
| 305 | |
| 306 | # Step 2: Build bundle assignments for worker rank placement |
| 307 | # while respecting VLLM_RAY_BUNDLE_INDICES. |
| 308 | if envs.VLLM_RAY_BUNDLE_INDICES: |
| 309 | bundle_to_node_id = get_bundles_for_indices( |
| 310 | placement_group, |
| 311 | list(map(int, envs.VLLM_RAY_BUNDLE_INDICES.split(","))), |
| 312 | self.world_size, |
| 313 | ) |
| 314 | else: |
| 315 | bundle_to_node_id = get_bundles_sorted_by_node(placement_group) |
| 316 | driver_node = ray.get_runtime_context().get_node_id() |
| 317 | |
| 318 | bundle_assignments: list[dict[str, Any]] = [] |
| 319 | for rank, (bundle_id_idx, node_id, node_ip) in enumerate(bundle_to_node_id): |
| 320 | bundle_assignments.append( |
| 321 | { |
| 322 | "rank": rank, |
| 323 | "bundle_id_idx": bundle_id_idx, |
| 324 | "node_id": node_id, |
| 325 | "node_ip": node_ip, |
| 326 | } |
| 327 | ) |
| 328 | |
| 329 | # Step 3: Create broadcast MessageQueue. |
| 330 | # Workers on the driver node use shared memory; the rest use TCP. |
| 331 | max_chunk_bytes = envs.VLLM_MQ_MAX_CHUNK_BYTES_MB * 1024 * 1024 |
| 332 | n_local = sum(1 for a in bundle_assignments if a["node_id"] == driver_node) |
| 333 | self.rpc_broadcast_mq = MessageQueue( |
| 334 | self.world_size, |
| 335 | n_local, |
| 336 | max_chunk_bytes=max_chunk_bytes, |
| 337 | connect_ip=ray.util.get_node_ip_address(), |
| 338 | ) |
| 339 | scheduler_output_handle = self.rpc_broadcast_mq.export_handle() |
| 340 | |
| 341 | # Step 4: Spawn RayWorkerProc actors into PG bundles (deferred init). |
| 342 | # Workers are created lightweight here; full initialization happens |
| 343 | # in Step 6 after GPU IDs are discovered. |
| 344 | self.ray_worker_handles: list[RayWorkerHandle] = [] |
| 345 | instance_id = self.vllm_config.instance_id |
| 346 | |
| 347 | # Collect driver env vars and apply but don't overwrite node-local values. |
| 348 | self.driver_env_vars = get_driver_env_vars( |
| 349 | worker_specific_vars=WORKER_SPECIFIC_ENV_VARS, |
| 350 | ) |
| 351 | |
| 352 | runtime_env = self._build_runtime_env() |
| 353 | resource_kwargs = self._get_actor_resource_kwargs() |
| 354 | |
| 355 | # The sharded_rdt weight transfer backend pulls weights via Ray's tensor |
| 356 | # transport (@ray.method(tensor_transport=...)). Ray requires the |
| 357 | # *calling* actor (the vLLM worker) to opt in via enable_tensor_transport. |
| 358 | wt_cfg = self.vllm_config.weight_transfer_config |
| 359 | extra_actor_options: dict[str, object] = {} |
| 360 | if wt_cfg is not None and wt_cfg.backend == "sharded_rdt": |
| 361 | from vllm.distributed.weight_transfer.sharded_rdt_common import ( |
| 362 | check_ray_rdt_version, |
| 363 | ) |
| 364 | |
| 365 | check_ray_rdt_version() |
| 366 | extra_actor_options["enable_tensor_transport"] = True |
| 367 | |
| 368 | for bundle_idx in range(self.world_size): |
| 369 | bundle = bundle_assignments[bundle_idx] |
| 370 | is_driver_worker = self._is_driver_worker(bundle["rank"]) |
| 371 | is_driver_node = bundle["node_id"] == driver_node |
| 372 | |
| 373 | scheduling_strategy = PlacementGroupSchedulingStrategy( |
| 374 | placement_group=placement_group, |
| 375 | placement_group_bundle_index=bundle["bundle_id_idx"], |
| 376 | ) |
| 377 | |
| 378 | actor_name = build_actor_name( |
| 379 | instance_id, bundle["rank"], tp_size, pp_size, pcp_size |
| 380 | ) |
| 381 | |
| 382 | actor = ( |
| 383 | ray.remote(RayWorkerProc) |
| 384 | .options( |
| 385 | name=actor_name, |
| 386 | num_cpus=0, |
| 387 | **resource_kwargs, |
| 388 | scheduling_strategy=scheduling_strategy, |
| 389 | runtime_env=runtime_env, |
| 390 | **extra_actor_options, |
| 391 | ) |
| 392 | .remote( |
| 393 | vllm_config=self.vllm_config, |
| 394 | rank=bundle["rank"], |
| 395 | input_shm_handle=scheduler_output_handle, |
| 396 | is_driver_worker=is_driver_worker, |
| 397 | is_driver_node=is_driver_node, |
| 398 | ) |
| 399 | ) |
| 400 | |
| 401 | handle = RayWorkerHandle( |
| 402 | actor=actor, |
| 403 | rank=bundle["rank"], |
| 404 | local_rank=-1, # Set in Step 7 after GPU ID discovery |
| 405 | node_id=bundle["node_id"], |
| 406 | bundle_id_idx=bundle["bundle_id_idx"], |
| 407 | ) |
| 408 | self.ray_worker_handles.append(handle) |
| 409 | |
| 410 | # Step 5: Discover physical GPU IDs assigned to each worker via Ray |
| 411 | # runtime context. |
| 412 | worker_node_and_physical_gpu_ids = ray.get( |
| 413 | [ |
| 414 | h.actor.get_node_and_physical_gpu_ids.remote() |
| 415 | for h in self.ray_worker_handles |
| 416 | ] |
| 417 | ) |
| 418 | |
| 419 | node_workers: dict[str, list[int]] = defaultdict(list) |
| 420 | node_physical_gpu_ids: dict[str, list[int]] = defaultdict(list) |
| 421 | for i, (node_id, physical_gpu_ids) in enumerate( |
| 422 | worker_node_and_physical_gpu_ids |
| 423 | ): |
| 424 | node_workers[node_id].append(i) |
| 425 | node_physical_gpu_ids[node_id].extend(physical_gpu_ids) |
| 426 | for node_id, physical_gpu_ids in node_physical_gpu_ids.items(): |
| 427 | node_physical_gpu_ids[node_id] = sorted(physical_gpu_ids) |
| 428 | |
| 429 | # Step 6: Initialize workers with local logical ranks and the |
| 430 | # logical-to-physical GPU mapping discovered from Ray placement. |
| 431 | distributed_init_method = ray.get( |
| 432 | self.ray_worker_handles[0].actor.create_dist_init_method.remote() |
| 433 | ) |
| 434 | init_worker_refs = [] |
| 435 | for i, (node_id, _) in enumerate(worker_node_and_physical_gpu_ids): |
| 436 | local_rank = node_workers[node_id].index(i) |
| 437 | assigned_physical_gpu_ids = sorted(node_physical_gpu_ids[node_id]) |
| 438 | worker_env_vars: dict[str, str] = {} |
| 439 | self.ray_worker_handles[i].local_rank = local_rank |
| 440 | init_worker_refs.append( |
| 441 | self.ray_worker_handles[i].actor.initialize_worker.remote( |
| 442 | local_rank, |
| 443 | worker_env_vars, |
| 444 | distributed_init_method, |
| 445 | self.driver_env_vars, |
| 446 | assigned_physical_gpu_ids=assigned_physical_gpu_ids, |
| 447 | ) |
| 448 | ) |
| 449 | # Also set on the executor-side config for consistency. The mapping |
| 450 | # is per-node, so only do this when all workers share one node. |
| 451 | if len(node_physical_gpu_ids) == 1: |
| 452 | node_id_0 = worker_node_and_physical_gpu_ids[0][0] |
| 453 | self.vllm_config.parallel_config.assigned_physical_gpu_ids = sorted( |
| 454 | node_physical_gpu_ids[node_id_0] |
| 455 | ) |
| 456 | ray.get(init_worker_refs) |
| 457 | |
| 458 | # Step 7: Collect response MQ handles |
| 459 | init_results = ray.get( |
| 460 | [h.actor.wait_for_init.remote() for h in self.ray_worker_handles] |
| 461 | ) |
| 462 | |
| 463 | self.response_mqs: list[MessageQueue] = [] |
| 464 | for i, result in enumerate(init_results): |
| 465 | if result["status"] != RayWorkerProc.READY_STR: |
| 466 | raise RuntimeError(f"Worker {i} failed to initialize: {result}") |
| 467 | self.response_mqs.append( |
| 468 | MessageQueue.create_from_handle(result["handle"], 0) |
| 469 | ) |
| 470 | |
| 471 | # Step 8: Start run() before wait_until_ready() to avoid |
| 472 | # deadlock — workers send subscriptions inside run(). |
| 473 | for handle in self.ray_worker_handles: |
| 474 | handle.run() |
| 475 | |
| 476 | # Step 9: wait_until_ready() barrier |
| 477 | self.rpc_broadcast_mq.wait_until_ready() |
| 478 | for response_mq in self.response_mqs: |
| 479 | response_mq.wait_until_ready() |
| 480 | |
| 481 | self.futures_queue = deque[FutureWrapper]() |
| 482 | self._post_init_executor() |
| 483 | |
| 484 | self.start_worker_monitor() |
| 485 | self.output_rank = self._get_output_rank() |
| 486 | |
| 487 | def start_worker_monitor(self, inline=False) -> None: |
| 488 | """Monitor worker liveness via ray.wait() on run() ObjectRefs.""" |
| 489 | run_refs = [h.run_ref for h in self.ray_worker_handles if h.run_ref is not None] |
| 490 | if not run_refs: |
| 491 | raise RuntimeError("Ray workers have not started successfully.") |
| 492 | |
| 493 | self_ref = weakref.ref(self) |
| 494 | ref_to_rank = { |
| 495 | h.run_ref: h.rank for h in self.ray_worker_handles if h.run_ref is not None |
| 496 | } |
| 497 | |
| 498 | def _should_stop() -> bool: |
| 499 | executor = self_ref() |
| 500 | return not executor or executor.shutting_down |
| 501 | |
| 502 | def monitor_workers(): |
| 503 | # Poll with a timeout rather than blocking on ray.wait() |
| 504 | # because a blocking call would segfault if Ray is torn down |
| 505 | # while this thread is inside it. |
| 506 | while not _should_stop() and ray.is_initialized(): |
| 507 | try: |
| 508 | done, _ = ray.wait(run_refs, num_returns=1, timeout=5.0) |
| 509 | except Exception: |
| 510 | logger.exception( |
| 511 | "RayWorkerMonitor: unexpected error, exiting monitor thread" |
| 512 | ) |
| 513 | return |
| 514 | if not done or _should_stop(): |
| 515 | continue |
| 516 | |
| 517 | dead_ranks = [ref_to_rank[r] for r in done] |
| 518 | executor = self_ref() |
| 519 | if not executor: |
| 520 | return |
| 521 | executor.is_failed = True |
| 522 | logger.error( |
| 523 | "RayWorkerProc rank=%s died unexpectedly, shutting down executor.", |
| 524 | dead_ranks, |
| 525 | ) |
| 526 | executor.shutdown() |
| 527 | if executor.failure_callback is not None: |
| 528 | callback = executor.failure_callback |
| 529 | executor.failure_callback = None |
| 530 | callback() |
| 531 | return |
| 532 | |
| 533 | t = threading.Thread( |
| 534 | target=monitor_workers, daemon=True, name="RayWorkerMonitor" |
| 535 | ) |
| 536 | t.start() |
| 537 | self._monitor_thread = t |
| 538 | |
| 539 | def _join_monitor_thread(self) -> None: |
| 540 | """Wait for the monitor thread to exit. |
| 541 | |
| 542 | Must be called before tearing down Ray resources — the monitor |
| 543 | may be inside ray.wait() which would segfault if Ray is shut |
| 544 | down underneath it. When the monitor itself calls shutdown() |
| 545 | on worker death, we skip the join because the thread is about |
| 546 | to return anyway. |
| 547 | """ |
| 548 | monitor = getattr(self, "_monitor_thread", None) |
| 549 | if ( |
| 550 | monitor is not None |
| 551 | and monitor.is_alive() |
| 552 | and threading.current_thread() is not monitor |
| 553 | ): |
| 554 | monitor.join(timeout=10) |
| 555 | |
| 556 | def shutdown(self) -> None: |
| 557 | """Properly shut down the executor and its workers.""" |
| 558 | lock = getattr(self, "shutdown_lock", None) |
| 559 | if lock is None: |
| 560 | return |
| 561 | |
| 562 | with lock: |
| 563 | if getattr(self, "shutting_down", False): |
| 564 | return |
| 565 | self.shutting_down = True |
| 566 | |
| 567 | self._join_monitor_thread() |
| 568 | |
| 569 | for handle in getattr(self, "ray_worker_handles", []): |
| 570 | try: |
| 571 | ray.kill(handle.actor) |
| 572 | logger.debug("Killed actor rank=%d", handle.rank) |
| 573 | except Exception: |
| 574 | logger.exception("Failed to kill actor rank=%d", handle.rank) |
| 575 | |
| 576 | if rpc_broadcast_mq := getattr(self, "rpc_broadcast_mq", None): |
| 577 | rpc_broadcast_mq.shutdown() |
| 578 | self.rpc_broadcast_mq = None |
| 579 | |
| 580 | for mq in getattr(self, "response_mqs", []): |
| 581 | mq.shutdown() |
| 582 | self.response_mqs = [] |