| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """EngineCoreSentinel and fault_tolerant_wrapper for the engine core.""" |
| 4 | |
| 5 | import json |
| 6 | import threading |
| 7 | from collections.abc import Callable |
| 8 | from typing import TYPE_CHECKING |
| 9 | |
| 10 | import msgspec |
| 11 | |
| 12 | from vllm.distributed import stateless_destroy_torch_distributed_process_group |
| 13 | from vllm.distributed.utils import stateless_init_torch_distributed_process_group |
| 14 | from vllm.logger import init_logger |
| 15 | from vllm.utils.network_utils import get_open_port |
| 16 | from vllm.v1.engine import ( |
| 17 | FT_STATUS_CALL_ID, |
| 18 | EngineCoreOutputs, |
| 19 | EngineStatusType, |
| 20 | UtilityOutput, |
| 21 | ) |
| 22 | from vllm.v1.fault_tolerance.utils import FaultToleranceRequest, FaultToleranceResult |
| 23 | from vllm.v1.request import RequestStatus |
| 24 | from vllm.v1.serial_utils import UtilityResult, run_method |
| 25 | |
| 26 | if TYPE_CHECKING: |
| 27 | from vllm.v1.engine.core import EngineCoreProc |
| 28 | |
| 29 | logger = init_logger(__name__) |
| 30 | |
| 31 | FT_UTILITY_METHOD = "handle_fault_tolerance" |
| 32 | |
| 33 | |
| 34 | class EngineCoreSentinel: |
| 35 | """Manages fault tolerance state for a single engine core.""" |
| 36 | |
| 37 | def __init__(self, engine: "EngineCoreProc", parallel_config): |
| 38 | self.engine = engine |
| 39 | self.engine_index = engine.engine_index |
| 40 | self.parallel_config = parallel_config |
| 41 | ft_config = parallel_config.fault_tolerance_config |
| 42 | self.engine_recovery_timeout_sec = ft_config.engine_recovery_timeout_sec |
| 43 | |
| 44 | self.resumed = threading.Event() |
| 45 | self.resumed.set() |
| 46 | self.status_type = EngineStatusType.HEALTHY |
| 47 | self.fault_info: str | None = None |
| 48 | self._dp_reinit_epoch = 0 |
| 49 | |
| 50 | def handle_command(self, client_idx: int, call_id: int, ft_args: dict): |
| 51 | """Dispatch an FT command by instruction name.""" |
| 52 | ft_request = FaultToleranceRequest(**ft_args) |
| 53 | if self.status_type != EngineStatusType.UNHEALTHY: |
| 54 | reason = ( |
| 55 | f"[FT] Rejecting {ft_request.instruction} on engine " |
| 56 | f"{self.engine_index}: status is {self.status_type.name}" |
| 57 | ) |
| 58 | logger.warning(reason) |
| 59 | result = FaultToleranceResult( |
| 60 | request_id=ft_request.request_id, |
| 61 | success=False, |
| 62 | reason=reason, |
| 63 | ) |
| 64 | else: |
| 65 | try: |
| 66 | result = run_method(self, ft_request.instruction, (ft_request,), {}) |
| 67 | except Exception as e: |
| 68 | logger.exception("[FT] Instruction '%s' failed", ft_request.instruction) |
| 69 | result = FaultToleranceResult( |
| 70 | request_id=ft_request.request_id, success=False, reason=str(e) |
| 71 | ) |
| 72 | |
| 73 | uo = UtilityOutput(call_id) |
| 74 | uo.result = UtilityResult(msgspec.structs.asdict(result)) |
| 75 | self.engine.output_queue.put_nowait( |
| 76 | (client_idx, EngineCoreOutputs(utility_output=uo)) |
| 77 | ) |
| 78 | |
| 79 | def on_fault(self, exc: Exception): |
| 80 | """Called by the wrapper when the busy loop raises an exception.""" |
| 81 | self.resumed.clear() |
| 82 | logger.warning( |
| 83 | "[FT] Busy loop raised %s. Waiting for recovery.", type(exc).__name__ |
| 84 | ) |
| 85 | |
| 86 | engine = self.engine |
| 87 | aborted = engine.scheduler.finish_requests(None, RequestStatus.FINISHED_ABORTED) |
| 88 | engine._send_abort_outputs(aborted) |
| 89 | if engine.batch_queue is not None: |
| 90 | engine.batch_queue.clear() |
| 91 | if ( |
| 92 | hasattr(engine.model_executor, "is_failed") |
| 93 | and engine.model_executor.is_failed |
| 94 | ): |
| 95 | self.status_type = EngineStatusType.DEAD |
| 96 | else: |
| 97 | self.status_type = EngineStatusType.UNHEALTHY |
| 98 | self.fault_info = f"{type(exc).__name__}" |
| 99 | logger.info( |
| 100 | "[FT] Engine %d status -> %s:", |
| 101 | self.engine_index, |
| 102 | self.status_type.name, |
| 103 | exc_info=exc, |
| 104 | ) |
| 105 | self._push_status() |
| 106 | |
| 107 | def _push_status(self): |
| 108 | """Push current health to the client so it can refresh its cache.""" |
| 109 | payload = {"id": self.engine_index, "status": self.status_type.name.lower()} |
| 110 | if self.status_type == EngineStatusType.UNHEALTHY: |
| 111 | payload["fault_info"] = self.fault_info |
| 112 | outputs = EngineCoreOutputs( |
| 113 | utility_output=UtilityOutput( |
| 114 | call_id=FT_STATUS_CALL_ID, |
| 115 | result=UtilityResult(payload), |
| 116 | ) |
| 117 | ) |
| 118 | outputs.engine_index = self.engine_index |
| 119 | self.engine.output_queue.put_nowait((0, outputs)) |
| 120 | |
| 121 | def retry(self, ft_request: FaultToleranceRequest) -> FaultToleranceResult: |
| 122 | engine = self.engine |
| 123 | executor = engine.model_executor |
| 124 | |
| 125 | ft_request.params.update(self._reinit_dp_group()) |
| 126 | if hasattr(engine, "step_counter"): |
| 127 | engine.step_counter = 0 |
| 128 | |
| 129 | executor.collective_rpc("handle_ft_command", args=(ft_request,)) |
| 130 | |
| 131 | self.status_type = EngineStatusType.HEALTHY |
| 132 | logger.info("[FT] Engine %d status -> HEALTHY", self.engine_index) |
| 133 | self.resumed.set() |
| 134 | self._push_status() |
| 135 | return FaultToleranceResult(request_id=ft_request.request_id, success=True) |
| 136 | |
| 137 | def _reinit_dp_group(self) -> dict: |
| 138 | """Reinit DP process group if in DP mode. Returns worker params.""" |
| 139 | engine = self.engine |
| 140 | if not hasattr(engine, "dp_group") or not hasattr(engine, "dp_store"): |
| 141 | return {} |
| 142 | |
| 143 | parallel_config = engine.vllm_config.parallel_config |
| 144 | worker_key = f"ft_worker_dp_ports_{self._dp_reinit_epoch}" |
| 145 | engine_key = f"ft_engine_dp_port_{self._dp_reinit_epoch}" |
| 146 | self._dp_reinit_epoch += 1 |
| 147 | |
| 148 | if parallel_config.data_parallel_rank == 0: |
| 149 | worker_ports = [get_open_port() for _ in range(parallel_config.world_size)] |
| 150 | engine_port = get_open_port() |
| 151 | engine.dp_store.set(worker_key, json.dumps(worker_ports).encode()) |
| 152 | engine.dp_store.set(engine_key, str(engine_port).encode()) |
| 153 | else: |
| 154 | worker_ports = json.loads(engine.dp_store.get(worker_key).decode()) |
| 155 | engine_port = int(engine.dp_store.get(engine_key).decode()) |
| 156 | |
| 157 | stateless_destroy_torch_distributed_process_group(engine.dp_group) |
| 158 | engine.dp_group, engine.dp_store = ( |
| 159 | stateless_init_torch_distributed_process_group( |
| 160 | parallel_config.data_parallel_master_ip, |
| 161 | engine_port, |
| 162 | parallel_config.data_parallel_rank, |
| 163 | parallel_config.data_parallel_size, |
| 164 | backend="gloo", |
| 165 | return_store=True, |
| 166 | timeout=parallel_config.cpu_distributed_timeout, |
| 167 | ) |
| 168 | ) |
| 169 | return {"new_stateless_dp_group_ports": worker_ports} |
| 170 | |
| 171 | |
| 172 | def fault_tolerant_wrapper(busy_loop_func: Callable): |
| 173 | """Wrap the busy loop to catch faults and delegate recovery.""" |
| 174 | |
| 175 | def run_with_fault_tolerance(self: "EngineCoreProc"): |
| 176 | while True: |
| 177 | try: |
| 178 | busy_loop_func(self) |
| 179 | except SystemExit: |
| 180 | raise |
| 181 | except Exception as exc: |
| 182 | if not self.enable_fault_tolerance: |
| 183 | raise |
| 184 | self.ft_sentinel.on_fault(exc) |
| 185 | recovered = self.ft_sentinel.resumed.wait( |
| 186 | timeout=self.ft_sentinel.engine_recovery_timeout_sec |
| 187 | ) |
| 188 | if recovered: |
| 189 | continue |
| 190 | logger.error( |
| 191 | "[FT] No recovery within %ds timeout.", |
| 192 | self.ft_sentinel.engine_recovery_timeout_sec, |
| 193 | ) |
| 194 | raise |
| 195 | |
| 196 | return run_with_fault_tolerance |