minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

fix: race condition between validation and download locks (#14761)

Browse Source
This commit is contained in:
Alison Shao
2025-12-09 20:36:54 -08:00
committed by GitHub
parent 4285e99da7
commit 01835998e1
1 changed files with 216 additions and 160 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/model_loader/weight_utils.py
+87 -31
View File
@@ -263,13 +263,19 @@ def get_quant_config(
return quant_cls.from_config(config)
def find_local_hf_snapshot_dir(
def _find_local_hf_snapshot_dir_unlocked(
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."""
"""Find local HF snapshot directory without locking.
IMPORTANT: Caller MUST hold the model lock before calling this function
to prevent race conditions during validation and cleanup.
If the weights are already local, skip downloading and returns the path.
"""
if os.path.isdir(model_name_or_path):
return None
@@ -315,20 +321,13 @@ def find_local_hf_snapshot_dir(
if found_local_snapshot_dir is None:
return None
# Use file lock to prevent multiple processes (TP ranks) from
# validating and cleaning up the same model cache simultaneously.
# This prevents race conditions where multiple ranks detect corruption
# and try to delete the same files at the same time.
with get_lock(model_name_or_path, cache_dir, suffix="-validation"):
# Re-check if snapshot dir still exists after acquiring lock
# (another process may have already cleaned it up)
# Check if snapshot dir exists (might have been cleaned by another process
# before we acquired the lock)
if not os.path.isdir(found_local_snapshot_dir):
return None
# Check for incomplete files and clean up if found
repo_folder = os.path.abspath(
os.path.join(found_local_snapshot_dir, "..", "..")
)
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
@@ -352,9 +351,7 @@ def find_local_hf_snapshot_dir(
local_weight_files: List[str] = []
try:
for pattern in allow_patterns:
matched_files = glob.glob(
os.path.join(found_local_snapshot_dir, pattern)
)
matched_files = glob.glob(os.path.join(found_local_snapshot_dir, pattern))
for f in matched_files:
# os.path.exists returns False for broken symlinks.
if not os.path.exists(f):
@@ -383,9 +380,7 @@ def find_local_hf_snapshot_dir(
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
)
_cleanup_corrupted_files_selective(model_name_or_path, corrupted_files)
return None
else:
# Cannot selectively clean (e.g., missing shards) - remove entire cache
@@ -435,11 +430,35 @@ def find_local_hf_snapshot_dir(
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,
) -> None:
) -> bool:
"""Validate downloaded weight files to catch corruption early.
This function validates safetensors files after download to catch
@@ -451,8 +470,8 @@ def _validate_weights_after_download(
allow_patterns: Patterns used to match weight files
model_name_or_path: Model identifier for error messages
Raises:
RuntimeError: If any weight files are corrupted
Returns:
True if all files are valid, False if corrupted files were found and cleaned up
"""
import glob as glob_module
@@ -462,7 +481,7 @@ def _validate_weights_after_download(
weight_files.extend(glob_module.glob(os.path.join(hf_folder, pattern)))
if not weight_files:
return # No weight files to validate
return True # No weight files to validate
# Validate safetensors files
corrupted_files = []
@@ -477,11 +496,15 @@ def _validate_weights_after_download(
model_name_or_path,
[os.path.join(hf_folder, f) for f in corrupted_files],
)
raise RuntimeError(
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. "
"Please retry to re-download the model."
"Will retry download.",
)
return False
return True
def download_weights_from_hf(
@@ -490,6 +513,7 @@ def download_weights_from_hf(
allow_patterns: List[str],
revision: Optional[str] = None,
ignore_patterns: Optional[Union[str, List[str]]] = None,
max_retries: int = 3,
) -> str:
"""Download model weights from Hugging Face Hub.
@@ -504,14 +528,25 @@ def download_weights_from_hf(
ignore_patterns (Optional[Union[str, List[str]]]): The patterns to
filter out the weight files. Files matched by any of the patterns
will be ignored.
max_retries (int): Maximum number of download retries if corruption
is detected. Defaults to 3.
Returns:
str: The path to the downloaded model weights.
"""
# For local paths, no HF operations needed
if os.path.isdir(model_name_or_path):
return model_name_or_path
# Always check for valid local cache first.
# This validates cached files and cleans up corrupted ones.
path = find_local_hf_snapshot_dir(
# Use a SINGLE lock for the entire operation (validation + cleanup + download)
# to prevent race conditions where:
# 1. Process A validates, finds corruption, deletes corrupted file
# 2. Process B validates, sees missing file, deletes ENTIRE cache
# 3. Process A tries to download but cache is gone
# By using one lock, validation/cleanup and download are atomic.
with get_lock(model_name_or_path, cache_dir):
# Check for valid local cache first (validates and cleans up if needed)
path = _find_local_hf_snapshot_dir_unlocked(
model_name_or_path, cache_dir, allow_patterns, revision
)
if path is not None:
@@ -533,9 +568,9 @@ def download_weights_from_hf(
break
log_info_on_rank0(logger, f"Using model weights format {allow_patterns}")
# Use file lock to prevent multiple processes from
# downloading the same model weights at the same time.
with get_lock(model_name_or_path, cache_dir):
# Retry loop for handling corrupted downloads
for attempt in range(max_retries):
hf_folder = snapshot_download(
model_name_or_path,
allow_patterns=allow_patterns,
@@ -547,8 +582,29 @@ def download_weights_from_hf(
)
# Validate downloaded files to catch corruption early
_validate_weights_after_download(hf_folder, allow_patterns, model_name_or_path)
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