wehub-resource-sync
265 lines
11 KiB
Python
265 lines
11 KiB
Python
"""Command line entrypoint of serve."""
|
|
|
|
import dataclasses
|
|
import json
|
|
from io import StringIO
|
|
from typing import Literal, Optional
|
|
|
|
from mlc_llm.interface.help import HELP
|
|
from mlc_llm.interface.serve import serve
|
|
from mlc_llm.support import argparse
|
|
from mlc_llm.support.argparse import ArgumentParser
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class EngineConfigOverride:
|
|
"""Arguments for overriding engine config."""
|
|
|
|
# Overrides for EngineConfig (runtime)
|
|
max_num_sequence: Optional[int] = None
|
|
max_total_seq_length: Optional[int] = None
|
|
prefill_chunk_size: Optional[int] = None
|
|
max_history_size: Optional[int] = None
|
|
gpu_memory_utilization: Optional[float] = None
|
|
spec_draft_length: Optional[int] = None
|
|
spec_tree_width: Optional[int] = None
|
|
prefix_cache_mode: Optional[Literal["disable", "radix"]] = None
|
|
prefix_cache_max_num_recycling_seqs: Optional[int] = None
|
|
prefill_mode: Optional[Literal["chunked", "hybrid"]] = None
|
|
context_window_size: Optional[int] = None
|
|
sliding_window_size: Optional[int] = None
|
|
attention_sink_size: Optional[int] = None
|
|
tensor_parallel_shards: Optional[int] = None
|
|
pipeline_parallel_stages: Optional[int] = None
|
|
opt: Optional[str] = None
|
|
|
|
def __repr__(self) -> str:
|
|
out = StringIO()
|
|
print(f"max_num_sequence={self.max_num_sequence}", file=out, end="")
|
|
print(f";max_total_seq_length={self.max_total_seq_length}", file=out, end="")
|
|
print(f";prefill_chunk_size={self.prefill_chunk_size}", file=out, end="")
|
|
print(f";max_history_size={self.max_history_size}", file=out, end="")
|
|
print(f";gpu_memory_utilization={self.gpu_memory_utilization}", file=out, end="")
|
|
print(f";spec_draft_length={self.spec_draft_length}", file=out, end="")
|
|
print(f";spec_tree_width={self.spec_tree_width}", file=out, end="")
|
|
print(f";prefix_cache_mode={self.prefix_cache_mode}", file=out, end="")
|
|
print(
|
|
f";prefix_cache_max_num_recycling_seqs={self.prefix_cache_max_num_recycling_seqs}",
|
|
file=out,
|
|
end="",
|
|
)
|
|
print(f";prefill_mode={self.prefill_mode}", file=out, end="")
|
|
print(f";context_window_size={self.context_window_size}", file=out, end="")
|
|
print(f";sliding_window_size={self.sliding_window_size}", file=out, end="")
|
|
print(f";attention_sink_size={self.attention_sink_size}", file=out, end="")
|
|
print(f";tensor_parallel_shards={self.tensor_parallel_shards}", file=out, end="")
|
|
print(
|
|
f";pipeline_parallel_stages={self.pipeline_parallel_stages}",
|
|
file=out,
|
|
end="",
|
|
)
|
|
print(f";opt={self.opt}", file=out, end="")
|
|
return out.getvalue().rstrip()
|
|
|
|
@staticmethod
|
|
def from_str(source: str) -> "EngineConfigOverride":
|
|
"""Parse engine config override values from a string."""
|
|
parser = argparse.ArgumentParser(description="Engine config override values")
|
|
|
|
parser.add_argument("--max_num_sequence", type=int, default=None)
|
|
parser.add_argument("--max_total_seq_length", type=int, default=None)
|
|
parser.add_argument("--prefill_chunk_size", type=int, default=None)
|
|
parser.add_argument("--max_history_size", type=int, default=None)
|
|
parser.add_argument("--gpu_memory_utilization", type=float, default=None)
|
|
parser.add_argument("--spec_draft_length", type=int, default=None)
|
|
parser.add_argument("--spec_tree_width", type=int, default=None)
|
|
parser.add_argument("--prefix_cache_mode", type=str, default="radix")
|
|
parser.add_argument("--prefix_cache_max_num_recycling_seqs", type=int, default=None)
|
|
parser.add_argument("--prefill_mode", type=str, default="hybrid")
|
|
parser.add_argument("--context_window_size", type=int, default=None)
|
|
parser.add_argument("--sliding_window_size", type=int, default=None)
|
|
parser.add_argument("--attention_sink_size", type=int, default=None)
|
|
parser.add_argument("--tensor_parallel_shards", type=int, default=None)
|
|
parser.add_argument("--pipeline_parallel_stages", type=int, default=None)
|
|
parser.add_argument("--opt", type=str, default=None)
|
|
results = parser.parse_args([f"--{i}" for i in source.split(";") if i])
|
|
return EngineConfigOverride(
|
|
max_num_sequence=results.max_num_sequence,
|
|
max_total_seq_length=results.max_total_seq_length,
|
|
prefill_chunk_size=results.prefill_chunk_size,
|
|
max_history_size=results.max_history_size,
|
|
gpu_memory_utilization=results.gpu_memory_utilization,
|
|
spec_draft_length=results.spec_draft_length,
|
|
spec_tree_width=results.spec_tree_width,
|
|
prefix_cache_mode=results.prefix_cache_mode,
|
|
prefix_cache_max_num_recycling_seqs=results.prefix_cache_max_num_recycling_seqs,
|
|
prefill_mode=results.prefill_mode,
|
|
context_window_size=results.context_window_size,
|
|
sliding_window_size=results.sliding_window_size,
|
|
attention_sink_size=results.attention_sink_size,
|
|
tensor_parallel_shards=results.tensor_parallel_shards,
|
|
pipeline_parallel_stages=results.pipeline_parallel_stages,
|
|
opt=results.opt,
|
|
)
|
|
|
|
|
|
def main(argv):
|
|
"""Parse command line arguments and call `mlc_llm.interface.serve`."""
|
|
parser = ArgumentParser("MLC LLM Serve CLI")
|
|
|
|
parser.add_argument(
|
|
"model",
|
|
type=str,
|
|
help=HELP["model"] + " (required)",
|
|
)
|
|
parser.add_argument(
|
|
"--device",
|
|
type=str,
|
|
default="auto",
|
|
help=HELP["device_deploy"] + ' (default: "%(default)s")',
|
|
)
|
|
parser.add_argument(
|
|
"--model-lib",
|
|
type=str,
|
|
default=None,
|
|
help=HELP["model_lib"] + ' (default: "%(default)s")',
|
|
)
|
|
parser.add_argument(
|
|
"--mode",
|
|
type=str,
|
|
choices=["local", "interactive", "server"],
|
|
default="local",
|
|
help=HELP["mode_serve"] + ' (default: "%(default)s")',
|
|
)
|
|
parser.add_argument(
|
|
"--enable-debug",
|
|
action="store_true",
|
|
help="whether we enable debug end points and debug config when accepting requests",
|
|
)
|
|
parser.add_argument(
|
|
"--additional-models", type=str, nargs="*", help=HELP["additional_models_serve"]
|
|
)
|
|
parser.add_argument(
|
|
"--embedding-model",
|
|
type=str,
|
|
default=None,
|
|
help="Path to the embedding model weight directory (enables /v1/embeddings endpoint)",
|
|
)
|
|
parser.add_argument(
|
|
"--embedding-model-lib",
|
|
type=str,
|
|
default=None,
|
|
help="Path to the compiled embedding model library (.so/.dylib file)",
|
|
)
|
|
parser.add_argument(
|
|
"--speculative-mode",
|
|
type=str,
|
|
choices=["disable", "small_draft", "eagle", "medusa"],
|
|
default="disable",
|
|
help=HELP["speculative_mode_serve"] + ' (default: "%(default)s")',
|
|
)
|
|
parser.add_argument(
|
|
"--prefix-cache-mode",
|
|
type=str,
|
|
choices=["disable", "radix"],
|
|
default="radix",
|
|
help=HELP["prefix_cache_mode_serve"] + ' (default: "%(default)s")',
|
|
)
|
|
parser.add_argument(
|
|
"--prefill-mode",
|
|
type=str,
|
|
choices=["hybrid", "chunked"],
|
|
default="hybrid",
|
|
help=HELP["prefill_mode"] + ' (default: "%(default)s")',
|
|
)
|
|
parser.add_argument(
|
|
"--overrides",
|
|
type=EngineConfigOverride.from_str,
|
|
default="",
|
|
help=HELP["overrides_serve"],
|
|
)
|
|
parser.add_argument("--enable-tracing", action="store_true", help=HELP["enable_tracing_serve"])
|
|
parser.add_argument(
|
|
"--host",
|
|
type=str,
|
|
default="127.0.0.1",
|
|
help="host name" + ' (default: "%(default)s")',
|
|
)
|
|
parser.add_argument(
|
|
"--port",
|
|
type=int,
|
|
default=8000,
|
|
help="port" + ' (default: "%(default)s")',
|
|
)
|
|
parser.add_argument("--allow-credentials", action="store_true", help="allow credentials")
|
|
parser.add_argument(
|
|
"--allow-origins",
|
|
type=json.loads,
|
|
default=["*"],
|
|
help="allowed origins" + ' (default: "%(default)s")',
|
|
)
|
|
parser.add_argument(
|
|
"--allow-methods",
|
|
type=json.loads,
|
|
default=["*"],
|
|
help="allowed methods" + ' (default: "%(default)s")',
|
|
)
|
|
parser.add_argument(
|
|
"--allow-headers",
|
|
type=json.loads,
|
|
default=["*"],
|
|
help="allowed headers" + ' (default: "%(default)s")',
|
|
)
|
|
parser.add_argument(
|
|
"--api-key",
|
|
type=str,
|
|
default=None,
|
|
help="API key for authentication. If not provided, authentication is disabled.",
|
|
)
|
|
parsed = parser.parse_args(argv)
|
|
|
|
additional_models = []
|
|
if parsed.additional_models is not None:
|
|
for additional_model in parsed.additional_models:
|
|
splits = additional_model.split(",", maxsplit=1)
|
|
if len(splits) == 2:
|
|
additional_models.append((splits[0], splits[1]))
|
|
else:
|
|
additional_models.append(splits[0])
|
|
|
|
serve(
|
|
model=parsed.model,
|
|
device=parsed.device,
|
|
model_lib=parsed.model_lib,
|
|
mode=parsed.mode,
|
|
enable_debug=parsed.enable_debug,
|
|
additional_models=additional_models,
|
|
embedding_model=parsed.embedding_model,
|
|
embedding_model_lib=parsed.embedding_model_lib,
|
|
tensor_parallel_shards=parsed.overrides.tensor_parallel_shards,
|
|
pipeline_parallel_stages=parsed.overrides.pipeline_parallel_stages,
|
|
opt=parsed.overrides.opt,
|
|
speculative_mode=parsed.speculative_mode,
|
|
prefix_cache_mode=parsed.prefix_cache_mode,
|
|
max_num_sequence=parsed.overrides.max_num_sequence,
|
|
max_total_sequence_length=parsed.overrides.max_total_seq_length,
|
|
max_single_sequence_length=parsed.overrides.context_window_size,
|
|
prefill_chunk_size=parsed.overrides.prefill_chunk_size,
|
|
sliding_window_size=parsed.overrides.sliding_window_size,
|
|
attention_sink_size=parsed.overrides.attention_sink_size,
|
|
max_history_size=parsed.overrides.max_history_size,
|
|
gpu_memory_utilization=parsed.overrides.gpu_memory_utilization,
|
|
spec_draft_length=parsed.overrides.spec_draft_length,
|
|
spec_tree_width=parsed.overrides.spec_tree_width,
|
|
prefix_cache_max_num_recycling_seqs=parsed.overrides.prefix_cache_max_num_recycling_seqs,
|
|
prefill_mode=parsed.prefill_mode,
|
|
enable_tracing=parsed.enable_tracing,
|
|
host=parsed.host,
|
|
port=parsed.port,
|
|
allow_credentials=parsed.allow_credentials,
|
|
allow_origins=parsed.allow_origins,
|
|
allow_methods=parsed.allow_methods,
|
|
allow_headers=parsed.allow_headers,
|
|
api_key=parsed.api_key,
|
|
)
|