| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | |
| 5 | import json |
| 6 | from http import HTTPStatus |
| 7 | |
| 8 | from fastapi import APIRouter, Depends, FastAPI, HTTPException, Request |
| 9 | from fastapi.responses import JSONResponse |
| 10 | |
| 11 | from vllm.engine.protocol import EngineClient |
| 12 | from vllm.entrypoints.serve.elastic_ep.middleware import get_scaling_elastic_ep |
| 13 | from vllm.entrypoints.serve.engine.protocol import ErrorResponse |
| 14 | from vllm.entrypoints.serve.utils.api_utils import validate_json_request |
| 15 | from vllm.logger import init_logger |
| 16 | |
| 17 | logger = init_logger(__name__) |
| 18 | |
| 19 | |
| 20 | def engine_client(request: Request) -> EngineClient: |
| 21 | return request.app.state.engine_client |
| 22 | |
| 23 | |
| 24 | router = APIRouter() |
| 25 | |
| 26 | |
| 27 | @router.post( |
| 28 | "/scale_elastic_ep", |
| 29 | dependencies=[Depends(validate_json_request)], |
| 30 | responses={ |
| 31 | HTTPStatus.OK.value: {"model": dict}, |
| 32 | HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse}, |
| 33 | HTTPStatus.REQUEST_TIMEOUT.value: {"model": ErrorResponse}, |
| 34 | HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse}, |
| 35 | }, |
| 36 | ) |
| 37 | async def scale_elastic_ep(raw_request: Request): |
| 38 | try: |
| 39 | body = await raw_request.json() |
| 40 | except json.JSONDecodeError as e: |
| 41 | raise HTTPException(status_code=400, detail="Invalid JSON format") from e |
| 42 | |
| 43 | new_data_parallel_size = body.get("new_data_parallel_size") |
| 44 | drain_timeout = body.get("drain_timeout", 120) # Default 2 minutes |
| 45 | |
| 46 | if new_data_parallel_size is None: |
| 47 | raise HTTPException( |
| 48 | status_code=400, detail="new_data_parallel_size is required" |
| 49 | ) |
| 50 | |
| 51 | if not isinstance(new_data_parallel_size, int) or new_data_parallel_size <= 0: |
| 52 | raise HTTPException( |
| 53 | status_code=400, |
| 54 | detail="new_data_parallel_size must be a positive integer", |
| 55 | ) |
| 56 | |
| 57 | if not isinstance(drain_timeout, int) or drain_timeout <= 0: |
| 58 | raise HTTPException( |
| 59 | status_code=400, detail="drain_timeout must be a positive integer" |
| 60 | ) |
| 61 | |
| 62 | client = engine_client(raw_request) |
| 63 | try: |
| 64 | await client.scale_elastic_ep(new_data_parallel_size, drain_timeout) |
| 65 | return JSONResponse( |
| 66 | { |
| 67 | "message": f"Scaled to {new_data_parallel_size} data parallel engines", |
| 68 | } |
| 69 | ) |
| 70 | except TimeoutError as e: |
| 71 | raise HTTPException( |
| 72 | status_code=408, |
| 73 | detail="Scale failed due to request drain timeout " |
| 74 | f"after {drain_timeout} seconds", |
| 75 | ) from e |
| 76 | except ValueError as e: |
| 77 | raise HTTPException(status_code=400, detail=str(e)) from e |
| 78 | except Exception as e: |
| 79 | logger.error("Scale failed: %s", e) |
| 80 | raise HTTPException(status_code=500, detail="Scale failed") from e |
| 81 | |
| 82 | |
| 83 | @router.post("/is_scaling_elastic_ep") |
| 84 | async def is_scaling_elastic_ep(raw_request: Request): |
| 85 | return JSONResponse({"is_scaling_elastic_ep": get_scaling_elastic_ep()}) |
| 86 | |
| 87 | |
| 88 | def attach_router(app: FastAPI): |
| 89 | app.include_router(router) |