Refactor: separate CI-specific weight validation into dedicated module (#15216)
Co-authored-by: Kangyan-Zhou <zky314343421@gmail.com>
This commit is contained in:
Alison Shao
629
python/sglang/srt/model_loader/ci_weight_validation.py
Normal file
629
python/sglang/srt/model_loader/ci_weight_validation.py
Normal file
@@ -0,0 +1,629 @@
|
||||
"""
|
||||
CI-specific weight validation and cache cleanup utilities.
|
||||
|
||||
This module contains validation and cleanup logic that is ONLY used in CI environments.
|
||||
These functions handle:
|
||||
- Validating safetensors files for corruption
|
||||
- Checking for missing shards in sharded models
|
||||
- Cleaning up corrupted files (selective or full cache deletion)
|
||||
- Automatic retry logic for corrupted downloads
|
||||
|
||||
For regular users, weight_utils.py provides simple download functionality without
|
||||
the overhead of validation and automatic cleanup. The CI-specific behavior is
|
||||
gated by is_in_ci() checks in weight_utils.py.
|
||||
"""
|
||||
|
||||
import glob as glob_module
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
import safetensors
|
||||
|
||||
from sglang.srt.utils import log_info_on_rank0
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _validate_safetensors_file(file_path: str) -> bool:
|
||||
"""
|
||||
Validate that a safetensors file is readable and not corrupted.
|
||||
|
||||
Args:
|
||||
file_path: Path to the safetensors file
|
||||
|
||||
Returns:
|
||||
True if the file is valid, False if corrupted
|
||||
"""
|
||||
try:
|
||||
# Attempt to open and read the header
|
||||
# This will fail if the file is corrupted or incomplete
|
||||
with safetensors.safe_open(file_path, framework="pt", device="cpu") as f:
|
||||
# Just accessing the keys validates the header is readable
|
||||
_ = list(f.keys())
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"Corrupted safetensors file detected: %s - %s: %s",
|
||||
file_path,
|
||||
type(e).__name__,
|
||||
str(e),
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
def _check_index_files_exist(snapshot_dir: str) -> Tuple[bool, Optional[str]]:
|
||||
"""
|
||||
Check if all files listed in safetensors index files actually exist on disk.
|
||||
|
||||
This catches cases where the snapshot directory exists but files are missing
|
||||
(e.g., due to incomplete downloads or corrupted cache).
|
||||
|
||||
Args:
|
||||
snapshot_dir: Path to the model snapshot directory
|
||||
|
||||
Returns:
|
||||
Tuple of (all_exist, error_message)
|
||||
"""
|
||||
# Find all safetensors index files
|
||||
index_files = [
|
||||
f for f in os.listdir(snapshot_dir) if f.endswith(".safetensors.index.json")
|
||||
]
|
||||
|
||||
if not index_files:
|
||||
# No index files means it's not a sharded model, skip this check
|
||||
return True, None
|
||||
|
||||
for index_file in index_files:
|
||||
index_path = os.path.join(snapshot_dir, index_file)
|
||||
|
||||
# Check if index file is a broken symlink (exists in listing but blob missing)
|
||||
if os.path.islink(index_path) and not os.path.exists(index_path):
|
||||
# Broken symlink - clean it up so download can proceed
|
||||
try:
|
||||
blob_path = os.path.realpath(index_path)
|
||||
os.remove(index_path)
|
||||
logger.warning(
|
||||
"Removed broken index symlink: %s (blob missing)", index_file
|
||||
)
|
||||
# Also try to remove dangling blob reference if it somehow exists
|
||||
if os.path.exists(blob_path):
|
||||
os.remove(blob_path)
|
||||
except Exception as e:
|
||||
logger.error("Failed to remove broken symlink %s: %s", index_file, e)
|
||||
return (
|
||||
False,
|
||||
f"Broken index file symlink: {index_file} (cleaned up, will re-download)",
|
||||
)
|
||||
|
||||
try:
|
||||
with open(index_path) as f:
|
||||
index_data = json.load(f)
|
||||
|
||||
weight_map = index_data.get("weight_map", {})
|
||||
if not weight_map:
|
||||
continue
|
||||
|
||||
# Check that all files in weight_map exist
|
||||
required_files = set(weight_map.values())
|
||||
missing_files = []
|
||||
|
||||
for file_name in required_files:
|
||||
file_path = os.path.join(snapshot_dir, file_name)
|
||||
# Check both existence and that it's not a broken symlink
|
||||
if not os.path.exists(file_path):
|
||||
missing_files.append(file_name)
|
||||
|
||||
if missing_files:
|
||||
return (
|
||||
False,
|
||||
f"Missing {len(missing_files)} file(s) from index {index_file}: {missing_files[:3]}{'...' if len(missing_files) > 3 else ''}",
|
||||
)
|
||||
|
||||
except FileNotFoundError as e:
|
||||
# Index file was listed but can't be read - could be race condition or broken state
|
||||
logger.warning("Failed to read index file %s: %s", index_file, e)
|
||||
return (
|
||||
False,
|
||||
f"Index file {index_file} unreadable (will re-download)",
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning("Failed to read index file %s: %s", index_file, e)
|
||||
continue
|
||||
|
||||
return True, None
|
||||
|
||||
|
||||
def _validate_sharded_model(
|
||||
snapshot_dir: str, weight_files: List[str]
|
||||
) -> Tuple[bool, Optional[str], List[str]]:
|
||||
"""
|
||||
Validate that all model shards are present and not corrupted.
|
||||
|
||||
Args:
|
||||
snapshot_dir: Path to the model snapshot directory
|
||||
weight_files: List of weight file paths
|
||||
|
||||
Returns:
|
||||
Tuple of (is_valid, error_message, corrupted_files)
|
||||
- corrupted_files: List of file paths that are corrupted (for selective cleanup)
|
||||
"""
|
||||
# First, check if all files from the index actually exist
|
||||
# This catches missing files that wouldn't be found by glob
|
||||
index_check_valid, index_error = _check_index_files_exist(snapshot_dir)
|
||||
if not index_check_valid:
|
||||
return False, index_error, []
|
||||
|
||||
# Pattern for sharded files: model-00001-of-00009.safetensors
|
||||
shard_pattern = re.compile(r"(.*?)-(\d+)-of-(\d+)\.(safetensors|bin)")
|
||||
|
||||
# Group files by shard pattern (prefix-*-of-N)
|
||||
shard_groups = {}
|
||||
for f in weight_files:
|
||||
base_name = os.path.basename(f)
|
||||
match = shard_pattern.match(base_name)
|
||||
if match:
|
||||
prefix = match.group(1)
|
||||
total_shards_str = match.group(3)
|
||||
suffix = match.group(4)
|
||||
|
||||
group_key = f"{prefix}-of-{total_shards_str}.{suffix}"
|
||||
if group_key not in shard_groups:
|
||||
shard_groups[group_key] = {
|
||||
"prefix": prefix,
|
||||
"total": int(total_shards_str),
|
||||
"suffix": suffix,
|
||||
"found_shards": [],
|
||||
"files": [],
|
||||
}
|
||||
|
||||
shard_id = int(match.group(2))
|
||||
shard_groups[group_key]["found_shards"].append(shard_id)
|
||||
shard_groups[group_key]["files"].append(f)
|
||||
|
||||
# Track corrupted files for selective cleanup
|
||||
corrupted_files = []
|
||||
|
||||
# Validate each shard group
|
||||
for group_key, group_info in shard_groups.items():
|
||||
total_shards = group_info["total"]
|
||||
found_shards = set(group_info["found_shards"])
|
||||
expected_shards = set(range(1, total_shards + 1))
|
||||
|
||||
# Check for missing shards
|
||||
missing_shards = expected_shards - found_shards
|
||||
if missing_shards:
|
||||
return (
|
||||
False,
|
||||
f"Missing shards in {group_key}: {sorted(missing_shards)}",
|
||||
[],
|
||||
)
|
||||
|
||||
# Validate safetensors files for corruption
|
||||
if group_info["suffix"] == "safetensors":
|
||||
for f in group_info["files"]:
|
||||
if not _validate_safetensors_file(f):
|
||||
corrupted_files.append(f)
|
||||
|
||||
# Check for required index file for safetensors shards
|
||||
if group_info["suffix"] == "safetensors":
|
||||
index_file = os.path.join(
|
||||
snapshot_dir, f"{group_info['prefix']}.safetensors.index.json"
|
||||
)
|
||||
if not os.path.exists(index_file):
|
||||
return (
|
||||
False,
|
||||
f"Missing index file: {os.path.basename(index_file)}",
|
||||
[],
|
||||
)
|
||||
|
||||
if corrupted_files:
|
||||
return (
|
||||
False,
|
||||
f"Corrupted shard files: {[os.path.basename(f) for f in corrupted_files]}",
|
||||
corrupted_files,
|
||||
)
|
||||
|
||||
return True, None, []
|
||||
|
||||
|
||||
def _cleanup_corrupted_files_selective(
|
||||
model_name_or_path: str, corrupted_files: List[str]
|
||||
) -> int:
|
||||
"""
|
||||
Selectively remove corrupted files and their blobs to force re-download.
|
||||
|
||||
This is more efficient than removing the entire model cache as it only
|
||||
re-downloads corrupted files rather than the entire model.
|
||||
|
||||
Args:
|
||||
model_name_or_path: Model identifier
|
||||
corrupted_files: List of corrupted file paths (symlinks in snapshot)
|
||||
|
||||
Returns:
|
||||
Number of files successfully cleaned up
|
||||
"""
|
||||
cleaned_count = 0
|
||||
|
||||
for file_path in corrupted_files:
|
||||
try:
|
||||
# Resolve symlink to get blob path before deleting symlink
|
||||
if os.path.islink(file_path):
|
||||
blob_path = os.path.realpath(file_path)
|
||||
|
||||
# Delete the symlink
|
||||
os.remove(file_path)
|
||||
logger.info(
|
||||
"Removed corrupted symlink: %s", os.path.basename(file_path)
|
||||
)
|
||||
|
||||
# Delete the blob (the actual corrupted data)
|
||||
if os.path.exists(blob_path):
|
||||
os.remove(blob_path)
|
||||
logger.info(
|
||||
"Removed corrupted blob: %s", os.path.basename(blob_path)
|
||||
)
|
||||
|
||||
cleaned_count += 1
|
||||
elif os.path.exists(file_path):
|
||||
# Not a symlink, just delete the file
|
||||
os.remove(file_path)
|
||||
logger.info("Removed corrupted file: %s", os.path.basename(file_path))
|
||||
cleaned_count += 1
|
||||
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Failed to remove corrupted file %s: %s",
|
||||
os.path.basename(file_path),
|
||||
e,
|
||||
)
|
||||
|
||||
if cleaned_count > 0:
|
||||
logger.warning(
|
||||
"Removed %d corrupted file(s) for %s. "
|
||||
"These will be re-downloaded on next load.",
|
||||
cleaned_count,
|
||||
model_name_or_path,
|
||||
)
|
||||
|
||||
return cleaned_count
|
||||
|
||||
|
||||
def _cleanup_corrupted_model_cache(
|
||||
model_name_or_path: str, snapshot_dir: str, reason: str
|
||||
) -> None:
|
||||
"""
|
||||
Remove entire corrupted model cache directory to force a clean re-download.
|
||||
|
||||
This is used when we cannot selectively clean (e.g., missing shards, incomplete
|
||||
downloads with unknown affected files).
|
||||
|
||||
Args:
|
||||
model_name_or_path: Model identifier
|
||||
snapshot_dir: Path to the snapshot directory
|
||||
reason: Reason for cleanup
|
||||
"""
|
||||
# Navigate up to the model root directory: snapshots/hash -> snapshots -> model_root
|
||||
repo_folder = os.path.abspath(os.path.join(snapshot_dir, "..", ".."))
|
||||
|
||||
try:
|
||||
logger.warning(
|
||||
"Removing entire cache for %s at %s. Reason: %s",
|
||||
model_name_or_path,
|
||||
repo_folder,
|
||||
reason,
|
||||
)
|
||||
shutil.rmtree(repo_folder)
|
||||
logger.info("Successfully removed corrupted cache directory")
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Failed to remove corrupted cache directory %s: %s. "
|
||||
"Manual cleanup may be required.",
|
||||
repo_folder,
|
||||
e,
|
||||
)
|
||||
|
||||
|
||||
def ci_validate_and_cleanup_local_snapshot(
|
||||
model_name_or_path: str,
|
||||
found_local_snapshot_dir: str,
|
||||
local_weight_files: List[str],
|
||||
) -> bool:
|
||||
"""
|
||||
CI-specific validation and cleanup for local model snapshots.
|
||||
|
||||
This function validates the local snapshot and performs automatic cleanup
|
||||
if corruption or missing files are detected. This behavior is only appropriate
|
||||
for CI environments where we want automatic recovery.
|
||||
|
||||
Args:
|
||||
model_name_or_path: Model identifier for logging
|
||||
found_local_snapshot_dir: Path to the local snapshot directory
|
||||
local_weight_files: List of weight file paths found in the snapshot
|
||||
|
||||
Returns:
|
||||
True if the snapshot is valid and can be used, False if it was invalid
|
||||
and cleanup was performed (caller should re-download)
|
||||
"""
|
||||
# Check for incomplete files and clean up if found
|
||||
repo_folder = os.path.abspath(os.path.join(found_local_snapshot_dir, "..", ".."))
|
||||
blobs_dir = os.path.join(repo_folder, "blobs")
|
||||
|
||||
# Check for incomplete download markers
|
||||
incomplete_files = []
|
||||
if os.path.isdir(blobs_dir):
|
||||
incomplete_files = glob_module.glob(os.path.join(blobs_dir, "*.incomplete"))
|
||||
|
||||
if incomplete_files:
|
||||
log_info_on_rank0(
|
||||
logger,
|
||||
f"Found {len(incomplete_files)} .incomplete files in {blobs_dir} for "
|
||||
f"{model_name_or_path}. Will clean up and re-download.",
|
||||
)
|
||||
_cleanup_corrupted_model_cache(
|
||||
model_name_or_path,
|
||||
found_local_snapshot_dir,
|
||||
f"Incomplete download detected ({len(incomplete_files)} incomplete files)",
|
||||
)
|
||||
return False
|
||||
|
||||
# Validate sharded models and check for corruption
|
||||
if local_weight_files:
|
||||
is_valid, error_msg, corrupted_files = _validate_sharded_model(
|
||||
found_local_snapshot_dir, local_weight_files
|
||||
)
|
||||
if not is_valid:
|
||||
if corrupted_files:
|
||||
# Selective cleanup: only remove corrupted files
|
||||
log_info_on_rank0(
|
||||
logger,
|
||||
f"Found {len(corrupted_files)} corrupted file(s) for "
|
||||
f"{model_name_or_path}: {error_msg}. "
|
||||
"Will selectively clean and re-download only these files.",
|
||||
)
|
||||
_cleanup_corrupted_files_selective(model_name_or_path, corrupted_files)
|
||||
return False
|
||||
else:
|
||||
# Missing shards (not corruption) - let snapshot_download handle it.
|
||||
# IMPORTANT: Do NOT delete the entire cache here, as other processes
|
||||
# (TP/EP ranks) may already be loading weights from these files.
|
||||
log_info_on_rank0(
|
||||
logger,
|
||||
f"Validation failed for {model_name_or_path}: {error_msg}. "
|
||||
"Will attempt to download missing files.",
|
||||
)
|
||||
return False
|
||||
|
||||
# Also validate single (non-sharded) safetensors files
|
||||
for f in local_weight_files:
|
||||
base_name = os.path.basename(f)
|
||||
# Check if this is a single model file (not sharded)
|
||||
# Include adapter_model.safetensors for LoRA adapters
|
||||
if base_name in [
|
||||
"model.safetensors",
|
||||
"pytorch_model.safetensors",
|
||||
"adapter_model.safetensors",
|
||||
]:
|
||||
if not _validate_safetensors_file(f):
|
||||
log_info_on_rank0(
|
||||
logger,
|
||||
f"Corrupted model file {base_name} for {model_name_or_path}. "
|
||||
"Will selectively clean and re-download this file.",
|
||||
)
|
||||
# Selective cleanup for single file
|
||||
_cleanup_corrupted_files_selective(model_name_or_path, [f])
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def _validate_weights_after_download(
|
||||
hf_folder: str,
|
||||
allow_patterns: List[str],
|
||||
model_name_or_path: str,
|
||||
) -> bool:
|
||||
"""
|
||||
Validate downloaded weight files to catch corruption early.
|
||||
|
||||
This function validates safetensors files after download to catch
|
||||
corruption issues (truncated downloads, network errors, etc.) before
|
||||
model loading fails with cryptic errors. If corruption is found,
|
||||
the corrupted files are automatically cleaned up.
|
||||
|
||||
Args:
|
||||
hf_folder: Path to the downloaded model folder
|
||||
allow_patterns: Patterns used to match weight files
|
||||
model_name_or_path: Model identifier for error messages
|
||||
|
||||
Returns:
|
||||
True if all files are valid, False if corrupted files were found and cleaned up
|
||||
"""
|
||||
# Find all weight files that were downloaded
|
||||
weight_files: List[str] = []
|
||||
for pattern in allow_patterns:
|
||||
weight_files.extend(glob_module.glob(os.path.join(hf_folder, pattern)))
|
||||
|
||||
if not weight_files:
|
||||
return True # No weight files to validate
|
||||
|
||||
# Validate safetensors files
|
||||
corrupted_files = []
|
||||
for f in weight_files:
|
||||
if f.endswith(".safetensors") and os.path.exists(f):
|
||||
if not _validate_safetensors_file(f):
|
||||
corrupted_files.append(os.path.basename(f))
|
||||
|
||||
if corrupted_files:
|
||||
# Clean up corrupted files so next attempt re-downloads them
|
||||
_cleanup_corrupted_files_selective(
|
||||
model_name_or_path,
|
||||
[os.path.join(hf_folder, f) for f in corrupted_files],
|
||||
)
|
||||
log_info_on_rank0(
|
||||
logger,
|
||||
f"Downloaded model files are corrupted for {model_name_or_path}: "
|
||||
f"{corrupted_files}. The corrupted files have been removed. "
|
||||
"Will retry download.",
|
||||
)
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def ci_download_with_validation_and_retry(
|
||||
model_name_or_path: str,
|
||||
allow_patterns: List[str],
|
||||
ignore_patterns,
|
||||
cache_dir: Optional[str],
|
||||
revision: Optional[str],
|
||||
max_retries: int = 3,
|
||||
) -> str:
|
||||
"""
|
||||
CI-specific download with validation and automatic retry on corruption.
|
||||
|
||||
This function handles the download of model weights in CI environments,
|
||||
with automatic validation and retry logic for handling corrupted downloads.
|
||||
|
||||
Args:
|
||||
model_name_or_path: The model name or path
|
||||
allow_patterns: The allowed patterns for weight files
|
||||
ignore_patterns: The patterns to filter out weight files
|
||||
cache_dir: The cache directory to store model weights
|
||||
revision: The revision of the model
|
||||
max_retries: Maximum number of download retries if corruption is detected
|
||||
|
||||
Returns:
|
||||
str: The path to the downloaded model weights
|
||||
|
||||
Raises:
|
||||
RuntimeError: If download fails after max_retries attempts
|
||||
"""
|
||||
# Lazy imports to avoid circular dependencies
|
||||
import huggingface_hub.constants
|
||||
from huggingface_hub import snapshot_download
|
||||
from tqdm.auto import tqdm
|
||||
|
||||
class DisabledTqdm(tqdm):
|
||||
def __init__(self, *args, **kwargs):
|
||||
kwargs["disable"] = True
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
log_info_on_rank0(logger, f"Using model weights format {allow_patterns}")
|
||||
|
||||
# Retry loop for handling corrupted downloads
|
||||
for attempt in range(max_retries):
|
||||
hf_folder = snapshot_download(
|
||||
model_name_or_path,
|
||||
allow_patterns=allow_patterns,
|
||||
ignore_patterns=ignore_patterns,
|
||||
cache_dir=cache_dir,
|
||||
tqdm_class=DisabledTqdm,
|
||||
revision=revision,
|
||||
local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,
|
||||
)
|
||||
|
||||
# Validate downloaded files to catch corruption early
|
||||
is_valid = _validate_weights_after_download(
|
||||
hf_folder, allow_patterns, model_name_or_path
|
||||
)
|
||||
|
||||
if is_valid:
|
||||
return hf_folder
|
||||
|
||||
# Validation failed, corrupted files were cleaned up
|
||||
if attempt < max_retries - 1:
|
||||
log_info_on_rank0(
|
||||
logger,
|
||||
f"Retrying download for {model_name_or_path} "
|
||||
f"(attempt {attempt + 2}/{max_retries})...",
|
||||
)
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"Downloaded model files are still corrupted for "
|
||||
f"{model_name_or_path} after {max_retries} attempts. "
|
||||
"This may indicate a persistent issue with the model files "
|
||||
"on Hugging Face Hub or network problems."
|
||||
)
|
||||
|
||||
# This should never be reached, but just in case
|
||||
return hf_folder
|
||||
|
||||
|
||||
def ci_validate_and_clean_hf_cache(model_path: str) -> None:
|
||||
"""
|
||||
Validate and clean corrupted safetensors files in HF cache before loading.
|
||||
|
||||
This function is needed because HFRunner (used in tests) calls transformers'
|
||||
from_pretrained() directly, which bypasses SGLang's weight validation.
|
||||
Corrupted cached files can cause cryptic errors like "EOF while parsing"
|
||||
from safetensors.
|
||||
|
||||
Only runs in CI to avoid overhead for regular users.
|
||||
|
||||
Args:
|
||||
model_path: Model identifier (e.g., "meta-llama/Llama-2-7b")
|
||||
"""
|
||||
from sglang.utils import is_in_ci
|
||||
|
||||
if not is_in_ci():
|
||||
return
|
||||
|
||||
# Skip for local paths
|
||||
if os.path.isdir(model_path):
|
||||
return
|
||||
|
||||
try:
|
||||
import huggingface_hub.constants
|
||||
|
||||
# Find the HF cache directory for this model
|
||||
cache_dir = huggingface_hub.constants.HF_HUB_CACHE
|
||||
repo_folder = os.path.join(
|
||||
cache_dir,
|
||||
huggingface_hub.constants.REPO_ID_SEPARATOR.join(
|
||||
["models", *model_path.split("/")]
|
||||
),
|
||||
)
|
||||
|
||||
if not os.path.isdir(repo_folder):
|
||||
return
|
||||
|
||||
# Find snapshot directories
|
||||
snapshots_dir = os.path.join(repo_folder, "snapshots")
|
||||
if not os.path.isdir(snapshots_dir):
|
||||
return
|
||||
|
||||
# Check each snapshot for corrupted files
|
||||
corrupted_files = []
|
||||
for snapshot_hash in os.listdir(snapshots_dir):
|
||||
snapshot_dir = os.path.join(snapshots_dir, snapshot_hash)
|
||||
if not os.path.isdir(snapshot_dir):
|
||||
continue
|
||||
|
||||
# Find all safetensors files
|
||||
safetensors_files = glob_module.glob(
|
||||
os.path.join(snapshot_dir, "*.safetensors")
|
||||
)
|
||||
|
||||
for sf_file in safetensors_files:
|
||||
# Skip broken symlinks (os.path.exists returns False for them)
|
||||
if not os.path.exists(sf_file):
|
||||
continue
|
||||
|
||||
if not _validate_safetensors_file(sf_file):
|
||||
corrupted_files.append(sf_file)
|
||||
|
||||
if corrupted_files:
|
||||
logger.warning(
|
||||
"HFRunner: Found %d corrupted safetensors file(s) for %s. "
|
||||
"Removing to force re-download.",
|
||||
len(corrupted_files),
|
||||
model_path,
|
||||
)
|
||||
_cleanup_corrupted_files_selective(model_path, corrupted_files)
|
||||
|
||||
except Exception as e:
|
||||
# Don't fail if validation itself fails - let HF handle it
|
||||
logger.debug("HF cache validation failed (non-fatal): %s", e)
|
||||
@@ -40,12 +40,6 @@ from sglang.srt.layers.quantization.modelopt_quant import (
|
||||
ModelOptFp4Config,
|
||||
ModelOptFp8Config,
|
||||
)
|
||||
from sglang.srt.model_loader.weight_validation import (
|
||||
_cleanup_corrupted_files_selective,
|
||||
_cleanup_corrupted_model_cache,
|
||||
_validate_safetensors_file,
|
||||
_validate_sharded_model,
|
||||
)
|
||||
from sglang.srt.utils import find_local_repo_dir, log_info_on_rank0, print_warning_once
|
||||
from sglang.utils import is_in_ci
|
||||
|
||||
@@ -342,33 +336,6 @@ def _find_local_hf_snapshot_dir_unlocked(
|
||||
if not os.path.isdir(found_local_snapshot_dir):
|
||||
return None
|
||||
|
||||
# Only perform cache validation and cleanup in CI to avoid
|
||||
# unnecessary overhead for regular users
|
||||
if is_in_ci():
|
||||
# Check for incomplete files and clean up if found
|
||||
repo_folder = os.path.abspath(
|
||||
os.path.join(found_local_snapshot_dir, "..", "..")
|
||||
)
|
||||
blobs_dir = os.path.join(repo_folder, "blobs")
|
||||
|
||||
# Check for incomplete download markers
|
||||
incomplete_files = []
|
||||
if os.path.isdir(blobs_dir):
|
||||
incomplete_files = glob.glob(os.path.join(blobs_dir, "*.incomplete"))
|
||||
|
||||
if incomplete_files:
|
||||
log_info_on_rank0(
|
||||
logger,
|
||||
f"Found {len(incomplete_files)} .incomplete files in {blobs_dir} for "
|
||||
f"{model_name_or_path}. Will clean up and re-download.",
|
||||
)
|
||||
_cleanup_corrupted_model_cache(
|
||||
model_name_or_path,
|
||||
found_local_snapshot_dir,
|
||||
f"Incomplete download detected ({len(incomplete_files)} incomplete files)",
|
||||
)
|
||||
return None
|
||||
|
||||
local_weight_files: List[str] = []
|
||||
try:
|
||||
for pattern in allow_patterns:
|
||||
@@ -387,60 +354,18 @@ def _find_local_hf_snapshot_dir_unlocked(
|
||||
)
|
||||
local_weight_files = []
|
||||
|
||||
# Only perform cache validation and cleanup in CI
|
||||
if is_in_ci():
|
||||
# Validate sharded models and check for corruption
|
||||
if local_weight_files:
|
||||
is_valid, error_msg, corrupted_files = _validate_sharded_model(
|
||||
found_local_snapshot_dir, local_weight_files
|
||||
)
|
||||
if not is_valid:
|
||||
if corrupted_files:
|
||||
# Selective cleanup: only remove corrupted files
|
||||
log_info_on_rank0(
|
||||
logger,
|
||||
f"Found {len(corrupted_files)} corrupted file(s) for "
|
||||
f"{model_name_or_path}: {error_msg}. "
|
||||
"Will selectively clean and re-download only these files.",
|
||||
)
|
||||
_cleanup_corrupted_files_selective(
|
||||
model_name_or_path, corrupted_files
|
||||
)
|
||||
return None
|
||||
else:
|
||||
# Missing shards (not corruption) - let snapshot_download handle it.
|
||||
# IMPORTANT: Do NOT delete the entire cache here, as other processes
|
||||
# (TP/EP ranks) may already be loading weights from these files.
|
||||
# Deleting the cache while other processes are using it causes
|
||||
# FileNotFoundError race conditions. Instead, just return None
|
||||
# to trigger a download - snapshot_download will only fetch
|
||||
# missing files without disturbing existing ones.
|
||||
log_info_on_rank0(
|
||||
logger,
|
||||
f"Validation failed for {model_name_or_path}: {error_msg}. "
|
||||
"Will attempt to download missing files.",
|
||||
)
|
||||
return None
|
||||
# Only perform cache validation and cleanup in CI to avoid
|
||||
# unnecessary overhead for regular users
|
||||
if is_in_ci() and local_weight_files:
|
||||
from sglang.srt.model_loader.ci_weight_validation import (
|
||||
ci_validate_and_cleanup_local_snapshot,
|
||||
)
|
||||
|
||||
# Also validate single (non-sharded) safetensors files
|
||||
for f in local_weight_files:
|
||||
base_name = os.path.basename(f)
|
||||
# Check if this is a single model file (not sharded)
|
||||
# Include adapter_model.safetensors for LoRA adapters
|
||||
if base_name in [
|
||||
"model.safetensors",
|
||||
"pytorch_model.safetensors",
|
||||
"adapter_model.safetensors",
|
||||
]:
|
||||
if not _validate_safetensors_file(f):
|
||||
log_info_on_rank0(
|
||||
logger,
|
||||
f"Corrupted model file {base_name} for {model_name_or_path}. "
|
||||
"Will selectively clean and re-download this file.",
|
||||
)
|
||||
# Selective cleanup for single file
|
||||
_cleanup_corrupted_files_selective(model_name_or_path, [f])
|
||||
return None
|
||||
is_valid = ci_validate_and_cleanup_local_snapshot(
|
||||
model_name_or_path, found_local_snapshot_dir, local_weight_files
|
||||
)
|
||||
if not is_valid:
|
||||
return None
|
||||
|
||||
if len(local_weight_files) > 0:
|
||||
log_info_on_rank0(
|
||||
@@ -458,83 +383,6 @@ def _find_local_hf_snapshot_dir_unlocked(
|
||||
return None
|
||||
|
||||
|
||||
def find_local_hf_snapshot_dir(
|
||||
model_name_or_path: str,
|
||||
cache_dir: Optional[str],
|
||||
allow_patterns: List[str],
|
||||
revision: Optional[str] = None,
|
||||
) -> Optional[str]:
|
||||
"""If the weights are already local, skip downloading and returns the path.
|
||||
|
||||
This function acquires a lock to prevent race conditions during validation
|
||||
and cleanup. For use within download_weights_from_hf, use
|
||||
_find_local_hf_snapshot_dir_unlocked instead with an external lock.
|
||||
"""
|
||||
# For local paths, no locking needed
|
||||
if os.path.isdir(model_name_or_path):
|
||||
return None
|
||||
|
||||
# Use file lock to prevent multiple processes (TP ranks) from
|
||||
# validating and cleaning up the same model cache simultaneously.
|
||||
with get_lock(model_name_or_path, cache_dir):
|
||||
return _find_local_hf_snapshot_dir_unlocked(
|
||||
model_name_or_path, cache_dir, allow_patterns, revision
|
||||
)
|
||||
|
||||
|
||||
def _validate_weights_after_download(
|
||||
hf_folder: str,
|
||||
allow_patterns: List[str],
|
||||
model_name_or_path: str,
|
||||
) -> bool:
|
||||
"""Validate downloaded weight files to catch corruption early.
|
||||
|
||||
This function validates safetensors files after download to catch
|
||||
corruption issues (truncated downloads, network errors, etc.) before
|
||||
model loading fails with cryptic errors.
|
||||
|
||||
Args:
|
||||
hf_folder: Path to the downloaded model folder
|
||||
allow_patterns: Patterns used to match weight files
|
||||
model_name_or_path: Model identifier for error messages
|
||||
|
||||
Returns:
|
||||
True if all files are valid, False if corrupted files were found and cleaned up
|
||||
"""
|
||||
import glob as glob_module
|
||||
|
||||
# Find all weight files that were downloaded
|
||||
weight_files: List[str] = []
|
||||
for pattern in allow_patterns:
|
||||
weight_files.extend(glob_module.glob(os.path.join(hf_folder, pattern)))
|
||||
|
||||
if not weight_files:
|
||||
return True # No weight files to validate
|
||||
|
||||
# Validate safetensors files
|
||||
corrupted_files = []
|
||||
for f in weight_files:
|
||||
if f.endswith(".safetensors") and os.path.exists(f):
|
||||
if not _validate_safetensors_file(f):
|
||||
corrupted_files.append(os.path.basename(f))
|
||||
|
||||
if corrupted_files:
|
||||
# Clean up corrupted files so next attempt re-downloads them
|
||||
_cleanup_corrupted_files_selective(
|
||||
model_name_or_path,
|
||||
[os.path.join(hf_folder, f) for f in corrupted_files],
|
||||
)
|
||||
log_info_on_rank0(
|
||||
logger,
|
||||
f"Downloaded model files are corrupted for {model_name_or_path}: "
|
||||
f"{corrupted_files}. The corrupted files have been removed. "
|
||||
"Will retry download.",
|
||||
)
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def download_weights_from_hf(
|
||||
model_name_or_path: str,
|
||||
cache_dir: Optional[str],
|
||||
@@ -595,49 +443,23 @@ def download_weights_from_hf(
|
||||
allow_patterns = [pattern]
|
||||
break
|
||||
|
||||
log_info_on_rank0(logger, f"Using model weights format {allow_patterns}")
|
||||
|
||||
# Only perform validation and retry in CI to avoid overhead for regular users
|
||||
if is_in_ci():
|
||||
# Retry loop for handling corrupted downloads
|
||||
for attempt in range(max_retries):
|
||||
hf_folder = snapshot_download(
|
||||
model_name_or_path,
|
||||
allow_patterns=allow_patterns,
|
||||
ignore_patterns=ignore_patterns,
|
||||
cache_dir=cache_dir,
|
||||
tqdm_class=DisabledTqdm,
|
||||
revision=revision,
|
||||
local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,
|
||||
)
|
||||
from sglang.srt.model_loader.ci_weight_validation import (
|
||||
ci_download_with_validation_and_retry,
|
||||
)
|
||||
|
||||
# Validate downloaded files to catch corruption early
|
||||
is_valid = _validate_weights_after_download(
|
||||
hf_folder, allow_patterns, model_name_or_path
|
||||
)
|
||||
|
||||
if is_valid:
|
||||
return hf_folder
|
||||
|
||||
# Validation failed, corrupted files were cleaned up
|
||||
if attempt < max_retries - 1:
|
||||
log_info_on_rank0(
|
||||
logger,
|
||||
f"Retrying download for {model_name_or_path} "
|
||||
f"(attempt {attempt + 2}/{max_retries})...",
|
||||
)
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"Downloaded model files are still corrupted for "
|
||||
f"{model_name_or_path} after {max_retries} attempts. "
|
||||
"This may indicate a persistent issue with the model files "
|
||||
"on Hugging Face Hub or network problems."
|
||||
)
|
||||
|
||||
# This should never be reached, but just in case
|
||||
return hf_folder
|
||||
return ci_download_with_validation_and_retry(
|
||||
model_name_or_path=model_name_or_path,
|
||||
allow_patterns=allow_patterns,
|
||||
ignore_patterns=ignore_patterns,
|
||||
cache_dir=cache_dir,
|
||||
revision=revision,
|
||||
max_retries=max_retries,
|
||||
)
|
||||
else:
|
||||
# Simple download without validation for non-CI environments
|
||||
log_info_on_rank0(logger, f"Using model weights format {allow_patterns}")
|
||||
hf_folder = snapshot_download(
|
||||
model_name_or_path,
|
||||
allow_patterns=allow_patterns,
|
||||
|
||||
@@ -1,309 +0,0 @@
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
import safetensors
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _validate_safetensors_file(file_path: str) -> bool:
|
||||
"""
|
||||
Validate that a safetensors file is readable and not corrupted.
|
||||
|
||||
Args:
|
||||
file_path: Path to the safetensors file
|
||||
|
||||
Returns:
|
||||
True if the file is valid, False if corrupted
|
||||
"""
|
||||
try:
|
||||
# Attempt to open and read the header
|
||||
# This will fail if the file is corrupted or incomplete
|
||||
with safetensors.safe_open(file_path, framework="pt", device="cpu") as f:
|
||||
# Just accessing the keys validates the header is readable
|
||||
_ = list(f.keys())
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"Corrupted safetensors file detected: %s - %s: %s",
|
||||
file_path,
|
||||
type(e).__name__,
|
||||
str(e),
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
def _check_index_files_exist(snapshot_dir: str) -> Tuple[bool, Optional[str]]:
|
||||
"""
|
||||
Check if all files listed in safetensors index files actually exist on disk.
|
||||
|
||||
This catches cases where the snapshot directory exists but files are missing
|
||||
(e.g., due to incomplete downloads or corrupted cache).
|
||||
|
||||
Args:
|
||||
snapshot_dir: Path to the model snapshot directory
|
||||
|
||||
Returns:
|
||||
Tuple of (all_exist, error_message)
|
||||
"""
|
||||
# Find all safetensors index files
|
||||
index_files = [
|
||||
f for f in os.listdir(snapshot_dir) if f.endswith(".safetensors.index.json")
|
||||
]
|
||||
|
||||
if not index_files:
|
||||
# No index files means it's not a sharded model, skip this check
|
||||
return True, None
|
||||
|
||||
for index_file in index_files:
|
||||
index_path = os.path.join(snapshot_dir, index_file)
|
||||
|
||||
# Check if index file is a broken symlink (exists in listing but blob missing)
|
||||
if os.path.islink(index_path) and not os.path.exists(index_path):
|
||||
# Broken symlink - clean it up so download can proceed
|
||||
try:
|
||||
blob_path = os.path.realpath(index_path)
|
||||
os.remove(index_path)
|
||||
logger.warning(
|
||||
"Removed broken index symlink: %s (blob missing)", index_file
|
||||
)
|
||||
# Also try to remove dangling blob reference if it somehow exists
|
||||
if os.path.exists(blob_path):
|
||||
os.remove(blob_path)
|
||||
except Exception as e:
|
||||
logger.error("Failed to remove broken symlink %s: %s", index_file, e)
|
||||
return (
|
||||
False,
|
||||
f"Broken index file symlink: {index_file} (cleaned up, will re-download)",
|
||||
)
|
||||
|
||||
try:
|
||||
with open(index_path) as f:
|
||||
index_data = json.load(f)
|
||||
|
||||
weight_map = index_data.get("weight_map", {})
|
||||
if not weight_map:
|
||||
continue
|
||||
|
||||
# Check that all files in weight_map exist
|
||||
required_files = set(weight_map.values())
|
||||
missing_files = []
|
||||
|
||||
for file_name in required_files:
|
||||
file_path = os.path.join(snapshot_dir, file_name)
|
||||
# Check both existence and that it's not a broken symlink
|
||||
if not os.path.exists(file_path):
|
||||
missing_files.append(file_name)
|
||||
|
||||
if missing_files:
|
||||
return (
|
||||
False,
|
||||
f"Missing {len(missing_files)} file(s) from index {index_file}: {missing_files[:3]}{'...' if len(missing_files) > 3 else ''}",
|
||||
)
|
||||
|
||||
except FileNotFoundError as e:
|
||||
# Index file was listed but can't be read - could be race condition or broken state
|
||||
logger.warning("Failed to read index file %s: %s", index_file, e)
|
||||
return (
|
||||
False,
|
||||
f"Index file {index_file} unreadable (will re-download)",
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning("Failed to read index file %s: %s", index_file, e)
|
||||
continue
|
||||
|
||||
return True, None
|
||||
|
||||
|
||||
def _validate_sharded_model(
|
||||
snapshot_dir: str, weight_files: List[str]
|
||||
) -> Tuple[bool, Optional[str], List[str]]:
|
||||
"""
|
||||
Validate that all model shards are present and not corrupted.
|
||||
|
||||
Args:
|
||||
snapshot_dir: Path to the model snapshot directory
|
||||
weight_files: List of weight file paths
|
||||
|
||||
Returns:
|
||||
Tuple of (is_valid, error_message, corrupted_files)
|
||||
- corrupted_files: List of file paths that are corrupted (for selective cleanup)
|
||||
"""
|
||||
# First, check if all files from the index actually exist
|
||||
# This catches missing files that wouldn't be found by glob
|
||||
index_check_valid, index_error = _check_index_files_exist(snapshot_dir)
|
||||
if not index_check_valid:
|
||||
return False, index_error, []
|
||||
|
||||
# Pattern for sharded files: model-00001-of-00009.safetensors
|
||||
shard_pattern = re.compile(r"(.*?)-(\d+)-of-(\d+)\.(safetensors|bin)")
|
||||
|
||||
# Group files by shard pattern (prefix-*-of-N)
|
||||
shard_groups = {}
|
||||
for f in weight_files:
|
||||
base_name = os.path.basename(f)
|
||||
match = shard_pattern.match(base_name)
|
||||
if match:
|
||||
prefix = match.group(1)
|
||||
total_shards_str = match.group(3)
|
||||
suffix = match.group(4)
|
||||
|
||||
group_key = f"{prefix}-of-{total_shards_str}.{suffix}"
|
||||
if group_key not in shard_groups:
|
||||
shard_groups[group_key] = {
|
||||
"prefix": prefix,
|
||||
"total": int(total_shards_str),
|
||||
"suffix": suffix,
|
||||
"found_shards": [],
|
||||
"files": [],
|
||||
}
|
||||
|
||||
shard_id = int(match.group(2))
|
||||
shard_groups[group_key]["found_shards"].append(shard_id)
|
||||
shard_groups[group_key]["files"].append(f)
|
||||
|
||||
# Track corrupted files for selective cleanup
|
||||
corrupted_files = []
|
||||
|
||||
# Validate each shard group
|
||||
for group_key, group_info in shard_groups.items():
|
||||
total_shards = group_info["total"]
|
||||
found_shards = set(group_info["found_shards"])
|
||||
expected_shards = set(range(1, total_shards + 1))
|
||||
|
||||
# Check for missing shards
|
||||
missing_shards = expected_shards - found_shards
|
||||
if missing_shards:
|
||||
return (
|
||||
False,
|
||||
f"Missing shards in {group_key}: {sorted(missing_shards)}",
|
||||
[],
|
||||
)
|
||||
|
||||
# Validate safetensors files for corruption
|
||||
if group_info["suffix"] == "safetensors":
|
||||
for f in group_info["files"]:
|
||||
if not _validate_safetensors_file(f):
|
||||
corrupted_files.append(f)
|
||||
|
||||
# Check for required index file for safetensors shards
|
||||
if group_info["suffix"] == "safetensors":
|
||||
index_file = os.path.join(
|
||||
snapshot_dir, f"{group_info['prefix']}.safetensors.index.json"
|
||||
)
|
||||
if not os.path.exists(index_file):
|
||||
return (
|
||||
False,
|
||||
f"Missing index file: {os.path.basename(index_file)}",
|
||||
[],
|
||||
)
|
||||
|
||||
if corrupted_files:
|
||||
return (
|
||||
False,
|
||||
f"Corrupted shard files: {[os.path.basename(f) for f in corrupted_files]}",
|
||||
corrupted_files,
|
||||
)
|
||||
|
||||
return True, None, []
|
||||
|
||||
|
||||
def _cleanup_corrupted_files_selective(
|
||||
model_name_or_path: str, corrupted_files: List[str]
|
||||
) -> int:
|
||||
"""
|
||||
Selectively remove corrupted files and their blobs to force re-download.
|
||||
|
||||
This is more efficient than removing the entire model cache as it only
|
||||
re-downloads corrupted files rather than the entire model.
|
||||
|
||||
Args:
|
||||
model_name_or_path: Model identifier
|
||||
corrupted_files: List of corrupted file paths (symlinks in snapshot)
|
||||
|
||||
Returns:
|
||||
Number of files successfully cleaned up
|
||||
"""
|
||||
cleaned_count = 0
|
||||
|
||||
for file_path in corrupted_files:
|
||||
try:
|
||||
# Resolve symlink to get blob path before deleting symlink
|
||||
if os.path.islink(file_path):
|
||||
blob_path = os.path.realpath(file_path)
|
||||
|
||||
# Delete the symlink
|
||||
os.remove(file_path)
|
||||
logger.info(
|
||||
"Removed corrupted symlink: %s", os.path.basename(file_path)
|
||||
)
|
||||
|
||||
# Delete the blob (the actual corrupted data)
|
||||
if os.path.exists(blob_path):
|
||||
os.remove(blob_path)
|
||||
logger.info(
|
||||
"Removed corrupted blob: %s", os.path.basename(blob_path)
|
||||
)
|
||||
|
||||
cleaned_count += 1
|
||||
elif os.path.exists(file_path):
|
||||
# Not a symlink, just delete the file
|
||||
os.remove(file_path)
|
||||
logger.info("Removed corrupted file: %s", os.path.basename(file_path))
|
||||
cleaned_count += 1
|
||||
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Failed to remove corrupted file %s: %s",
|
||||
os.path.basename(file_path),
|
||||
e,
|
||||
)
|
||||
|
||||
if cleaned_count > 0:
|
||||
logger.warning(
|
||||
"Removed %d corrupted file(s) for %s. "
|
||||
"These will be re-downloaded on next load.",
|
||||
cleaned_count,
|
||||
model_name_or_path,
|
||||
)
|
||||
|
||||
return cleaned_count
|
||||
|
||||
|
||||
def _cleanup_corrupted_model_cache(
|
||||
model_name_or_path: str, snapshot_dir: str, reason: str
|
||||
) -> None:
|
||||
"""
|
||||
Remove entire corrupted model cache directory to force a clean re-download.
|
||||
|
||||
This is used when we cannot selectively clean (e.g., missing shards, incomplete
|
||||
downloads with unknown affected files).
|
||||
|
||||
Args:
|
||||
model_name_or_path: Model identifier
|
||||
snapshot_dir: Path to the snapshot directory
|
||||
reason: Reason for cleanup
|
||||
"""
|
||||
# Navigate up to the model root directory: snapshots/hash -> snapshots -> model_root
|
||||
repo_folder = os.path.abspath(os.path.join(snapshot_dir, "..", ".."))
|
||||
|
||||
try:
|
||||
logger.warning(
|
||||
"Removing entire cache for %s at %s. Reason: %s",
|
||||
model_name_or_path,
|
||||
repo_folder,
|
||||
reason,
|
||||
)
|
||||
shutil.rmtree(repo_folder)
|
||||
logger.info("Successfully removed corrupted cache directory")
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Failed to remove corrupted cache directory %s: %s. "
|
||||
"Manual cleanup may be required.",
|
||||
repo_folder,
|
||||
e,
|
||||
)
|
||||
@@ -31,6 +31,7 @@ from transformers import (
|
||||
)
|
||||
|
||||
from sglang.srt.entrypoints.engine import Engine
|
||||
from sglang.srt.model_loader.ci_weight_validation import ci_validate_and_clean_hf_cache
|
||||
from sglang.srt.utils import is_npu, load_image
|
||||
from sglang.srt.utils.hf_transformers_utils import get_tokenizer
|
||||
from sglang.test.test_utils import DEFAULT_PORT_FOR_SRT_TEST_RUNNER, calculate_rouge_l
|
||||
@@ -251,6 +252,10 @@ class HFRunner:
|
||||
# Apply model-specific patches
|
||||
monkey_patch_gemma2_sdpa()
|
||||
|
||||
# Validate and clean corrupted files in HF cache (CI only)
|
||||
# This is needed because HFRunner bypasses SGLang's weight validation
|
||||
ci_validate_and_clean_hf_cache(model_path)
|
||||
|
||||
# Load the model and tokenizer
|
||||
if self.model_type == "generation":
|
||||
config = AutoConfig.from_pretrained(
|
||||
|
||||
@@ -11,7 +11,7 @@ import struct
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
from sglang.srt.model_loader.weight_validation import (
|
||||
from sglang.srt.model_loader.ci_weight_validation import (
|
||||
_check_index_files_exist,
|
||||
_validate_sharded_model,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user