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,131 +321,144 @@ def find_local_hf_snapshot_dir(
     if found_local_snapshot_dir is None:
         return None
 
+    # 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, "..", ".."))
+    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:
+            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):
+                    continue
+                local_weight_files.append(f)
+    except Exception as e:
+        logger.warning(
+            "Failed to scan local snapshot %s with patterns %s: %s",
+            found_local_snapshot_dir,
+            allow_patterns,
+            e,
+        )
+        local_weight_files = []
+
+    # 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:
+                # Cannot selectively clean (e.g., missing shards) - remove entire cache
+                log_info_on_rank0(
+                    logger,
+                    f"Validation failed for {model_name_or_path}: {error_msg}. "
+                    "Will remove entire cache and re-download.",
+                )
+                _cleanup_corrupted_model_cache(
+                    model_name_or_path, found_local_snapshot_dir, error_msg
+                )
+                return None
+
+        # 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
+
+    if len(local_weight_files) > 0:
+        log_info_on_rank0(
+            logger,
+            f"Found local HF snapshot for {model_name_or_path} at "
+            f"{found_local_snapshot_dir}; skipping download.",
+        )
+        return found_local_snapshot_dir
+    else:
+        log_info_on_rank0(
+            logger,
+            f"Local HF snapshot at {found_local_snapshot_dir} has no files matching "
+            f"{allow_patterns}; will attempt download.",
+        )
+        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.
-    # This prevents race conditions where multiple ranks detect corruption
+    with get_lock(model_name_or_path, cache_dir):
-    # and try to delete the same files at the same time.
+        return _find_local_hf_snapshot_dir_unlocked(
-    with get_lock(model_name_or_path, cache_dir, suffix="-validation"):
+            model_name_or_path, cache_dir, allow_patterns, revision
-        # 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(
-            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:
-                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):
-                        continue
-                    local_weight_files.append(f)
-        except Exception as e:
-            logger.warning(
-                "Failed to scan local snapshot %s with patterns %s: %s",
-                found_local_snapshot_dir,
-                allow_patterns,
-                e,
-            )
-            local_weight_files = []
-
-        # 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:
-                    # Cannot selectively clean (e.g., missing shards) - remove entire cache
-                    log_info_on_rank0(
-                        logger,
-                        f"Validation failed for {model_name_or_path}: {error_msg}. "
-                        "Will remove entire cache and re-download.",
-                    )
-                    _cleanup_corrupted_model_cache(
-                        model_name_or_path, found_local_snapshot_dir, error_msg
-                    )
-                    return None
-
-            # 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
-
-        if len(local_weight_files) > 0:
-            log_info_on_rank0(
-                logger,
-                f"Found local HF snapshot for {model_name_or_path} at "
-                f"{found_local_snapshot_dir}; skipping download.",
-            )
-            return found_local_snapshot_dir
-        else:
-            log_info_on_rank0(
-                logger,
-                f"Local HF snapshot at {found_local_snapshot_dir} has no files matching "
-                f"{allow_patterns}; will attempt download.",
-            )
-            return None
 
 
 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,52 +528,84 @@ 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
-        model_name_or_path, cache_dir, allow_patterns, revision
+    # 2. Process B validates, sees missing file, deletes ENTIRE cache
-    )
+    # 3. Process A tries to download but cache is gone
-    if path is not None:
+    # By using one lock, validation/cleanup and download are atomic.
-        # Valid local cache found, skip download
-        return path
-
-    # In CI, skip HF API calls if we're in offline mode or want to avoid rate limits
-    # But we already checked for local cache above, so if we're here we need to download
-    if not huggingface_hub.constants.HF_HUB_OFFLINE:
-        # Before we download we look at what is available:
-        fs = HfFileSystem()
-        file_list = fs.ls(model_name_or_path, detail=False, revision=revision)
-
-        # depending on what is available we download different things
-        for pattern in allow_patterns:
-            matching = fnmatch.filter(file_list, pattern)
-            if len(matching) > 0:
-                allow_patterns = [pattern]
-                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):
-        hf_folder = snapshot_download(
+        # Check for valid local cache first (validates and cleans up if needed)
-            model_name_or_path,
+        path = _find_local_hf_snapshot_dir_unlocked(
-            allow_patterns=allow_patterns,
+            model_name_or_path, cache_dir, allow_patterns, revision
-            ignore_patterns=ignore_patterns,
-            cache_dir=cache_dir,
-            tqdm_class=DisabledTqdm,
-            revision=revision,
-            local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,
         )
+        if path is not None:
+            # Valid local cache found, skip download
+            return path
 
-    # Validate downloaded files to catch corruption early
+        # In CI, skip HF API calls if we're in offline mode or want to avoid rate limits
-    _validate_weights_after_download(hf_folder, allow_patterns, model_name_or_path)
+        # But we already checked for local cache above, so if we're here we need to download
+        if not huggingface_hub.constants.HF_HUB_OFFLINE:
+            # Before we download we look at what is available:
+            fs = HfFileSystem()
+            file_list = fs.ls(model_name_or_path, detail=False, revision=revision)
 
-    return hf_folder
+            # depending on what is available we download different things
+            for pattern in allow_patterns:
+                matching = fnmatch.filter(file_list, pattern)
+                if len(matching) > 0:
+                    allow_patterns = [pattern]
+                    break
+
+        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 download_safetensors_index_file_from_hf(