YancyandAsthenia
73893e73ea
Signed-off-by: Asthenia <asthenia0412@gmail.com> Co-authored-by: Asthenia <asthenia0412@gmail.com>
84 lines
2.5 KiB
Python
84 lines
2.5 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Endpoint restriction policy for omni pipelines."""
|
|
|
|
from dataclasses import dataclass
|
|
from enum import Enum
|
|
from typing import NamedTuple
|
|
|
|
from fastapi import FastAPI, Request
|
|
from fastapi.responses import JSONResponse
|
|
|
|
# vLLM < 0.28 keeps create_error_response under serve.utils (it is not
|
|
# re-exported from vllm.entrypoints.serve); 0.28+ moved it under
|
|
# serve.exception_handling and re-exports it from the package root.
|
|
try:
|
|
from vllm.entrypoints.serve import create_error_response
|
|
except ImportError:
|
|
from vllm.entrypoints.serve.utils.error_response import create_error_response
|
|
|
|
|
|
class RouteTarget(NamedTuple):
|
|
"""A server path & supported methods."""
|
|
|
|
path: str
|
|
methods: frozenset[str]
|
|
|
|
|
|
class OmniServingCapability(Enum):
|
|
"""Serving capabilities that pipelines can shut down."""
|
|
|
|
COMPLETIONS = RouteTarget("/v1/completions", frozenset({"POST"}))
|
|
|
|
@property
|
|
def path(self) -> str:
|
|
return self.value.path
|
|
|
|
@property
|
|
def methods(self) -> frozenset[str]:
|
|
return self.value.methods
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class EndpointRestriction:
|
|
capability: OmniServingCapability
|
|
reason: str
|
|
|
|
|
|
def build_rejection_handler(reason: str):
|
|
"""Build a rejection handler for a given endpoint for the provided reason."""
|
|
|
|
async def rejection_handler(raw_request: Request):
|
|
error = create_error_response(message=reason)
|
|
return JSONResponse(
|
|
content=error.model_dump(),
|
|
status_code=error.error.code,
|
|
)
|
|
|
|
return rejection_handler
|
|
|
|
|
|
def shutdown_unsupported_routes(
|
|
app: FastAPI,
|
|
endpoint_restrictions: tuple[EndpointRestriction, ...],
|
|
):
|
|
"""Given an initialized FastAPI server instance and a set of model specific endpoint
|
|
restrictions, remove the restricted routes and patch a handler that returns 400.
|
|
"""
|
|
from vllm_omni.entrypoints.openai.api_server import _remove_route_from_app
|
|
|
|
for end_restrict in endpoint_restrictions:
|
|
capability = end_restrict.capability
|
|
# Remove the route from the app
|
|
_remove_route_from_app(app, capability.path, capability.methods)
|
|
|
|
# Patch the bad request error with the model specific
|
|
# reason for shutting down this endpoint
|
|
rejection_handler = build_rejection_handler(end_restrict.reason)
|
|
|
|
app.add_api_route(
|
|
capability.path,
|
|
rejection_handler,
|
|
methods=list(capability.methods),
|
|
)
|