wehub-resource-sync
132 lines
4.6 KiB
Python
132 lines
4.6 KiB
Python
"""Python entrypoint of serve."""
|
|
|
|
from typing import Any, List, Literal, Optional, Tuple, Union # noqa: UP035
|
|
|
|
import fastapi
|
|
import uvicorn
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
|
|
from mlc_llm.protocol import error_protocol
|
|
from mlc_llm.serve import engine
|
|
from mlc_llm.serve.embedding_engine import AsyncEmbeddingEngine
|
|
from mlc_llm.serve.entrypoints import (
|
|
debug_entrypoints,
|
|
metrics_entrypoints,
|
|
microserving_entrypoints,
|
|
openai_entrypoints,
|
|
)
|
|
from mlc_llm.serve.server import ServerContext
|
|
from mlc_llm.support import logging
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def serve(
|
|
model: str,
|
|
device: str,
|
|
model_lib: Optional[str],
|
|
mode: Literal["local", "interactive", "server"],
|
|
enable_debug: bool,
|
|
additional_models: List[Union[str, Tuple[str, str]]], # noqa: UP006
|
|
embedding_model: Optional[str],
|
|
embedding_model_lib: Optional[str],
|
|
tensor_parallel_shards: Optional[int],
|
|
pipeline_parallel_stages: Optional[int],
|
|
opt: Optional[str],
|
|
max_num_sequence: Optional[int],
|
|
max_total_sequence_length: Optional[int],
|
|
max_single_sequence_length: Optional[int],
|
|
prefill_chunk_size: Optional[int],
|
|
sliding_window_size: Optional[int],
|
|
attention_sink_size: Optional[int],
|
|
max_history_size: Optional[int],
|
|
gpu_memory_utilization: Optional[float],
|
|
speculative_mode: Literal["disable", "small_draft", "eagle", "medusa"],
|
|
spec_draft_length: Optional[int],
|
|
spec_tree_width: Optional[int],
|
|
prefix_cache_mode: Literal["disable", "radix"],
|
|
prefix_cache_max_num_recycling_seqs: Optional[int],
|
|
prefill_mode: Literal["hybrid", "chunked"],
|
|
enable_tracing: bool,
|
|
host: str,
|
|
port: int,
|
|
allow_credentials: bool,
|
|
allow_origins: Any,
|
|
allow_methods: Any,
|
|
allow_headers: Any,
|
|
api_key: Optional[str] = None,
|
|
):
|
|
"""Serve the model with the specified configuration."""
|
|
# Create engine and start the background loop
|
|
async_engine = engine.AsyncMLCEngine(
|
|
model=model,
|
|
device=device,
|
|
model_lib=model_lib,
|
|
mode=mode,
|
|
engine_config=engine.EngineConfig(
|
|
additional_models=additional_models,
|
|
tensor_parallel_shards=tensor_parallel_shards,
|
|
pipeline_parallel_stages=pipeline_parallel_stages,
|
|
opt=opt,
|
|
max_num_sequence=max_num_sequence,
|
|
max_total_sequence_length=max_total_sequence_length,
|
|
max_single_sequence_length=max_single_sequence_length,
|
|
prefill_chunk_size=prefill_chunk_size,
|
|
sliding_window_size=sliding_window_size,
|
|
attention_sink_size=attention_sink_size,
|
|
max_history_size=max_history_size,
|
|
gpu_memory_utilization=gpu_memory_utilization,
|
|
speculative_mode=speculative_mode,
|
|
spec_draft_length=spec_draft_length,
|
|
spec_tree_width=spec_tree_width,
|
|
prefix_cache_mode=prefix_cache_mode,
|
|
prefix_cache_max_num_recycling_seqs=prefix_cache_max_num_recycling_seqs,
|
|
prefill_mode=prefill_mode,
|
|
),
|
|
enable_tracing=enable_tracing,
|
|
)
|
|
|
|
# Set up embedding model if specified
|
|
emb_engine = None
|
|
if embedding_model is not None:
|
|
if embedding_model_lib is None:
|
|
raise ValueError(
|
|
"--embedding-model-lib is required when --embedding-model is specified."
|
|
)
|
|
emb_engine = AsyncEmbeddingEngine(
|
|
model=embedding_model,
|
|
model_lib=embedding_model_lib,
|
|
device=device,
|
|
)
|
|
logger.info("Embedding model %s loaded successfully.", embedding_model)
|
|
|
|
with ServerContext() as server_context:
|
|
server_context.add_model(model, async_engine)
|
|
if emb_engine is not None:
|
|
server_context.add_embedding_engine(embedding_model, emb_engine)
|
|
server_context.api_key = api_key
|
|
|
|
app = fastapi.FastAPI()
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_credentials=allow_credentials,
|
|
allow_origins=allow_origins,
|
|
allow_methods=allow_methods,
|
|
allow_headers=allow_headers,
|
|
)
|
|
|
|
app.include_router(openai_entrypoints.app)
|
|
app.include_router(metrics_entrypoints.app)
|
|
app.include_router(microserving_entrypoints.app)
|
|
|
|
server_context.enable_debug = enable_debug
|
|
|
|
if enable_debug:
|
|
app.include_router(debug_entrypoints.app)
|
|
logger.info("Enable debug endpoint and debug_config in requests...")
|
|
|
|
app.exception_handler(error_protocol.BadRequestError)(
|
|
error_protocol.bad_request_error_handler
|
|
)
|
|
uvicorn.run(app, host=host, port=port, log_level="info")
|