From 83e7207763c0682a6289fd5618f4a9aa23b325d6 Mon Sep 17 00:00:00 2001
From: alisonshao <54658187+alisonshao@users.noreply.github.com>
Date: Mon, 24 Nov 2025 20:02:43 -0800
Subject: [PATCH] [diffusion] CI: add validation and cleanup for corrupted
 safetensors in multimodal loader (#13870)

Co-authored-by: Mick <mickjagger19@icloud.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
---
 .../runtime/loader/weight_utils.py            | 57 +++++++++++++++++++
 1 file changed, 57 insertions(+)

diff --git a/python/sglang/multimodal_gen/runtime/loader/weight_utils.py b/python/sglang/multimodal_gen/runtime/loader/weight_utils.py
index f0e45a95f..a1d9eade8 100644
--- a/python/sglang/multimodal_gen/runtime/loader/weight_utils.py
+++ b/python/sglang/multimodal_gen/runtime/loader/weight_utils.py
@@ -115,6 +115,30 @@ def filter_files_not_needed_for_inference(hf_weights_files: list[str]) -> list[s
 _BAR_FORMAT = "{desc}: {percentage:3.0f}% Completed | {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]\n"  # noqa: E501
 
 
+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 file is valid, False if corrupted
+    """
+    try:
+        with safe_open(file_path, framework="pt", device="cpu") as f:
+            _ = list(f.keys())
+        return True
+    except Exception as e:
+        logger.error(
+            "Corrupted safetensors file detected: %s - %s: %s",
+            file_path,
+            type(e).__name__,
+            str(e),
+        )
+        return False
+
+
 def safetensors_weights_iterator(
     hf_weights_files: list[str],
     to_cpu: bool = True,
@@ -124,6 +148,39 @@ def safetensors_weights_iterator(
         not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0
     )
     device = "cpu" if to_cpu else str(get_local_torch_device())
+
+    # Validate files before loading
+    corrupted_files = [st_file for st_file in hf_weights_files if not _validate_safetensors_file(st_file)]
+
+    if corrupted_files:
+        # Delete corrupted files (both symlink and blob if applicable)
+        for file_path in corrupted_files:
+            try:
+                if os.path.islink(file_path):
+                    blob_path = os.path.realpath(file_path)
+                    os.remove(file_path)
+                    logger.info(
+                        "Removed corrupted symlink: %s", os.path.basename(file_path)
+                    )
+                    if os.path.exists(blob_path):
+                        os.remove(blob_path)
+                        logger.info(
+                            "Removed corrupted blob: %s", os.path.basename(blob_path)
+                        )
+                elif os.path.isfile(file_path):
+                    os.remove(file_path)
+                    logger.info(
+                        "Removed corrupted file: %s", os.path.basename(file_path)
+                    )
+            except Exception as e:
+                logger.warning("Failed to remove corrupted file %s: %s", file_path, e)
+
+        raise RuntimeError(
+            f"Found {len(corrupted_files)} corrupted safetensors file(s). "
+            f"Files have been removed: {[os.path.basename(f) for f in corrupted_files]}. "
+            "Please retry - the files will be re-downloaded automatically."
+        )
+
     for st_file in tqdm(
         hf_weights_files,
         desc="Loading safetensors checkpoint shards",