diff --git a/python/sglang/test/test_utils.py b/python/sglang/test/test_utils.py index 23d44eb35..132a796d4 100644 --- a/python/sglang/test/test_utils.py +++ b/python/sglang/test/test_utils.py @@ -2,6 +2,7 @@ import argparse import copy +import json import logging import os import random @@ -102,6 +103,15 @@ def is_in_amd_ci(): return get_bool_env_var("SGLANG_AMD_CI") +def _use_cached_default_models(model_repo: str): + cache_dir = os.getenv("DEFAULT_MODEL_CACHE_DIR") + if cache_dir and model_repo: + model_path = os.path.join(cache_dir, model_repo) + if os.path.isdir(model_path): + return os.path.abspath(model_path) + return "" + + if is_in_ci(): DEFAULT_PORT_FOR_SRT_TEST_RUNNER = ( 5000 + int(os.environ.get("CUDA_VISIBLE_DEVICES", "0")[0]) * 100 @@ -419,6 +429,31 @@ def get_call_select(args: argparse.Namespace): return func +def _get_default_models(): + import inspect + + current_module = inspect.getmodule(_get_default_models) + default_models = set() + for name, value in current_module.__dict__.items(): + if ( + isinstance(name, str) + and "DEFAULT_" in name + and "MODEL_" in name + and isinstance(value, str) + ): + if "," in value: + parts = [part.strip() for part in value.split(",")] + default_models.update(parts) + else: + default_models.add(value.strip()) + return json.dumps(list(default_models)) + + +def try_cached_model(model_repo: str): + model_dir = _use_cached_default_models(model_repo) + return model_dir if model_dir else model_repo + + def popen_launch_server( model: str, base_url: str, diff --git a/scripts/ci_cache_models.sh b/scripts/ci_cache_models.sh new file mode 100755 index 000000000..0ebe6c055 --- /dev/null +++ b/scripts/ci_cache_models.sh @@ -0,0 +1,40 @@ +#!/bin/bash +set -euxo pipefail + +mapfile -t models < <(python3 -c "from sglang.test.test_utils import _get_default_models; print(_get_default_models())" | jq -r '.[]') + +if [ ${#models[@]} -eq 0 ]; then + echo "Failed to get default models." + exit 1 +fi + +cache_dir="${DEFAULT_MODEL_CACHE_DIR:-}" + +if [ -z "$cache_dir" ]; then + echo "DEFAULT_MODEL_CACHE_DIR environment variable is not set." + exit 1 +fi + +failed_models=() +for model in "${models[@]}"; do + local_model_dir="$cache_dir/$model" + echo "Caching model: $model to $local_model_dir" + mkdir -p "$local_model_dir" + + if ! huggingface-cli download "$model" \ + --local-dir "$local_model_dir" \ + --local-dir-use-symlinks False 2>/dev/null; then + echo "WARNING: Failed to cache model: $model" + rm -rf "$local_model_dir" + failed_models+=("$model") + continue + fi + echo "Successfully cached model: $model" +done + +if [ ${#failed_models[@]} -gt 0 ]; then + echo -e "\n[Summary] Failed to cache following models:" + printf ' - %s\n' "${failed_models[@]}" +else + echo -e "\n[Summary] All models cached successfully" +fi