@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 159 lines 5.82 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import json
4 import logging
5 from collections.abc import Awaitable, Callable
6 from http import HTTPStatus
7 from typing import Any
8
9 import pydantic
10 from fastapi import APIRouter, Depends, FastAPI, HTTPException, Request
11 from fastapi.responses import JSONResponse, Response
12
13 from vllm.config import ModelConfig
14 from vllm.entrypoints.generate.factories import get_generate_invocation_types
15 from vllm.entrypoints.pooling.factories import get_pooling_invocation_types
16 from vllm.entrypoints.serve.engine.protocol import ErrorResponse
17 from vllm.entrypoints.serve.engine.serving import BaseServing
18 from vllm.entrypoints.serve.instrumentator.basic import base
19 from vllm.entrypoints.serve.instrumentator.health import health
20 from vllm.entrypoints.serve.utils.api_utils import validate_json_request
21 from vllm.tasks import SupportedTask
22
23 # TODO: RequestType = TypeForm[BaseModel] when recognized by type checkers
24 # (requires typing_extensions >= 4.13)
25 RequestType = Any
26 GetHandlerFn = Callable[[Request], BaseServing | None]
27 EndpointFn = Callable[[RequestType, Request], Awaitable[Any]]
28
29
30 def _snapshot_handler_levels() -> list[tuple[logging.Handler, int]]:
31 """Snapshot handler levels for loggers that may be affected by third-party
32 logging configuration side effects (e.g. model_hosting_container_standards
33 calling configure_root_logger() at import time)."""
34 loggers = [logging.getLogger(), logging.getLogger("vllm")]
35 seen: set[int] = set()
36 snapshot: list[tuple[logging.Handler, int]] = []
37 for logger in loggers:
38 for handler in logger.handlers:
39 if id(handler) not in seen:
40 seen.add(id(handler))
41 snapshot.append((handler, handler.level))
42 return snapshot
43
44
45 def _restore_handler_levels(
46 snapshot: list[tuple[logging.Handler, int]],
47 ) -> None:
48 """Restore handler levels from a snapshot."""
49 for handler, level in snapshot:
50 handler.setLevel(level)
51
52
53 def attach_router(
54 app: FastAPI,
55 supported_tasks: tuple["SupportedTask", ...],
56 model_config: ModelConfig | None = None,
57 ):
58 """Attach the SageMaker hosting endpoints to the API server.
59
60 Handler levels are snapshotted and restored because importing
61 model_hosting_container_standards may reconfigure root logging.
62 """
63 snapshot = _snapshot_handler_levels()
64 try:
65 _attach_router(app, supported_tasks, model_config)
66 finally:
67 _restore_handler_levels(snapshot)
68
69
70 def _attach_router(
71 app: FastAPI,
72 supported_tasks: tuple["SupportedTask", ...],
73 model_config: ModelConfig | None = None,
74 ):
75 """Register the SageMaker hosting routes (/ping, /invocations) on the
76 app."""
77 import model_hosting_container_standards.sagemaker as sagemaker_standards
78
79 router = APIRouter()
80
81 # NOTE: Construct the TypeAdapters only once
82 INVOCATION_TYPES = get_generate_invocation_types(
83 supported_tasks, model_config
84 ) + get_pooling_invocation_types(supported_tasks, model_config)
85
86 INVOCATION_VALIDATORS = [
87 (pydantic.TypeAdapter(request_type), (get_handler, endpoint))
88 for request_type, (get_handler, endpoint) in INVOCATION_TYPES
89 ]
90
91 @router.post("/ping", response_class=Response)
92 @router.get("/ping", response_class=Response)
93 @sagemaker_standards.register_ping_handler
94 async def ping(raw_request: Request) -> Response:
95 """Ping check. Endpoint required for SageMaker."""
96 return await health(raw_request)
97
98 @router.post(
99 "/invocations",
100 dependencies=[Depends(validate_json_request)],
101 responses={
102 HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
103 HTTPStatus.UNSUPPORTED_MEDIA_TYPE.value: {"model": ErrorResponse},
104 HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
105 },
106 )
107 @sagemaker_standards.register_invocation_handler
108 @sagemaker_standards.stateful_session_manager()
109 @sagemaker_standards.inject_adapter_id(adapter_path="model")
110 async def invocations(raw_request: Request):
111 """For SageMaker, routes requests based on the request type."""
112 try:
113 body = await raw_request.json()
114 except json.JSONDecodeError as e:
115 raise HTTPException(
116 status_code=HTTPStatus.BAD_REQUEST.value,
117 detail=f"JSON decode error: {e}",
118 ) from e
119
120 valid_endpoints = [
121 (validator, endpoint)
122 for validator, (get_handler, endpoint) in INVOCATION_VALIDATORS
123 if get_handler(raw_request) is not None
124 ]
125
126 for request_validator, endpoint in valid_endpoints:
127 try:
128 request = request_validator.validate_python(body)
129 except pydantic.ValidationError:
130 continue
131
132 return await endpoint(request, raw_request)
133
134 type_names = [
135 t.__name__ if isinstance(t := validator._type, type) else str(t)
136 for validator, _ in valid_endpoints
137 ]
138 msg = f"Cannot find suitable handler for request. Expected one of: {type_names}"
139 res = base(raw_request).create_error_response(message=msg)
140 return JSONResponse(content=res.model_dump(), status_code=res.error.code)
141
142 app.include_router(router)
143
144
145 def sagemaker_standards_bootstrap(app: FastAPI) -> FastAPI:
146 """Bootstrap the app with the SageMaker hosting standards.
147
148 Handler levels are restored right after the import because importing
149 model_hosting_container_standards may reconfigure root logging, and
150 bootstrap must run with the original levels.
151 """
152 snapshot = _snapshot_handler_levels()
153 try:
154 import model_hosting_container_standards.sagemaker as sagemaker_standards
155
156 app = sagemaker_standards.bootstrap(app)
157 finally:
158 _restore_handler_levels(snapshot)
159 return app