[CLI] Add --model-type override and keep launch_server supported (#19523)
This commit is contained in:
Liangsheng Yin
@@ -10,6 +10,39 @@ from sglang.srt.utils import kill_process_tree
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _extract_model_type_override(extra_argv):
|
||||
"""Extract and remove --model-type override from argv."""
|
||||
model_type = "auto"
|
||||
filtered_argv = []
|
||||
i = 0
|
||||
while i < len(extra_argv):
|
||||
arg = extra_argv[i]
|
||||
if arg == "--model-type":
|
||||
if i + 1 >= len(extra_argv):
|
||||
raise Exception(
|
||||
"Error: --model-type requires a value. "
|
||||
"Valid values are: auto, llm, diffusion."
|
||||
)
|
||||
model_type = extra_argv[i + 1]
|
||||
i += 2
|
||||
continue
|
||||
|
||||
if arg.startswith("--model-type="):
|
||||
model_type = arg.split("=", 1)[1]
|
||||
i += 1
|
||||
continue
|
||||
|
||||
filtered_argv.append(arg)
|
||||
i += 1
|
||||
|
||||
if model_type not in ("auto", "llm", "diffusion"):
|
||||
raise Exception(
|
||||
f"Error: invalid --model-type '{model_type}'. "
|
||||
"Valid values are: auto, llm, diffusion."
|
||||
)
|
||||
return model_type, filtered_argv
|
||||
|
||||
|
||||
def serve(args, extra_argv):
|
||||
if any(h in extra_argv for h in ("-h", "--help")):
|
||||
# Since the server type is determined by the model, and we don't have a model path,
|
||||
@@ -22,6 +55,10 @@ def serve(args, extra_argv):
|
||||
"This command can launch either a standard language model server or a diffusion model server."
|
||||
)
|
||||
print("The server type is determined by the model path.\n")
|
||||
print(
|
||||
"Optional override: --model-type {auto,llm,diffusion} "
|
||||
"(default: auto, fallback to LLM on detection failure).\n"
|
||||
)
|
||||
print("For specific arguments, please provide a model_path.")
|
||||
print("\n--- Help for Standard Language Model Server ---")
|
||||
from sglang.srt.server_args import prepare_server_args
|
||||
@@ -41,11 +78,19 @@ def serve(args, extra_argv):
|
||||
parser.print_help()
|
||||
return
|
||||
|
||||
model_path = get_model_path(extra_argv)
|
||||
model_type, dispatch_argv = _extract_model_type_override(extra_argv)
|
||||
model_path = get_model_path(dispatch_argv)
|
||||
try:
|
||||
is_diffusion_model = get_is_diffusion_model(model_path)
|
||||
if is_diffusion_model:
|
||||
logger.info("Diffusion model detected")
|
||||
if model_type == "auto":
|
||||
is_diffusion_model = get_is_diffusion_model(model_path)
|
||||
if is_diffusion_model:
|
||||
logger.info("Diffusion model detected")
|
||||
else:
|
||||
is_diffusion_model = model_type == "diffusion"
|
||||
logger.info(
|
||||
"Dispatch override enabled: --model-type=%s " "(skip auto detection)",
|
||||
model_type,
|
||||
)
|
||||
|
||||
if is_diffusion_model:
|
||||
# Logic for Diffusion Models
|
||||
@@ -58,7 +103,7 @@ def serve(args, extra_argv):
|
||||
description="SGLang Diffusion Model Serving"
|
||||
)
|
||||
add_multimodal_gen_serve_args(parser)
|
||||
parsed_args, remaining_argv = parser.parse_known_args(extra_argv)
|
||||
parsed_args, remaining_argv = parser.parse_known_args(dispatch_argv)
|
||||
|
||||
execute_serve_cmd(parsed_args, remaining_argv)
|
||||
else:
|
||||
@@ -68,7 +113,7 @@ def serve(args, extra_argv):
|
||||
|
||||
# Add a dummy argument for the program name, expected by prepare_server_args
|
||||
# as it typically processes sys.argv
|
||||
server_args = prepare_server_args(extra_argv)
|
||||
server_args = prepare_server_args(dispatch_argv)
|
||||
|
||||
run_server(server_args)
|
||||
finally:
|
||||
|
||||
@@ -1,160 +1,53 @@
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import subprocess
|
||||
import tempfile
|
||||
from functools import lru_cache
|
||||
from typing import Optional
|
||||
|
||||
import filelock
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
temp_dir = tempfile.gettempdir()
|
||||
|
||||
|
||||
def _get_lock(model_name_or_path: str, cache_dir: Optional[str] = None):
|
||||
lock_dir = cache_dir or temp_dir
|
||||
os.makedirs(os.path.dirname(lock_dir), exist_ok=True)
|
||||
model_name = model_name_or_path.replace("/", "-")
|
||||
hash_name = hashlib.sha256(model_name.encode()).hexdigest()
|
||||
lock_file_name = hash_name + model_name + ".lock"
|
||||
lock = filelock.FileLock(os.path.join(lock_dir, lock_file_name), mode=0o666)
|
||||
return lock
|
||||
|
||||
|
||||
def _hf_hub_download(
|
||||
repo_id: str,
|
||||
filename: str,
|
||||
local_dir: Optional[str] = None,
|
||||
**kwargs,
|
||||
) -> str:
|
||||
"""Unified hf_hub_download supporting both Hugging Face Hub and ModelScope."""
|
||||
if envs.SGLANG_USE_MODELSCOPE.get():
|
||||
from modelscope import model_file_download
|
||||
|
||||
return model_file_download(
|
||||
model_id=repo_id,
|
||||
file_path=filename,
|
||||
cache_dir=local_dir,
|
||||
**kwargs,
|
||||
)
|
||||
else:
|
||||
from huggingface_hub import hf_hub_download as _hf_hub_download_impl
|
||||
|
||||
return _hf_hub_download_impl(
|
||||
repo_id=repo_id,
|
||||
filename=filename,
|
||||
local_dir=local_dir,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
# Copied and adapted from hf_diffusers_utils.py
|
||||
def _maybe_download_model(
|
||||
model_name_or_path: str, local_dir: str | None = None, download: bool = True
|
||||
) -> str:
|
||||
"""
|
||||
Resolve a model path. If it's a local directory, return it.
|
||||
If it's a Hugging Face Hub ID, download only the config file
|
||||
(`model_index.json` or `config.json`) and return its directory.
|
||||
|
||||
Args:
|
||||
model_name_or_path: Local path or Hugging Face Hub model ID
|
||||
local_dir: Local directory to save the downloaded file (if any)
|
||||
download: Whether to download from Hugging Face Hub when needed
|
||||
|
||||
Returns:
|
||||
Local directory path that contains the downloaded config file, or the original local directory.
|
||||
"""
|
||||
|
||||
if os.path.exists(model_name_or_path):
|
||||
logger.debug("Model already exists locally")
|
||||
return model_name_or_path
|
||||
|
||||
if not download:
|
||||
return model_name_or_path
|
||||
|
||||
with _get_lock(model_name_or_path):
|
||||
# Try `model_index.json` first (diffusers models)
|
||||
source_hub = "MS Hub" if envs.SGLANG_USE_MODELSCOPE.get() else "HF Hub"
|
||||
try:
|
||||
logger.debug(
|
||||
"Downloading model_index.json from %s for %s...",
|
||||
source_hub,
|
||||
model_name_or_path,
|
||||
)
|
||||
file_path = _hf_hub_download(
|
||||
repo_id=model_name_or_path,
|
||||
filename="model_index.json",
|
||||
local_dir=local_dir,
|
||||
)
|
||||
logger.debug("Downloaded to %s", file_path)
|
||||
return os.path.dirname(file_path)
|
||||
except Exception as e_index:
|
||||
logger.debug("model_index.json not found or failed: %s", e_index)
|
||||
|
||||
# Fallback to `config.json`
|
||||
try:
|
||||
logger.debug(
|
||||
"Downloading config.json from %s for %s...",
|
||||
source_hub,
|
||||
model_name_or_path,
|
||||
)
|
||||
file_path = _hf_hub_download(
|
||||
repo_id=model_name_or_path,
|
||||
filename="config.json",
|
||||
local_dir=local_dir,
|
||||
)
|
||||
logger.debug("Downloaded to %s", file_path)
|
||||
return os.path.dirname(file_path)
|
||||
except Exception as e_config:
|
||||
raise ValueError(
|
||||
(
|
||||
"Could not find model locally at %s and failed to download "
|
||||
"model_index.json/config.json from %s: %s"
|
||||
)
|
||||
% (model_name_or_path, source_hub, e_config)
|
||||
) from e_config
|
||||
|
||||
|
||||
# Copied and adapted from hf_diffusers_utils.py
|
||||
def is_diffusers_model_path(model_path: str) -> True:
|
||||
"""
|
||||
Verify if the model directory contains a valid diffusers configuration.
|
||||
|
||||
Args:
|
||||
model_path: Path to the model directory
|
||||
|
||||
Returns:
|
||||
The loaded model configuration as a dictionary if the model is a diffusers model
|
||||
None if the model is not a diffusers model
|
||||
"""
|
||||
|
||||
# Prefer model_index.json which indicates a diffusers pipeline
|
||||
config_path = os.path.join(model_path, "model_index.json")
|
||||
def _is_diffusers_model_dir(model_dir: str) -> bool:
|
||||
"""Check if a local directory contains a valid diffusers model_index.json."""
|
||||
config_path = os.path.join(model_dir, "model_index.json")
|
||||
if not os.path.exists(config_path):
|
||||
return False
|
||||
|
||||
# Load the config
|
||||
with open(config_path) as f:
|
||||
config = json.load(f)
|
||||
|
||||
# Verify diffusers version exists
|
||||
if "_diffusers_version" not in config:
|
||||
return "_diffusers_version" in config
|
||||
|
||||
|
||||
def get_is_diffusion_model(model_path: str) -> bool:
|
||||
"""Detect whether model_path points to a diffusion model.
|
||||
|
||||
For local directories, checks the filesystem directly.
|
||||
For HF/ModelScope model IDs, attempts to fetch only model_index.json.
|
||||
Returns False on any failure (network error, 404, offline mode, etc.)
|
||||
so that the caller falls through to the standard LLM server path.
|
||||
"""
|
||||
if os.path.isdir(model_path):
|
||||
return _is_diffusers_model_dir(model_path)
|
||||
|
||||
try:
|
||||
if envs.SGLANG_USE_MODELSCOPE.get():
|
||||
from modelscope import model_file_download
|
||||
|
||||
file_path = model_file_download(
|
||||
model_id=model_path, file_path="model_index.json"
|
||||
)
|
||||
else:
|
||||
from huggingface_hub import hf_hub_download
|
||||
|
||||
file_path = hf_hub_download(repo_id=model_path, filename="model_index.json")
|
||||
|
||||
return _is_diffusers_model_dir(os.path.dirname(file_path))
|
||||
except Exception as e:
|
||||
logger.debug("Failed to auto-detect diffusion model for %s: %s", model_path, e)
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def get_is_diffusion_model(model_path: str):
|
||||
model_path = _maybe_download_model(model_path)
|
||||
is_diffusion_model = is_diffusers_model_path(model_path)
|
||||
if is_diffusion_model:
|
||||
logger.info("Diffusion model detected")
|
||||
return is_diffusion_model
|
||||
|
||||
|
||||
def get_model_path(extra_argv):
|
||||
|
||||
@@ -29,6 +29,16 @@ def run_server(server_args):
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import warnings
|
||||
|
||||
warnings.warn(
|
||||
"'python -m sglang.launch_server' is still supported, but "
|
||||
"'sglang serve' is the recommended entrypoint.\n"
|
||||
" Example: sglang serve --model-path <model> [options]",
|
||||
UserWarning,
|
||||
stacklevel=1,
|
||||
)
|
||||
|
||||
server_args = prepare_server_args(sys.argv[1:])
|
||||
|
||||
try:
|
||||
|
||||
Reference in New Issue
Block a user