| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import json |
| 4 | import uuid |
| 5 | from http import HTTPStatus |
| 6 | |
| 7 | from fastapi import APIRouter, BackgroundTasks, Depends, FastAPI, HTTPException, Request |
| 8 | from fastapi.responses import JSONResponse |
| 9 | |
| 10 | from vllm.engine.protocol import EngineClient |
| 11 | from vllm.entrypoints.serve.engine.protocol import ErrorResponse |
| 12 | from vllm.entrypoints.serve.utils.api_utils import validate_json_request |
| 13 | from vllm.logger import init_logger |
| 14 | from vllm.v1.fault_tolerance.utils import FaultToleranceRequest |
| 15 | |
| 16 | logger = init_logger(__name__) |
| 17 | |
| 18 | router = APIRouter() |
| 19 | |
| 20 | _ALLOWED_INSTRUCTIONS = {"retry"} |
| 21 | |
| 22 | |
| 23 | def _validate_payload(body: dict) -> tuple[str, dict]: |
| 24 | if not isinstance(body, dict): |
| 25 | raise HTTPException(400, "Request body must be a JSON object.") |
| 26 | instruction = body.get("instruction") |
| 27 | if not instruction: |
| 28 | raise HTTPException(400, "'instruction' is required.") |
| 29 | if instruction not in _ALLOWED_INSTRUCTIONS: |
| 30 | raise HTTPException(400, f"Invalid instruction: '{instruction}'.") |
| 31 | params = body.get("params", {}) |
| 32 | if not isinstance(params, dict): |
| 33 | raise HTTPException(400, "'params' must be an object.") |
| 34 | return instruction, params |
| 35 | |
| 36 | |
| 37 | @router.post( |
| 38 | "/fault_tolerance/apply", |
| 39 | dependencies=[Depends(validate_json_request)], |
| 40 | responses={ |
| 41 | HTTPStatus.ACCEPTED.value: {"model": dict}, |
| 42 | HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse}, |
| 43 | }, |
| 44 | ) |
| 45 | async def process_fault_tolerance_instruction( |
| 46 | raw_request: Request, background_tasks: BackgroundTasks |
| 47 | ): |
| 48 | try: |
| 49 | body = await raw_request.json() |
| 50 | except json.JSONDecodeError as e: |
| 51 | raise HTTPException(400, "Invalid JSON format") from e |
| 52 | |
| 53 | instruction, params = _validate_payload(body) |
| 54 | ft_request = FaultToleranceRequest( |
| 55 | instruction=instruction, |
| 56 | params=params, |
| 57 | request_id=str(uuid.uuid4()), |
| 58 | ) |
| 59 | |
| 60 | client: EngineClient = raw_request.app.state.engine_client |
| 61 | # Recovery runs cross-rank collective ops that only complete once every rank |
| 62 | # has been dispatched. Run it in the background and return immediately so the |
| 63 | # orchestrator can dispatch to all ranks without blocking; completion is |
| 64 | # observed by polling GET /fault_tolerance/status. |
| 65 | background_tasks.add_task(_run_fault_recovery, client, ft_request) |
| 66 | return JSONResponse( |
| 67 | status_code=HTTPStatus.ACCEPTED.value, |
| 68 | content={ |
| 69 | "message": "Request accepted; poll /fault_tolerance/status for updates.", |
| 70 | "request_id": ft_request.request_id, |
| 71 | }, |
| 72 | background=background_tasks, |
| 73 | ) |
| 74 | |
| 75 | |
| 76 | async def _run_fault_recovery( |
| 77 | client: EngineClient, ft_request: FaultToleranceRequest |
| 78 | ) -> None: |
| 79 | """Drive recovery to completion after the 202 response is sent.""" |
| 80 | try: |
| 81 | result = await client.handle_fault(ft_request) |
| 82 | except Exception: |
| 83 | logger.exception("[FT] Recovery dispatch failed.") |
| 84 | return |
| 85 | if not result.success: |
| 86 | logger.error( |
| 87 | "[FT] Recovery failed for request %s: %s", |
| 88 | ft_request.request_id, |
| 89 | result.reason, |
| 90 | ) |
| 91 | |
| 92 | |
| 93 | @router.get("/fault_tolerance/status") |
| 94 | async def get_status(raw_request: Request): |
| 95 | client: EngineClient = raw_request.app.state.engine_client |
| 96 | return JSONResponse(content=await client.get_status()) |
| 97 | |
| 98 | |
| 99 | def register_fault_tolerance_api_router(app: FastAPI): |
| 100 | app.include_router(router) |