| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | |
| 5 | from fastapi import APIRouter, FastAPI, Request |
| 6 | from fastapi.responses import Response |
| 7 | |
| 8 | from vllm.config import ProfilerConfig |
| 9 | from vllm.engine.protocol import EngineClient |
| 10 | from vllm.logger import init_logger |
| 11 | |
| 12 | logger = init_logger(__name__) |
| 13 | |
| 14 | router = APIRouter() |
| 15 | |
| 16 | |
| 17 | def engine_client(request: Request) -> EngineClient: |
| 18 | return request.app.state.engine_client |
| 19 | |
| 20 | |
| 21 | @router.post("/start_profile") |
| 22 | async def start_profile(raw_request: Request): |
| 23 | logger.info("Starting profiler...") |
| 24 | await engine_client(raw_request).start_profile() |
| 25 | logger.info("Profiler started.") |
| 26 | return Response(status_code=200) |
| 27 | |
| 28 | |
| 29 | @router.post("/stop_profile") |
| 30 | async def stop_profile(raw_request: Request): |
| 31 | logger.info("Stopping profiler...") |
| 32 | await engine_client(raw_request).stop_profile() |
| 33 | logger.info("Profiler stopped.") |
| 34 | return Response(status_code=200) |
| 35 | |
| 36 | |
| 37 | def attach_router(app: FastAPI): |
| 38 | profiler_config = getattr(app.state.args, "profiler_config", None) |
| 39 | assert profiler_config is None or isinstance(profiler_config, ProfilerConfig) |
| 40 | if profiler_config is not None and profiler_config.profiler is not None: |
| 41 | logger.warning_once( |
| 42 | "Profiler with mode '%s' is enabled in the " |
| 43 | "API server. This should ONLY be used for local development!", |
| 44 | profiler_config.profiler, |
| 45 | ) |
| 46 | app.include_router(router) |