chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,264 @@
|
||||
"""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,
|
||||
)
|
||||
Reference in New Issue
Block a user