| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | |
| 5 | from typing import Annotated |
| 6 | |
| 7 | from fastapi import APIRouter, FastAPI, Query, Request |
| 8 | from fastapi.responses import JSONResponse |
| 9 | |
| 10 | from vllm.engine.protocol import EngineClient |
| 11 | from vllm.logger import init_logger |
| 12 | from vllm.v1.engine import PauseMode |
| 13 | |
| 14 | from .metrics import sleep_mode_operation_metrics |
| 15 | |
| 16 | logger = init_logger(__name__) |
| 17 | |
| 18 | |
| 19 | def engine_client(request: Request) -> EngineClient: |
| 20 | return request.app.state.engine_client |
| 21 | |
| 22 | |
| 23 | router = APIRouter() |
| 24 | |
| 25 | |
| 26 | @router.post("/sleep") |
| 27 | async def sleep( |
| 28 | raw_request: Request, |
| 29 | level: Annotated[int, Query(ge=0, le=2)] = 1, |
| 30 | mode: Annotated[PauseMode, Query()] = "abort", |
| 31 | ) -> JSONResponse: |
| 32 | with sleep_mode_operation_metrics().record("sleep"): |
| 33 | await engine_client(raw_request).sleep(level, mode) |
| 34 | return JSONResponse(content={"status": "sleeping", "level": level}) |
| 35 | |
| 36 | |
| 37 | @router.post("/release_kv_cache_memory") |
| 38 | async def release_kv_cache_memory(raw_request: Request) -> JSONResponse: |
| 39 | with sleep_mode_operation_metrics().record("release_kv_cache_memory"): |
| 40 | await engine_client(raw_request).release_kv_cache_memory() |
| 41 | return JSONResponse(content={"status": "kv_cache_released"}) |
| 42 | |
| 43 | |
| 44 | @router.post("/wake_up") |
| 45 | async def wake_up(raw_request: Request) -> JSONResponse: |
| 46 | tags = raw_request.query_params.getlist("tags") |
| 47 | if tags == []: |
| 48 | # set to None to wake up all tags if no tags are provided |
| 49 | tags = None |
| 50 | logger.info("wake up the engine with tags: %s", tags) |
| 51 | with sleep_mode_operation_metrics().record("wake"): |
| 52 | fully_awake = await engine_client(raw_request).wake_up(tags) |
| 53 | return JSONResponse( |
| 54 | content={"status": "awake" if fully_awake else "sleeping", "tags": tags} |
| 55 | ) |
| 56 | |
| 57 | |
| 58 | @router.get("/is_sleeping") |
| 59 | async def is_sleeping(raw_request: Request): |
| 60 | is_sleeping = await engine_client(raw_request).is_sleeping() |
| 61 | return JSONResponse(content={"is_sleeping": is_sleeping}) |
| 62 | |
| 63 | |
| 64 | def attach_router(app: FastAPI): |
| 65 | app.include_router(router) |