ci: enable offline mode when local cache is complete to avoid HF Hub … (#16121)
This commit is contained in:
@@ -590,39 +590,56 @@ class ModelConfig:
|
||||
|
||||
hf_api = HfApi()
|
||||
try:
|
||||
# Retry HF API call up to 3 times
|
||||
file_exists = retry(
|
||||
lambda: hf_api.file_exists(
|
||||
self.model_path, "hf_quant_config.json"
|
||||
),
|
||||
max_retry=2,
|
||||
initial_delay=1.0,
|
||||
max_delay=5.0,
|
||||
# In offline mode, skip file_exists check to avoid OfflineModeIsEnabled error
|
||||
# Instead, directly try to download/read from cache with local_files_only
|
||||
file_exists = False # Initialize to avoid UnboundLocalError
|
||||
if not huggingface_hub.constants.HF_HUB_OFFLINE:
|
||||
# Online mode: check if file exists before attempting download (optimization)
|
||||
file_exists = retry(
|
||||
lambda: hf_api.file_exists(
|
||||
self.model_path, "hf_quant_config.json"
|
||||
),
|
||||
max_retry=2,
|
||||
initial_delay=1.0,
|
||||
max_delay=5.0,
|
||||
)
|
||||
if not file_exists:
|
||||
# File doesn't exist on hub, no need to try downloading
|
||||
return quant_cfg # None
|
||||
|
||||
# Download (online mode) or read from cache (offline mode)
|
||||
if envs.SGLANG_USE_MODELSCOPE.get():
|
||||
quant_config_file = model_file_download(
|
||||
model_id=self.model_path,
|
||||
file_path="hf_quant_config.json",
|
||||
revision=self.revision,
|
||||
)
|
||||
else:
|
||||
quant_config_file = hf_hub_download(
|
||||
repo_id=self.model_path,
|
||||
filename="hf_quant_config.json",
|
||||
revision=self.revision,
|
||||
local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,
|
||||
)
|
||||
with open(quant_config_file) as f:
|
||||
quant_config_dict = json.load(f)
|
||||
quant_cfg = self._parse_modelopt_quant_config(quant_config_dict)
|
||||
except huggingface_hub.errors.LocalEntryNotFoundError:
|
||||
# Offline mode and file not in cache - this is normal for non-quantized models
|
||||
logger.debug(
|
||||
f"hf_quant_config.json not found in cache for {self.model_path} "
|
||||
"(offline mode, normal for non-quantized models)"
|
||||
)
|
||||
if file_exists:
|
||||
# Download and parse the quantization config for remote models
|
||||
if envs.SGLANG_USE_MODELSCOPE.get():
|
||||
quant_config_file = model_file_download(
|
||||
model_id=self.model_path,
|
||||
file_path="hf_quant_config.json",
|
||||
revision=self.revision,
|
||||
)
|
||||
else:
|
||||
quant_config_file = hf_hub_download(
|
||||
repo_id=self.model_path,
|
||||
filename="hf_quant_config.json",
|
||||
revision=self.revision,
|
||||
)
|
||||
with open(quant_config_file) as f:
|
||||
quant_config_dict = json.load(f)
|
||||
quant_cfg = self._parse_modelopt_quant_config(quant_config_dict)
|
||||
except huggingface_hub.errors.OfflineModeIsEnabled:
|
||||
# Should not reach here after our changes, but keep for safety
|
||||
logger.warning(
|
||||
"Offline mode is enabled, skipping hf_quant_config.json check"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"Failed to check hf_quant_config.json: {self.model_path} {e}"
|
||||
"Failed to load hf_quant_config.json for model %s: %s",
|
||||
self.model_path,
|
||||
e,
|
||||
)
|
||||
elif os.path.exists(os.path.join(self.model_path, "hf_quant_config.json")):
|
||||
quant_config_file = os.path.join(
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -566,7 +566,9 @@ class DefaultModelLoader(BaseModelLoader):
|
||||
)
|
||||
|
||||
hf_config = AutoConfig.from_pretrained(
|
||||
model_config.model_path, trust_remote_code=True
|
||||
model_config.model_path,
|
||||
trust_remote_code=True,
|
||||
local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,
|
||||
)
|
||||
with init_empty_weights():
|
||||
torch_dtype = getattr(hf_config, "torch_dtype", torch.float16)
|
||||
@@ -599,6 +601,7 @@ class DefaultModelLoader(BaseModelLoader):
|
||||
device_map=device_map,
|
||||
**model_kwargs,
|
||||
trust_remote_code=True,
|
||||
local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,
|
||||
)
|
||||
# Handle both legacy modelopt_quant and unified quantization flags
|
||||
if hasattr(model_config, "modelopt_quant") and model_config.modelopt_quant:
|
||||
|
||||
@@ -562,6 +562,249 @@ def popen_with_error_check(command: list[str], allow_exit: bool = False):
|
||||
return process
|
||||
|
||||
|
||||
def _try_enable_offline_mode_if_cache_complete(
|
||||
model_name_or_path: str, env: dict, other_args: Optional[list[str]] = None
|
||||
) -> Optional[str]:
|
||||
"""
|
||||
CI helper: Check if model cache is complete and enable offline mode.
|
||||
|
||||
Uses per-run validation markers that are NOT shared across runners.
|
||||
Each runner independently validates its cache using lightweight checks
|
||||
before enabling offline mode.
|
||||
|
||||
IMPORTANT: Even if a per-run marker exists, this function ALWAYS validates
|
||||
the current launch's requirements (e.g., hf_quant_config.json for modelopt).
|
||||
The marker is only a hint that this snapshot was validated earlier in the run.
|
||||
|
||||
Args:
|
||||
model_name_or_path: Model identifier or path
|
||||
env: Environment dict to modify (will add HF_HUB_OFFLINE=1 if validation passes)
|
||||
other_args: Launch command arguments (used to detect quantization requirement)
|
||||
|
||||
Returns:
|
||||
Per-run marker path if offline mode was enabled, None otherwise
|
||||
"""
|
||||
from sglang.srt.model_loader.ci_weight_validation import (
|
||||
_get_per_run_marker_path,
|
||||
_read_per_run_marker,
|
||||
_write_per_run_marker,
|
||||
validate_cache_lightweight,
|
||||
)
|
||||
from sglang.srt.utils import find_local_repo_dir
|
||||
|
||||
other_args = other_args or []
|
||||
|
||||
# Fast-path: If subprocess env already has HF_HUB_OFFLINE=1, skip
|
||||
if env.get("HF_HUB_OFFLINE") == "1":
|
||||
print(
|
||||
f"CI_OFFLINE: Subprocess env already has HF_HUB_OFFLINE=1, skip - {model_name_or_path}"
|
||||
)
|
||||
return None
|
||||
|
||||
# Skip if already a local path
|
||||
if os.path.isdir(model_name_or_path):
|
||||
return None
|
||||
|
||||
# Try to find local snapshot
|
||||
try:
|
||||
snapshot_dir = find_local_repo_dir(model_name_or_path, revision=None)
|
||||
if not snapshot_dir or not os.path.isdir(snapshot_dir):
|
||||
return None
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
# Detect if quantization requires hf_quant_config.json
|
||||
# Do this BEFORE checking marker to ensure current launch requirements are known
|
||||
requires_hf_quant_config = False
|
||||
for i, arg in enumerate(other_args):
|
||||
if arg == "--quantization" and i + 1 < len(other_args):
|
||||
quant_value = other_args[i + 1].lower()
|
||||
if quant_value in ["modelopt_fp4", "modelopt_fp8", "modelopt"]:
|
||||
requires_hf_quant_config = True
|
||||
break
|
||||
|
||||
# Check per-run marker (fast hint - snapshot validated earlier in this run)
|
||||
per_run_marker = _read_per_run_marker(snapshot_dir)
|
||||
if per_run_marker is not None:
|
||||
# Marker exists, but STILL validate for current launch requirements
|
||||
# This prevents a test without --quantization from enabling offline
|
||||
# for a later test with --quantization that needs hf_quant_config.json
|
||||
is_valid = validate_cache_lightweight(snapshot_dir, requires_hf_quant_config)
|
||||
|
||||
if not is_valid:
|
||||
# Current launch requirements not met, ignore marker
|
||||
print(
|
||||
f"CI_OFFLINE: Per-run marker found but current validation failed "
|
||||
f"(requires_hf_quant_config={requires_hf_quant_config}), "
|
||||
f"will use online mode - {model_name_or_path}"
|
||||
)
|
||||
return None
|
||||
|
||||
# Marker exists and current validation passed
|
||||
env["HF_HUB_OFFLINE"] = "1"
|
||||
marker_path = _get_per_run_marker_path(snapshot_dir)
|
||||
print(
|
||||
f"CI_OFFLINE: Per-run marker found and current validation passed "
|
||||
f"(requires_hf_quant_config={requires_hf_quant_config}), "
|
||||
f"enabling offline mode - {model_name_or_path}"
|
||||
)
|
||||
return marker_path
|
||||
|
||||
# No per-run marker - perform lightweight validation
|
||||
is_valid = validate_cache_lightweight(snapshot_dir, requires_hf_quant_config)
|
||||
|
||||
if not is_valid:
|
||||
# Validation failed - cache is incomplete on this runner
|
||||
print(
|
||||
f"CI_OFFLINE: Cache validation failed "
|
||||
f"(requires_hf_quant_config={requires_hf_quant_config}), "
|
||||
f"will use online mode - {model_name_or_path}"
|
||||
)
|
||||
return None
|
||||
|
||||
# Validation passed - enable offline mode and write per-run marker
|
||||
env["HF_HUB_OFFLINE"] = "1"
|
||||
|
||||
# Write per-run marker for subsequent tests in this run
|
||||
_write_per_run_marker(snapshot_dir, model_name_or_path)
|
||||
|
||||
# Return marker path for potential invalidation if offline launch fails
|
||||
marker_path = _get_per_run_marker_path(snapshot_dir)
|
||||
|
||||
snapshot_basename = os.path.basename(snapshot_dir)
|
||||
print(
|
||||
f"CI_OFFLINE: Enabled HF_HUB_OFFLINE=1 for subprocess - "
|
||||
f"validation passed for {model_name_or_path} "
|
||||
f"(snapshot={snapshot_basename}, requires_hf_quant_config={requires_hf_quant_config})"
|
||||
)
|
||||
|
||||
return marker_path
|
||||
|
||||
|
||||
def _create_clean_subprocess_env(env: dict) -> dict:
|
||||
"""Create a clean subprocess environment without internal CI keys.
|
||||
|
||||
Removes all keys starting with '_CI_OFFLINE_' or 'CI_OFFLINE' to prevent
|
||||
leaking implementation details to the server subprocess.
|
||||
|
||||
Args:
|
||||
env: Source environment dict
|
||||
|
||||
Returns:
|
||||
Clean copy of environment dict
|
||||
"""
|
||||
child_env = env.copy()
|
||||
keys_to_remove = [
|
||||
k for k in child_env if k.startswith(("_CI_OFFLINE_", "CI_OFFLINE_"))
|
||||
]
|
||||
for k in keys_to_remove:
|
||||
del child_env[k]
|
||||
return child_env
|
||||
|
||||
|
||||
def _launch_server_process(
|
||||
command: List[str],
|
||||
env: dict,
|
||||
return_stdout_stderr: Optional[tuple],
|
||||
model: str,
|
||||
) -> subprocess.Popen:
|
||||
"""Launch server subprocess with clean environment.
|
||||
|
||||
Args:
|
||||
command: Command list for subprocess
|
||||
env: Environment dict (will be cleaned before use)
|
||||
return_stdout_stderr: Optional tuple of (stdout_file, stderr_file) for output capture
|
||||
model: Model name for logging
|
||||
|
||||
Returns:
|
||||
Started subprocess.Popen object
|
||||
"""
|
||||
child_env = _create_clean_subprocess_env(env)
|
||||
|
||||
hf_hub_offline = child_env.get("HF_HUB_OFFLINE", "0")
|
||||
print(f"CI_OFFLINE: Launching server HF_HUB_OFFLINE={hf_hub_offline} model={model}")
|
||||
|
||||
if return_stdout_stderr:
|
||||
proc = subprocess.Popen(
|
||||
command,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
env=child_env,
|
||||
text=True,
|
||||
bufsize=1,
|
||||
)
|
||||
|
||||
def _dump(src, sinks):
|
||||
for line in iter(src.readline, ""):
|
||||
for sink in sinks:
|
||||
sink.write(line)
|
||||
sink.flush()
|
||||
src.close()
|
||||
|
||||
threading.Thread(
|
||||
target=_dump,
|
||||
args=(proc.stdout, [return_stdout_stderr[0], sys.stdout]),
|
||||
daemon=True,
|
||||
).start()
|
||||
threading.Thread(
|
||||
target=_dump,
|
||||
args=(proc.stderr, [return_stdout_stderr[1], sys.stderr]),
|
||||
daemon=True,
|
||||
).start()
|
||||
else:
|
||||
proc = subprocess.Popen(command, stdout=None, stderr=None, env=child_env)
|
||||
|
||||
return proc
|
||||
|
||||
|
||||
def _wait_for_server_health(
|
||||
proc: subprocess.Popen,
|
||||
base_url: str,
|
||||
api_key: Optional[str],
|
||||
timeout_duration: float,
|
||||
) -> Tuple[bool, Optional[str]]:
|
||||
"""Wait for server health check to pass.
|
||||
|
||||
Args:
|
||||
proc: Server subprocess
|
||||
base_url: Base URL for health check
|
||||
api_key: Optional API key for authorization
|
||||
timeout_duration: Maximum wait time in seconds
|
||||
|
||||
Returns:
|
||||
Tuple of (success, error_message)
|
||||
"""
|
||||
start_time = time.perf_counter()
|
||||
with requests.Session() as session:
|
||||
while time.perf_counter() - start_time < timeout_duration:
|
||||
return_code = proc.poll()
|
||||
if return_code is not None:
|
||||
return False, f"Server process exited with code {return_code}"
|
||||
|
||||
try:
|
||||
headers = {
|
||||
"Content-Type": "application/json; charset=utf-8",
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
}
|
||||
response = session.get(
|
||||
f"{base_url}/health_generate",
|
||||
headers=headers,
|
||||
timeout=5,
|
||||
)
|
||||
if response.status_code == 200:
|
||||
return True, None
|
||||
except requests.RequestException:
|
||||
pass
|
||||
|
||||
return_code = proc.poll()
|
||||
if return_code is not None:
|
||||
return False, f"Server unexpectedly exited (return_code={return_code})"
|
||||
|
||||
time.sleep(10)
|
||||
|
||||
return False, "Server failed to start within the timeout period"
|
||||
|
||||
|
||||
def popen_launch_server(
|
||||
model: str,
|
||||
base_url: str,
|
||||
@@ -574,11 +817,22 @@ def popen_launch_server(
|
||||
pd_separated: bool = False,
|
||||
num_replicas: Optional[int] = None,
|
||||
):
|
||||
"""Launch a server process with automatic device detection.
|
||||
"""Launch a server process with automatic device detection and offline/online retry.
|
||||
|
||||
Args:
|
||||
device: Device type ("auto", "cuda", "rocm" or "cpu").
|
||||
If "auto", will detect available platforms automatically.
|
||||
model: Model path or identifier
|
||||
base_url: Base URL for the server
|
||||
timeout: Timeout for server startup
|
||||
api_key: Optional API key for authentication
|
||||
other_args: Additional command line arguments
|
||||
env: Environment dict for subprocess
|
||||
return_stdout_stderr: Optional tuple for output capture
|
||||
device: Device type ("auto", "cuda", "rocm" or "cpu")
|
||||
pd_separated: Whether to use PD separated mode
|
||||
num_replicas: Number of replicas for mixed PD mode
|
||||
|
||||
Returns:
|
||||
Started subprocess.Popen object
|
||||
"""
|
||||
other_args = other_args or []
|
||||
|
||||
@@ -588,6 +842,25 @@ def popen_launch_server(
|
||||
other_args = list(other_args)
|
||||
other_args += ["--device", str(device)]
|
||||
|
||||
# CI-specific: Validate cache and enable offline mode if complete
|
||||
if env is None:
|
||||
env = os.environ.copy()
|
||||
else:
|
||||
env = env.copy()
|
||||
|
||||
# Store per-run marker path for potential invalidation
|
||||
per_run_marker_path = None
|
||||
try:
|
||||
from sglang.utils import is_in_ci
|
||||
|
||||
if is_in_ci():
|
||||
per_run_marker_path = _try_enable_offline_mode_if_cache_complete(
|
||||
model, env, other_args
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"CI cache validation failed (non-fatal): {e}")
|
||||
|
||||
# Build server command
|
||||
_, host, port = base_url.split(":")
|
||||
host = host[2:]
|
||||
|
||||
@@ -607,104 +880,82 @@ def popen_launch_server(
|
||||
]
|
||||
|
||||
if pd_separated or use_mixed_pd_engine:
|
||||
command.extend(
|
||||
[
|
||||
"--lb-host",
|
||||
host,
|
||||
"--lb-port",
|
||||
port,
|
||||
]
|
||||
)
|
||||
command.extend(["--lb-host", host, "--lb-port", port])
|
||||
else:
|
||||
command.extend(
|
||||
[
|
||||
"--host",
|
||||
host,
|
||||
"--port",
|
||||
port,
|
||||
]
|
||||
)
|
||||
command.extend(["--host", host, "--port", port])
|
||||
|
||||
if use_mixed_pd_engine:
|
||||
command.extend(
|
||||
[
|
||||
"--mixed",
|
||||
"--num-replicas",
|
||||
str(num_replicas),
|
||||
]
|
||||
)
|
||||
command.extend(["--mixed", "--num-replicas", str(num_replicas)])
|
||||
|
||||
if api_key:
|
||||
command += ["--api-key", api_key]
|
||||
|
||||
print(f"command={shlex.join(command)}")
|
||||
|
||||
if return_stdout_stderr:
|
||||
process = subprocess.Popen(
|
||||
command,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
env=env,
|
||||
text=True,
|
||||
bufsize=1,
|
||||
# Track if offline mode was enabled for potential retry
|
||||
offline_enabled = env.get("HF_HUB_OFFLINE") == "1"
|
||||
|
||||
# First launch attempt
|
||||
process = _launch_server_process(command, env, return_stdout_stderr, model)
|
||||
success, error_msg = _wait_for_server_health(process, base_url, api_key, timeout)
|
||||
|
||||
# If offline launch failed and offline was enabled, retry with online mode
|
||||
if not success and offline_enabled:
|
||||
print(
|
||||
f"CI_OFFLINE: Offline launch failed ({error_msg}), retrying with online mode..."
|
||||
)
|
||||
|
||||
def _dump(src, sinks):
|
||||
for line in iter(src.readline, ""):
|
||||
for sink in sinks:
|
||||
sink.write(line)
|
||||
sink.flush()
|
||||
src.close()
|
||||
|
||||
threading.Thread(
|
||||
target=_dump,
|
||||
args=(process.stdout, [return_stdout_stderr[0], sys.stdout]),
|
||||
daemon=True,
|
||||
).start()
|
||||
threading.Thread(
|
||||
target=_dump,
|
||||
args=(process.stderr, [return_stdout_stderr[1], sys.stderr]),
|
||||
daemon=True,
|
||||
).start()
|
||||
else:
|
||||
process = subprocess.Popen(command, stdout=None, stderr=None, env=env)
|
||||
|
||||
start_time = time.perf_counter()
|
||||
with requests.Session() as session:
|
||||
while time.perf_counter() - start_time < timeout:
|
||||
return_code = process.poll()
|
||||
if return_code is not None:
|
||||
# Server failed to start (non-zero exit code) or crashed
|
||||
raise Exception(
|
||||
f"Server process exited with code {return_code}. "
|
||||
"Check server logs for errors."
|
||||
)
|
||||
# Kill failed process
|
||||
try:
|
||||
if process.poll() is None:
|
||||
kill_process_tree(process.pid)
|
||||
else:
|
||||
process.wait(timeout=5)
|
||||
except Exception as e:
|
||||
print(f"CI_OFFLINE: Error cleaning up failed offline process: {e}")
|
||||
|
||||
# Invalidate per-run marker to prevent subsequent tests from using offline
|
||||
if per_run_marker_path and os.path.exists(per_run_marker_path):
|
||||
try:
|
||||
headers = {
|
||||
"Content-Type": "application/json; charset=utf-8",
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
}
|
||||
response = session.get(
|
||||
f"{base_url}/health_generate",
|
||||
headers=headers,
|
||||
timeout=5,
|
||||
)
|
||||
if response.status_code == 200:
|
||||
return process
|
||||
except requests.RequestException:
|
||||
pass
|
||||
os.remove(per_run_marker_path)
|
||||
print("CI_OFFLINE: Invalidated per-run marker due to offline failure")
|
||||
except Exception as e:
|
||||
print(f"CI_OFFLINE: Failed to remove per-run marker: {e}")
|
||||
|
||||
return_code = process.poll()
|
||||
if return_code is not None:
|
||||
raise Exception(
|
||||
f"Server unexpectedly exits ({return_code=}). Usually there will be error logs describing the cause far above this line."
|
||||
)
|
||||
# Retry with online mode
|
||||
env["HF_HUB_OFFLINE"] = "0"
|
||||
process = _launch_server_process(command, env, return_stdout_stderr, model)
|
||||
success, error_msg = _wait_for_server_health(
|
||||
process, base_url, api_key, timeout
|
||||
)
|
||||
|
||||
time.sleep(10)
|
||||
if success:
|
||||
print("CI_OFFLINE: Online retry succeeded")
|
||||
return process
|
||||
|
||||
kill_process_tree(process.pid)
|
||||
raise TimeoutError("Server failed to start within the timeout period.")
|
||||
# Online retry also failed
|
||||
try:
|
||||
kill_process_tree(process.pid)
|
||||
except Exception as e:
|
||||
print(f"CI_OFFLINE: Error killing process after online retry failure: {e}")
|
||||
|
||||
if "exited" in error_msg:
|
||||
raise Exception(error_msg + ". Check server logs for errors.")
|
||||
raise TimeoutError(error_msg)
|
||||
|
||||
# First attempt succeeded or offline was not enabled
|
||||
if success:
|
||||
return process
|
||||
|
||||
# First attempt failed and offline was not enabled
|
||||
try:
|
||||
kill_process_tree(process.pid)
|
||||
except Exception as e:
|
||||
print(f"CI_OFFLINE: Error killing process after first attempt failure: {e}")
|
||||
|
||||
if "exited" in error_msg:
|
||||
raise Exception(error_msg + ". Check server logs for errors.")
|
||||
raise TimeoutError(error_msg)
|
||||
|
||||
|
||||
def popen_launch_pd_server(
|
||||
|
||||
@@ -11,4 +11,9 @@ echo ""
|
||||
python3 "${SCRIPT_DIR}/cleanup_hf_cache.py"
|
||||
echo ""
|
||||
|
||||
# Pre-validate cached models and write markers for offline mode
|
||||
# This allows tests to run with HF_HUB_OFFLINE=1 for models that are fully cached
|
||||
python3 "${SCRIPT_DIR}/prevalidate_cached_models.py"
|
||||
echo ""
|
||||
|
||||
echo "CI runner preparation complete!"
|
||||
|
||||
407
scripts/ci/prevalidate_cached_models.py
Executable file
407
scripts/ci/prevalidate_cached_models.py
Executable file
@@ -0,0 +1,407 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Pre-validate all cached HuggingFace models to provide detailed feedback.
|
||||
|
||||
This script runs once during CI initialization (in prepare_runner.sh) to:
|
||||
1. Scan snapshots in ~/.cache/huggingface/hub/ (with time/quantity limits)
|
||||
2. Validate completeness (config/tokenizer/weights)
|
||||
3. Output detailed failure reasons for debugging
|
||||
|
||||
NOTE: This script no longer writes shared validation markers. Each test run
|
||||
independently validates its cache using per-run markers to avoid cross-runner
|
||||
cache state pollution.
|
||||
"""
|
||||
|
||||
import glob
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
# Add python directory to path to import sglang modules
|
||||
REPO_ROOT = Path(__file__).parent.parent.parent
|
||||
sys.path.insert(0, str(REPO_ROOT / "python"))
|
||||
|
||||
from sglang.srt.model_loader.ci_weight_validation import ( # noqa: E402
|
||||
_validate_diffusion_model,
|
||||
validate_cache_with_detailed_reason,
|
||||
)
|
||||
|
||||
# Limits to avoid spending too much time on validation
|
||||
MAX_VALIDATION_TIME_SECONDS = 300 # Max 5 minutes total
|
||||
|
||||
|
||||
def find_all_hf_snapshots():
|
||||
"""
|
||||
Find all HuggingFace snapshots in cache.
|
||||
|
||||
Returns:
|
||||
List of (model_name, snapshot_dir) tuples, sorted by mtime (newest first)
|
||||
"""
|
||||
hf_home = os.environ.get("HF_HOME", os.path.expanduser("~/.cache/huggingface"))
|
||||
hub_dir = os.path.join(hf_home, "hub")
|
||||
|
||||
if not os.path.isdir(hub_dir):
|
||||
print(f"HF hub directory not found: {hub_dir}")
|
||||
return []
|
||||
|
||||
snapshots = []
|
||||
|
||||
# Pattern: models--org--model/snapshots/hash
|
||||
for model_dir in glob.glob(os.path.join(hub_dir, "models--*")):
|
||||
# Extract model name from directory (models--org--model -> org/model)
|
||||
dir_name = os.path.basename(model_dir)
|
||||
if not dir_name.startswith("models--"):
|
||||
continue
|
||||
|
||||
# models--meta-llama--Llama-2-7b-hf -> meta-llama/Llama-2-7b-hf
|
||||
# Handle multi-part names: models--a--b--c -> a/b-c (join parts 1+ with /)
|
||||
parts = dir_name.split("--")
|
||||
if len(parts) < 3 or parts[0] != "models":
|
||||
# Invalid format, skip
|
||||
continue
|
||||
# Standard format: models--org--repo -> org/repo
|
||||
# Extended format: models--org--repo--extra -> org/repo-extra (join with -)
|
||||
model_name = parts[1] + "/" + "-".join(parts[2:])
|
||||
|
||||
snapshots_dir = os.path.join(model_dir, "snapshots")
|
||||
if not os.path.isdir(snapshots_dir):
|
||||
continue
|
||||
|
||||
# Find all snapshot hashes
|
||||
for snapshot_hash_dir in os.listdir(snapshots_dir):
|
||||
snapshot_path = os.path.join(snapshots_dir, snapshot_hash_dir)
|
||||
if os.path.isdir(snapshot_path):
|
||||
try:
|
||||
mtime = os.path.getmtime(snapshot_path)
|
||||
snapshots.append((model_name, snapshot_path, mtime))
|
||||
except OSError:
|
||||
continue
|
||||
|
||||
# Sort by mtime (newest first) - prioritize recently used models
|
||||
snapshots.sort(key=lambda x: x[2], reverse=True)
|
||||
|
||||
# Return without mtime
|
||||
return [(name, path) for name, path, _ in snapshots]
|
||||
|
||||
|
||||
def is_transformers_text_model(snapshot_dir):
|
||||
"""
|
||||
Check if a snapshot is a transformers text model.
|
||||
|
||||
Only excludes (returns False) for models with STRONG evidence of being
|
||||
diffusers/generation pipelines. Uses conservative heuristics to avoid
|
||||
false negatives on multimodal LLMs with tokenizers.
|
||||
|
||||
Args:
|
||||
snapshot_dir: Path to snapshot directory
|
||||
|
||||
Returns:
|
||||
True if this looks like a transformers text model, False otherwise (N/A)
|
||||
"""
|
||||
# Check for diffusers pipeline markers (strong evidence)
|
||||
diffusers_markers = [
|
||||
"model_index.json", # Diffusers pipeline config
|
||||
"scheduler", # Scheduler directory (diffusers)
|
||||
]
|
||||
if any(
|
||||
os.path.exists(os.path.join(snapshot_dir, marker))
|
||||
for marker in diffusers_markers
|
||||
):
|
||||
return False
|
||||
|
||||
config_path = os.path.join(snapshot_dir, "config.json")
|
||||
if not os.path.exists(config_path):
|
||||
# No config.json - likely not a transformers model
|
||||
return False
|
||||
|
||||
try:
|
||||
with open(config_path, "r", encoding="utf-8") as f:
|
||||
config = json.load(f)
|
||||
|
||||
# Check for explicit diffusers/generation model types (conservative keywords)
|
||||
model_type = config.get("_class_name") or config.get("model_type")
|
||||
if model_type:
|
||||
model_type_lower = str(model_type).lower()
|
||||
# Only exclude clear diffusion/generation models
|
||||
if any(
|
||||
keyword in model_type_lower
|
||||
for keyword in [
|
||||
"diffusion",
|
||||
"unet",
|
||||
"vae",
|
||||
"controlnet",
|
||||
"stable-diffusion",
|
||||
"latent-diffusion",
|
||||
]
|
||||
):
|
||||
return False
|
||||
|
||||
# Check architectures for explicit generation/diffusion classes
|
||||
architectures = config.get("architectures", [])
|
||||
if architectures:
|
||||
arch_str = " ".join(architectures).lower()
|
||||
# Conservative: only exclude obvious diffusion/generation architectures
|
||||
# Use word boundaries to avoid false positives (e.g., "dit" in "conditional")
|
||||
for keyword in [
|
||||
"diffusion",
|
||||
"unet2d",
|
||||
"unet3d",
|
||||
"vaedecoder", # More specific than "vae"
|
||||
"vaeencoder",
|
||||
"controlnet",
|
||||
"autoencoder",
|
||||
"ditmodel", # Diffusion Transformer - use more specific pattern
|
||||
"pixart", # PixArt diffusion model
|
||||
]:
|
||||
if keyword in arch_str:
|
||||
return False
|
||||
|
||||
# Check for standalone vision encoder/image processor (no text component)
|
||||
# Only if model name explicitly indicates non-text usage
|
||||
model_name = config.get("_name_or_path", "").lower()
|
||||
|
||||
if any(
|
||||
keyword in model_name
|
||||
for keyword in [
|
||||
"image-edit-", # Pure image editing (e.g., Qwen-Image-Edit)
|
||||
"-image-editing",
|
||||
"dit-", # DiT generation models
|
||||
"pixart-", # PixArt generation models
|
||||
]
|
||||
):
|
||||
# Additional check: does it have tokenizer? If yes, might be multimodal LLM
|
||||
has_tokenizer = any(
|
||||
os.path.exists(os.path.join(snapshot_dir, fname))
|
||||
for fname in ["tokenizer.json", "tokenizer.model", "tiktoken.model"]
|
||||
)
|
||||
if not has_tokenizer:
|
||||
# Image-edit model without tokenizer -> likely pure vision pipeline
|
||||
return False
|
||||
|
||||
# Default: assume it's a transformers text/multimodal model
|
||||
# Even if it lacks tokenizer, let validation report the actual error
|
||||
# (better false positive than false negative for text models)
|
||||
return True
|
||||
|
||||
except (json.JSONDecodeError, OSError, KeyError):
|
||||
# Can't parse config - assume it's transformers and let validation report failure
|
||||
return True
|
||||
|
||||
|
||||
def scan_weight_files(snapshot_dir):
|
||||
"""
|
||||
Scan for weight files in a snapshot.
|
||||
|
||||
Returns:
|
||||
List of weight file paths, or empty list if scan fails
|
||||
"""
|
||||
weight_files = []
|
||||
|
||||
# First, look for index files
|
||||
index_patterns = ["*.safetensors.index.json", "pytorch_model.bin.index.json"]
|
||||
index_files = []
|
||||
for pattern in index_patterns:
|
||||
index_files.extend(glob.glob(os.path.join(snapshot_dir, pattern)))
|
||||
|
||||
# If we have safetensors index, collect shards from it
|
||||
for index_file in index_files:
|
||||
if index_file.endswith(".safetensors.index.json"):
|
||||
try:
|
||||
with open(index_file, "r", encoding="utf-8") as f:
|
||||
index_data = json.load(f)
|
||||
weight_map = index_data.get("weight_map", {})
|
||||
for weight_file in set(weight_map.values()):
|
||||
weight_path = os.path.join(snapshot_dir, weight_file)
|
||||
if os.path.exists(weight_path):
|
||||
weight_files.append(weight_path)
|
||||
except Exception as e:
|
||||
print(
|
||||
f" Warning: Failed to parse index {os.path.basename(index_file)}: {e}"
|
||||
)
|
||||
|
||||
# If no index found or no shards from index, do recursive glob
|
||||
if not weight_files:
|
||||
matched = glob.glob(
|
||||
os.path.join(snapshot_dir, "**/*.safetensors"), recursive=True
|
||||
)
|
||||
MAX_WEIGHT_FILES = 1000
|
||||
if len(matched) > MAX_WEIGHT_FILES:
|
||||
print(
|
||||
f" Warning: Too many safetensors files ({len(matched)} > {MAX_WEIGHT_FILES})"
|
||||
)
|
||||
return []
|
||||
|
||||
for f in matched:
|
||||
if os.path.exists(f): # Filter out broken symlinks
|
||||
weight_files.append(f)
|
||||
|
||||
return weight_files
|
||||
|
||||
|
||||
def validate_snapshot(model_name, snapshot_dir, weight_files, validated_cache):
|
||||
"""
|
||||
Validate a snapshot and return detailed status.
|
||||
|
||||
Uses in-process cache to avoid duplicate validation within the same run.
|
||||
|
||||
Args:
|
||||
model_name: Model identifier
|
||||
snapshot_dir: Path to snapshot directory
|
||||
weight_files: List of weight files to validate
|
||||
validated_cache: Dict to track already-validated snapshots in this run
|
||||
|
||||
Returns:
|
||||
Tuple of (result, reason):
|
||||
- (True, None) if validation passed
|
||||
- (False, reason_str) if validation failed
|
||||
- (None, None) if skipped (already validated in this run)
|
||||
"""
|
||||
# Fast path: check in-process cache first
|
||||
if snapshot_dir in validated_cache:
|
||||
return None, None # Already validated in this run, skip
|
||||
|
||||
try:
|
||||
# Perform validation with detailed reason
|
||||
is_complete, reason = validate_cache_with_detailed_reason(
|
||||
snapshot_dir=snapshot_dir,
|
||||
weight_files=weight_files,
|
||||
model_name_or_path=model_name,
|
||||
)
|
||||
|
||||
# Cache result to avoid re-validation in this run
|
||||
validated_cache[snapshot_dir] = (is_complete, reason)
|
||||
|
||||
return is_complete, reason
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"Validation raised exception: {e}"
|
||||
return False, error_msg
|
||||
|
||||
|
||||
def main():
|
||||
start_time = time.time()
|
||||
|
||||
print("=" * 70)
|
||||
print("CI_OFFLINE: Pre-validating cached HuggingFace models")
|
||||
print("=" * 70)
|
||||
print(f"Max time: {MAX_VALIDATION_TIME_SECONDS}s")
|
||||
print()
|
||||
|
||||
print("Scanning HuggingFace cache for models...")
|
||||
snapshots = find_all_hf_snapshots()
|
||||
|
||||
if not snapshots:
|
||||
print("No cached models found, skipping validation")
|
||||
print("=" * 70)
|
||||
return
|
||||
|
||||
print(f"Found {len(snapshots)} snapshot(s) in cache")
|
||||
print()
|
||||
|
||||
validated_count = 0
|
||||
failed_count = 0
|
||||
skipped_count = 0
|
||||
processed_count = 0
|
||||
|
||||
# In-process cache to avoid re-validating same snapshot in this run
|
||||
validated_cache = {}
|
||||
|
||||
for model_name, snapshot_dir in snapshots:
|
||||
# Check time limit
|
||||
elapsed = time.time() - start_time
|
||||
if elapsed > MAX_VALIDATION_TIME_SECONDS:
|
||||
print()
|
||||
print(
|
||||
f"Time limit reached ({elapsed:.1f}s > {MAX_VALIDATION_TIME_SECONDS}s)"
|
||||
)
|
||||
print(
|
||||
f"Stopping validation, {len(snapshots) - processed_count} snapshots remaining"
|
||||
)
|
||||
break
|
||||
|
||||
snapshot_hash = os.path.basename(snapshot_dir)
|
||||
print(
|
||||
f"[{processed_count + 1}/{len(snapshots)}] {model_name} ({snapshot_hash[:8]}...)"
|
||||
)
|
||||
processed_count += 1
|
||||
|
||||
# Determine model type by checking for model_index.json (diffusers pipeline marker)
|
||||
model_index_path = os.path.join(snapshot_dir, "model_index.json")
|
||||
is_diffusion_model = os.path.exists(model_index_path)
|
||||
|
||||
if is_diffusion_model:
|
||||
# This is a diffusers pipeline - use diffusion validation
|
||||
try:
|
||||
is_valid, reason = _validate_diffusion_model(snapshot_dir)
|
||||
|
||||
if is_valid:
|
||||
print(" PASS (diffusion) - Cache complete & valid")
|
||||
validated_count += 1
|
||||
else:
|
||||
print(f" FAIL (diffusion) - {reason}")
|
||||
failed_count += 1
|
||||
|
||||
except Exception as e:
|
||||
print(f" FAIL (diffusion) - Validation raised exception: {e}")
|
||||
failed_count += 1
|
||||
|
||||
continue
|
||||
|
||||
# Transformers model - use standard validation
|
||||
# First check if this looks like a transformers text model
|
||||
if not is_transformers_text_model(snapshot_dir):
|
||||
# Not a recognized model type, skip
|
||||
print(
|
||||
" SKIP (unknown type) - Not a diffusers pipeline or transformers model"
|
||||
)
|
||||
skipped_count += 1
|
||||
continue
|
||||
|
||||
# Scan weight files
|
||||
weight_files = scan_weight_files(snapshot_dir)
|
||||
|
||||
if not weight_files:
|
||||
print(" SKIP (no weights) - empty or incomplete download")
|
||||
skipped_count += 1
|
||||
continue
|
||||
|
||||
# Validate
|
||||
try:
|
||||
result, reason = validate_snapshot(
|
||||
model_name, snapshot_dir, weight_files, validated_cache
|
||||
)
|
||||
|
||||
if result is True:
|
||||
print(" PASS - Cache complete & valid")
|
||||
validated_count += 1
|
||||
elif result is False:
|
||||
# Print detailed failure reason
|
||||
if reason:
|
||||
print(f" FAIL (incomplete) - {reason}")
|
||||
else:
|
||||
print(" FAIL (incomplete) - cache validation failed")
|
||||
failed_count += 1
|
||||
else: # None (skipped)
|
||||
print(" SKIP (already validated in this run)")
|
||||
skipped_count += 1
|
||||
|
||||
except Exception as e:
|
||||
print(f" FAIL (error) - Validation raised exception: {e}")
|
||||
failed_count += 1
|
||||
|
||||
elapsed_total = time.time() - start_time
|
||||
|
||||
print()
|
||||
print("=" * 70)
|
||||
print(f"Validation summary (completed in {elapsed_total:.1f}s):")
|
||||
print(f" PASS (complete & valid): {validated_count}")
|
||||
print(f" FAIL (incomplete/corrupted): {failed_count}")
|
||||
print(f" SKIP (no weights/duplicate): {skipped_count}")
|
||||
print(f" Total processed: {processed_count}/{len(snapshots)}")
|
||||
print("=" * 70)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user