| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import json |
| 5 | from http import HTTPStatus |
| 6 | from typing import Annotated, Any, Literal |
| 7 | |
| 8 | from fastapi import APIRouter, Body, FastAPI, HTTPException, Query, Request |
| 9 | from fastapi.responses import JSONResponse |
| 10 | |
| 11 | from vllm.distributed.weight_transfer.base import ( |
| 12 | WeightTransferInitRequest, |
| 13 | WeightTransferUpdatePayload, |
| 14 | WeightTransferUpdateRequest, |
| 15 | ) |
| 16 | from vllm.engine.protocol import EngineClient |
| 17 | from vllm.entrypoints.serve.dev.rlhf.metrics import weight_operation_metrics |
| 18 | from vllm.logger import init_logger |
| 19 | from vllm.v1.engine import PauseMode |
| 20 | |
| 21 | logger = init_logger(__name__) |
| 22 | |
| 23 | |
| 24 | def engine_client(request: Request) -> EngineClient: |
| 25 | return request.app.state.engine_client |
| 26 | |
| 27 | |
| 28 | router = APIRouter() |
| 29 | |
| 30 | |
| 31 | @router.post("/pause") |
| 32 | async def pause_generation( |
| 33 | raw_request: Request, |
| 34 | mode: Annotated[PauseMode, Query()] = "abort", |
| 35 | wait_for_inflight_requests: bool = Query(False), |
| 36 | clear_cache: Annotated[bool, Query()] = True, |
| 37 | ) -> JSONResponse: |
| 38 | """Pause generation requests to allow weight updates. |
| 39 | |
| 40 | Args: |
| 41 | raw_request: The incoming FastAPI request, used to reach the engine |
| 42 | client on the app state. |
| 43 | mode: How to handle in-flight requests: |
| 44 | - ``"abort"``: Abort all in-flight requests immediately (default). |
| 45 | - ``"wait"``: Wait for in-flight requests to complete. |
| 46 | - ``"keep"``: Freeze requests in queue; they resume on /resume. |
| 47 | wait_for_inflight_requests: DEPRECATED. Use ``mode="wait"`` instead. |
| 48 | clear_cache: DEPRECATED. Whether to clear KV/prefix caches after |
| 49 | draining. Ignored when mode="keep". |
| 50 | |
| 51 | """ |
| 52 | engine = engine_client(raw_request) |
| 53 | |
| 54 | try: |
| 55 | await engine.pause_generation( |
| 56 | mode=mode, |
| 57 | clear_cache=clear_cache, |
| 58 | wait_for_inflight_requests=wait_for_inflight_requests, |
| 59 | ) |
| 60 | return JSONResponse( |
| 61 | content={"status": "paused"}, |
| 62 | status_code=HTTPStatus.OK.value, |
| 63 | ) |
| 64 | |
| 65 | except ValueError as err: |
| 66 | return JSONResponse( |
| 67 | content={"error": str(err)}, |
| 68 | status_code=HTTPStatus.BAD_REQUEST.value, |
| 69 | ) |
| 70 | except Exception as err: # pragma: no cover - defensive |
| 71 | logger.exception("Failed to pause generation") |
| 72 | return JSONResponse( |
| 73 | content={"error": f"Failed to pause generation: {err}"}, |
| 74 | status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, |
| 75 | ) |
| 76 | |
| 77 | |
| 78 | @router.post("/resume") |
| 79 | async def resume_generation(raw_request: Request) -> JSONResponse: |
| 80 | """Resume generation after a pause.""" |
| 81 | engine = engine_client(raw_request) |
| 82 | |
| 83 | try: |
| 84 | await engine.resume_generation() |
| 85 | return JSONResponse( |
| 86 | content={"status": "resumed"}, |
| 87 | status_code=HTTPStatus.OK.value, |
| 88 | ) |
| 89 | except Exception as err: # pragma: no cover - defensive |
| 90 | logger.exception("Failed to resume generation") |
| 91 | return JSONResponse( |
| 92 | content={"error": f"Failed to resume generation: {err}"}, |
| 93 | status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, |
| 94 | ) |
| 95 | |
| 96 | |
| 97 | @router.post("/abort_requests") |
| 98 | async def abort_requests(raw_request: Request) -> JSONResponse: |
| 99 | """Abort in-flight requests without pausing the scheduler. |
| 100 | |
| 101 | Empty/missing ``request_ids`` aborts all in-flight requests. |
| 102 | """ |
| 103 | engine = engine_client(raw_request) |
| 104 | |
| 105 | try: |
| 106 | body = await raw_request.json() |
| 107 | except json.JSONDecodeError as e: |
| 108 | raise HTTPException(status_code=400, detail="Invalid JSON format") from e # noqa: B904 |
| 109 | if not isinstance(body, dict): |
| 110 | raise HTTPException( |
| 111 | status_code=400, detail="Request body must be a JSON object" |
| 112 | ) |
| 113 | |
| 114 | request_ids = body.get("request_ids") |
| 115 | |
| 116 | try: |
| 117 | if request_ids: |
| 118 | # Body ids are external (user-supplied) request ids. |
| 119 | await engine.abort(request_ids) |
| 120 | else: |
| 121 | # The dev RL server runs AsyncLLM; abort everything it is tracking. |
| 122 | # request_states is keyed by internal ids; parent_requests holds |
| 123 | # parallel-sampling parents. Abort both as internal ids. |
| 124 | from vllm.v1.engine.async_llm import AsyncLLM |
| 125 | |
| 126 | assert isinstance(engine, AsyncLLM) |
| 127 | op = engine.output_processor |
| 128 | request_ids = [ |
| 129 | *op.request_states.keys(), |
| 130 | *op.parent_requests.keys(), |
| 131 | ] |
| 132 | await engine.abort(request_ids, internal=True) |
| 133 | return JSONResponse( |
| 134 | content={"status": "aborted", "aborted": len(request_ids)}, |
| 135 | status_code=HTTPStatus.OK.value, |
| 136 | ) |
| 137 | except Exception as err: # pragma: no cover - defensive |
| 138 | logger.exception("Failed to abort requests") |
| 139 | return JSONResponse( |
| 140 | content={"error": f"Failed to abort requests: {err}"}, |
| 141 | status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, |
| 142 | ) |
| 143 | |
| 144 | |
| 145 | @router.get("/is_paused") |
| 146 | async def is_paused(raw_request: Request) -> JSONResponse: |
| 147 | """Return the current pause status.""" |
| 148 | engine = engine_client(raw_request) |
| 149 | |
| 150 | try: |
| 151 | paused = await engine.is_paused() |
| 152 | except Exception as err: # pragma: no cover - defensive |
| 153 | logger.exception("Failed to fetch pause status") |
| 154 | return JSONResponse( |
| 155 | content={"error": f"Failed to fetch pause status: {err}"}, |
| 156 | status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, |
| 157 | ) |
| 158 | |
| 159 | return JSONResponse(content={"is_paused": paused}) |
| 160 | |
| 161 | |
| 162 | @router.post("/init_weight_transfer_engine") |
| 163 | async def init_weight_transfer_engine( |
| 164 | raw_request: Request, |
| 165 | init_info: Annotated[dict[str, Any], Body(embed=True)], |
| 166 | ): |
| 167 | with weight_operation_metrics().record("init"): |
| 168 | await engine_client(raw_request).init_weight_transfer_engine( |
| 169 | WeightTransferInitRequest(init_info=init_info) |
| 170 | ) |
| 171 | return JSONResponse(content={"message": "Weight transfer initialized"}) |
| 172 | |
| 173 | |
| 174 | @router.post("/start_weight_update") |
| 175 | async def start_weight_update(raw_request: Request): |
| 176 | with weight_operation_metrics().record("start"): |
| 177 | await engine_client(raw_request).start_weight_update() |
| 178 | return JSONResponse(content={"message": "Weight update started"}) |
| 179 | |
| 180 | |
| 181 | @router.post("/start_draft_weight_update") |
| 182 | async def start_draft_weight_update(raw_request: Request): |
| 183 | with weight_operation_metrics().record("start_draft"): |
| 184 | await engine_client(raw_request).start_draft_weight_update() |
| 185 | return JSONResponse(content={"message": "Draft weight update started"}) |
| 186 | |
| 187 | |
| 188 | @router.post("/update_weights") |
| 189 | async def update_weights( |
| 190 | raw_request: Request, |
| 191 | update_info: Annotated[WeightTransferUpdatePayload, Body(embed=True)], |
| 192 | ): |
| 193 | with weight_operation_metrics().record("update"): |
| 194 | await engine_client(raw_request).update_weights( |
| 195 | request=WeightTransferUpdateRequest(update_info=update_info) |
| 196 | ) |
| 197 | return JSONResponse(content={"message": "Weights updated"}) |
| 198 | |
| 199 | |
| 200 | @router.post("/finish_weight_update") |
| 201 | async def finish_weight_update( |
| 202 | raw_request: Request, |
| 203 | weight_version: Annotated[str | None, Body(embed=True)] = None, |
| 204 | ): |
| 205 | with weight_operation_metrics().record("finish"): |
| 206 | await engine_client(raw_request).finish_weight_update(weight_version) |
| 207 | return JSONResponse(content={"message": "Weight update finished"}) |
| 208 | |
| 209 | |
| 210 | @router.post("/update_weight_version") |
| 211 | async def update_weight_version( |
| 212 | raw_request: Request, |
| 213 | new_version: Annotated[str, Body(embed=True)], |
| 214 | ): |
| 215 | await engine_client(raw_request).update_weight_version(new_version) |
| 216 | return JSONResponse(content={"success": True, "new_version": new_version}) |
| 217 | |
| 218 | |
| 219 | @router.get("/weight_info") |
| 220 | async def weight_info(raw_request: Request): |
| 221 | weight_version = await engine_client(raw_request).get_weight_version() |
| 222 | return JSONResponse(content={"weight_version": weight_version}) |
| 223 | |
| 224 | |
| 225 | @router.post("/weight_checker") |
| 226 | async def weight_checker( |
| 227 | raw_request: Request, |
| 228 | action: Annotated[Literal["checksum", "reset", "compare"], Body(embed=True)], |
| 229 | baseline: Annotated[dict[str, str] | None, Body(embed=True)] = None, |
| 230 | ) -> JSONResponse: |
| 231 | """Checksum, reset, or compare model weights against a baseline.""" |
| 232 | client = engine_client(raw_request) |
| 233 | if action == "reset": |
| 234 | await client.collective_rpc("reset_weights") |
| 235 | return JSONResponse(content={"status": "reset"}) |
| 236 | if action == "compare" and baseline is None: |
| 237 | raise HTTPException( |
| 238 | status_code=HTTPStatus.BAD_REQUEST.value, |
| 239 | detail="action='compare' requires a 'baseline' object", |
| 240 | ) |
| 241 | |
| 242 | checksums: dict[str, str] = {} |
| 243 | for worker_checksums in await client.compute_weight_checksums(): |
| 244 | if duplicates := checksums.keys() & worker_checksums.keys(): |
| 245 | raise HTTPException( |
| 246 | status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, |
| 247 | detail=f"Duplicate weight checksum keys: {sorted(duplicates)}", |
| 248 | ) |
| 249 | checksums.update(worker_checksums) |
| 250 | if action == "checksum": |
| 251 | return JSONResponse(content={"checksums": checksums}) |
| 252 | |
| 253 | assert baseline is not None |
| 254 | mismatches = sorted( |
| 255 | key |
| 256 | for key in checksums.keys() | baseline.keys() |
| 257 | if checksums.get(key) != baseline.get(key) |
| 258 | ) |
| 259 | return JSONResponse(content={"match": not mismatches, "mismatches": mismatches}) |
| 260 | |
| 261 | |
| 262 | @router.get("/get_world_size") |
| 263 | async def get_world_size( |
| 264 | raw_request: Request, |
| 265 | include_dp: bool = Query(True), |
| 266 | ): |
| 267 | """Get the world size from the parallel config. |
| 268 | |
| 269 | Args: |
| 270 | raw_request: The incoming FastAPI request, used to reach the engine |
| 271 | client on the app state. |
| 272 | include_dp: If True (default), returns the world size including |
| 273 | data parallelism (TP * PP * DP). If False, returns the world |
| 274 | size without data parallelism (TP * PP). |
| 275 | |
| 276 | """ |
| 277 | parallel_config = engine_client(raw_request).vllm_config.parallel_config |
| 278 | if include_dp: |
| 279 | world_size = parallel_config.world_size_across_dp |
| 280 | else: |
| 281 | world_size = parallel_config.world_size |
| 282 | return JSONResponse(content={"world_size": world_size}) |
| 283 | |
| 284 | |
| 285 | def attach_router(app: FastAPI): |
| 286 | app.include_router(router) |