[ci] CI supports use cached models (#7874)
This commit is contained in:
@@ -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,
|
||||
|
||||
40
scripts/ci_cache_models.sh
Executable file
40
scripts/ci_cache_models.sh
Executable file
@@ -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
|
||||
Reference in New Issue
Block a user