main
py 68 lines 2.07 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
5 from fastapi import APIRouter, FastAPI, Query, Request
6 from fastapi.responses import JSONResponse, Response
7
8 from vllm.engine.protocol import EngineClient
9 from vllm.logger import init_logger
10
11 logger = init_logger(__name__)
12
13 router = APIRouter()
14
15
16 def engine_client(request: Request) -> EngineClient:
17 return request.app.state.engine_client
18
19
20 @router.post("/reset_prefix_cache")
21 async def reset_prefix_cache(
22 raw_request: Request,
23 reset_running_requests: bool = Query(default=False),
24 reset_external: bool = Query(default=False),
25 ):
26 """Reset the local prefix cache.
27
28 Optionally, if the query parameter `reset_external=true`
29 also resets the external (connector-managed) prefix cache.
30
31 Returns `{"success": bool}`. The reset fails (`success=false`) while
32 blocks are still held, e.g. by running requests or in-flight async KV
33 offload transfers; callers may retry.
34
35 Example:
36 POST /reset_prefix_cache?reset_external=true
37
38 """
39 logger.info("Resetting prefix cache...")
40
41 success = await engine_client(raw_request).reset_prefix_cache(
42 reset_running_requests, reset_external
43 )
44 return JSONResponse(content={"success": bool(success)})
45
46
47 @router.post("/reset_mm_cache")
48 async def reset_mm_cache(raw_request: Request):
49 """Reset the multi-modal cache. Note that we currently do not check if the
50 multi-modal cache is successfully reset in the API server.
51 """
52 logger.info("Resetting multi-modal cache...")
53 await engine_client(raw_request).reset_mm_cache()
54 return Response(status_code=200)
55
56
57 @router.post("/reset_encoder_cache")
58 async def reset_encoder_cache(raw_request: Request):
59 """Reset the encoder cache. Note that we currently do not check if the
60 encoder cache is successfully reset in the API server.
61 """
62 logger.info("Resetting encoder cache...")
63 await engine_client(raw_request).reset_encoder_cache()
64 return Response(status_code=200)
65
66
67 def attach_router(app: FastAPI):
68 app.include_router(router)