| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from typing import TYPE_CHECKING, cast |
| 4 | |
| 5 | import torch |
| 6 | |
| 7 | from vllm.config import set_current_vllm_config |
| 8 | from vllm.distributed import ( |
| 9 | get_dp_group, |
| 10 | stateless_destroy_torch_distributed_process_group, |
| 11 | stateless_init_torch_distributed_process_group, |
| 12 | ) |
| 13 | from vllm.logger import init_logger |
| 14 | from vllm.model_executor.layers.fused_moe.all2all_utils import get_ep_all2all_manager |
| 15 | from vllm.v1.fault_tolerance.utils import FaultToleranceRequest |
| 16 | from vllm.v1.serial_utils import run_method |
| 17 | |
| 18 | if TYPE_CHECKING: |
| 19 | from vllm.v1.worker.gpu.model_runner import GPUModelRunner as GPUModelRunnerV2 |
| 20 | from vllm.v1.worker.gpu_worker import Worker |
| 21 | |
| 22 | logger = init_logger(__name__) |
| 23 | |
| 24 | # All2all backends that support fault-tolerant timeout + rank masking, |
| 25 | # required for FT under DP+EP MoE deployments. |
| 26 | FT_BACKEND_SET = frozenset({"deepep_low_latency", "nixl_ep"}) |
| 27 | |
| 28 | |
| 29 | class WorkerSentinel: |
| 30 | """Holds FT state for a single worker (mask tensors, DP config). |
| 31 | |
| 32 | Methods are called via collective_rpc from EngineCoreSentinel. |
| 33 | """ |
| 34 | |
| 35 | def __init__(self, worker: "Worker"): |
| 36 | self.worker = worker |
| 37 | self.dp_rank = worker.parallel_config.data_parallel_rank |
| 38 | self.dp_size = worker.parallel_config.data_parallel_size |
| 39 | self.data_parallel_master_ip = worker.parallel_config.data_parallel_master_ip |
| 40 | all2all_backend = worker.parallel_config.all2all_backend |
| 41 | if all2all_backend not in FT_BACKEND_SET: |
| 42 | raise ValueError( |
| 43 | f"Fault tolerance requires an FT-capable all2all backend " |
| 44 | f"(one of {sorted(FT_BACKEND_SET)}), but got '{all2all_backend}'." |
| 45 | ) |
| 46 | |
| 47 | def handle_command(self, ft_request: FaultToleranceRequest): |
| 48 | """Dispatch an FT command by instruction name.""" |
| 49 | with set_current_vllm_config(self.worker.vllm_config): |
| 50 | return run_method(self, ft_request.instruction, (ft_request,), {}) |
| 51 | |
| 52 | def retry(self, ft_request: FaultToleranceRequest): |
| 53 | torch.accelerator.synchronize() |
| 54 | params = ft_request.params |
| 55 | self._clean_worker_state() |
| 56 | if self.dp_size > 1: |
| 57 | get_ep_all2all_manager().clean_buffers() |
| 58 | old_cpu_group = get_dp_group().cpu_group |
| 59 | stateless_destroy_torch_distributed_process_group(old_cpu_group) |
| 60 | world_size = self.worker.parallel_config.world_size |
| 61 | port = params["new_stateless_dp_group_ports"][self.worker.rank % world_size] |
| 62 | get_dp_group().cpu_group = stateless_init_torch_distributed_process_group( |
| 63 | self.data_parallel_master_ip, |
| 64 | port, |
| 65 | self.dp_rank, |
| 66 | self.dp_size, |
| 67 | backend="gloo", |
| 68 | timeout=self.worker.parallel_config.cpu_distributed_timeout, |
| 69 | ) |
| 70 | |
| 71 | def _clean_worker_state(self): |
| 72 | model_runner = self.worker.model_runner |
| 73 | model_runner.execute_model_state = None |
| 74 | if self.worker.use_v2_model_runner: |
| 75 | runner = cast("GPUModelRunnerV2", model_runner) |
| 76 | for req_id in list(runner.req_states.req_id_to_index): |
| 77 | runner._remove_request(req_id) |
| 78 | else: |
| 79 | model_runner.kv_connector_output = None |
| 80 | |
| 81 | input_batch = model_runner.input_batch |
| 82 | cached_req_ids = list(input_batch.req_id_to_index) |
| 83 | for req_id in cached_req_ids: |
| 84 | model_runner.requests.pop(req_id, None) |
| 85 | model_runner.num_prompt_logprobs.pop(req_id, None) |
| 86 | input_batch.remove_request(req_id) |
| 87 | |
| 88 | input_batch.condense() |
| 89 | input_batch.refresh_metadata() |
| 90 | input_batch.req_prompt_embeds.clear() |