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

python/sglang/srt/model_loader/weight_utils.py

@@ -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
+    # Check if snapshot dir exists (might have been cleaned by another process
-    # validating and cleaning up the same model cache simultaneously.
+    # before we acquired the lock)
-    # 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)
     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(
+    repo_folder = os.path.abspath(os.path.join(found_local_snapshot_dir, "..", ".."))
-            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(
+            matched_files = glob.glob(os.path.join(found_local_snapshot_dir, pattern))
-                    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(
+                _cleanup_corrupted_files_selective(model_name_or_path, corrupted_files)
-                        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:
+    Returns:
-        RuntimeError: If any weight files are corrupted
+        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.
+    # Use a SINGLE lock for the entire operation (validation + cleanup + download)
-    # This validates cached files and cleans up corrupted ones.
+    # to prevent race conditions where:
-    path = find_local_hf_snapshot_dir(
+    # 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.
+        # Retry loop for handling corrupted downloads
-    with get_lock(model_name_or_path, cache_dir):
+        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