[Misc] Add argument to disable FastAPI docs

vllm/engine/arg_utils.py

@@ -784,6 +784,11 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
             },
             default=None,
             help="override or set neuron device configuration.")
+        parser.add_argument(
+            "--disable-fastapi-docs",
+            action='store_true',
+            default=False,
+            help="Disable OpenAPI schema, Swagger UI, and ReDoc documentation")
 
         return parser
 

vllm/entrypoints/api_server.py

@@ -11,7 +11,7 @@
 from argparse import Namespace
 from typing import Any, AsyncGenerator, Optional
 
-from fastapi import FastAPI, Request
+from fastapi import APIRouter, FastAPI, Request
 from fastapi.responses import JSONResponse, Response, StreamingResponse
 
 from vllm.engine.arg_utils import AsyncEngineArgs
@@ -27,17 +27,17 @@
 logger = init_logger("vllm.entrypoints.api_server")
 
 TIMEOUT_KEEP_ALIVE = 5  # seconds.
-app = FastAPI()
+router = APIRouter()
 engine = None
 
 
-@app.get("/health")
+@router.get("/health")
 async def health() -> Response:
     """Health check."""
     return Response(status_code=200)
 
 
-@app.post("/generate")
+@router.post("/generate")
 async def generate(request: Request) -> Response:
     """Generate completion for the request.
 
@@ -88,8 +88,11 @@ async def stream_results() -> AsyncGenerator[bytes, None]:
 
 
 def build_app(args: Namespace) -> FastAPI:
-    global app
-
+    if args.disable_fastapi_docs:
+        app = FastAPI(openapi_url=None, docs_url=None, redoc_url=None)
+    else:
+        app = FastAPI()
+    app.include_router(router)
     app.root_path = args.root_path
     return app
 

vllm/entrypoints/openai/api_server.py

@@ -417,7 +417,13 @@ async def unload_lora_adapter(request: UnloadLoraAdapterRequest,
 
 
 def build_app(args: Namespace) -> FastAPI:
-    app = FastAPI(lifespan=lifespan)
+    if args.disable_fastapi_docs:
+        app = FastAPI(openapi_url=None,
+                      docs_url=None,
+                      redoc_url=None,
+                      lifespan=lifespan)
+    else:
+        app = FastAPI(lifespan=lifespan)
     app.include_router(router)
     app.root_path = args.root_path