From 5b4f790200929beb52ba51a6c7b82fbf45e8c9b1 Mon Sep 17 00:00:00 2001 From: Alison Shao <54658187+alisonshao@users.noreply.github.com> Date: Fri, 2 Jan 2026 22:12:45 -0800 Subject: [PATCH] [diffusion] CI: add CI validation for diffusion model downloads (#16311) --- .../runtime/utils/hf_diffusers_utils.py | 226 +++++++++++++++++- 1 file changed, 222 insertions(+), 4 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py b/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py index a4b0dc25f..e661459c3 100644 --- a/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py +++ b/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py @@ -19,8 +19,10 @@ """Utilities for Huggingface Transformers.""" import contextlib +import glob import json import os +import shutil import time from functools import reduce from pathlib import Path @@ -43,8 +45,166 @@ from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_N from sglang.multimodal_gen.runtime.loader.weight_utils import get_lock from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger +from sglang.utils import is_in_ci logger = init_logger(__name__) + + +def _check_index_files_for_missing_shards( + model_path: str, +) -> tuple[bool, list[str], list[str]]: + """ + Check all subdirectories for missing shards based on index files. + + This catches cases where a model download was interrupted, leaving + some safetensors shards missing while the index file exists. + + Args: + model_path: Path to the model directory + + Returns: + Tuple of (all_valid, missing_files, checked_subdirs) + """ + missing_files = [] + checked_subdirs = [] + + # Check the root directory and all subdirectories that might contain model weights + dirs_to_check = [model_path] + + # Add common subdirectories for diffusers models + try: + subdirs = os.listdir(model_path) + except OSError as e: + logger.warning("Failed to list model directory %s: %s", model_path, e) + return True, [], [] # Assume valid if we can't check + + for subdir in subdirs: + subdir_path = os.path.join(model_path, subdir) + if os.path.isdir(subdir_path): + dirs_to_check.append(subdir_path) + + for dir_path in dirs_to_check: + # Find all safetensors index files + index_files = glob.glob(os.path.join(dir_path, "*.safetensors.index.json")) + + for index_file in index_files: + checked_subdirs.append(os.path.basename(dir_path)) + try: + with open(index_file) as f: + index_data = json.load(f) + + weight_map = index_data.get("weight_map", {}) + if not weight_map: + continue + + # Get unique files referenced in weight_map + required_files = set(weight_map.values()) + + for file_name in required_files: + file_path = os.path.join(dir_path, file_name) + if not os.path.exists(file_path): + relative_path = os.path.relpath(file_path, model_path) + missing_files.append(relative_path) + + except Exception as e: + logger.warning("Failed to read index file %s: %s", index_file, e) + continue + + return len(missing_files) == 0, missing_files, checked_subdirs + + +def _cleanup_model_cache(model_path: str, reason: str) -> bool: + """ + Remove the model cache directory to force a clean re-download. + + Args: + model_path: Path to the model directory (snapshot path) + reason: Reason for cleanup (for logging) + + Returns: + True if cleanup was performed, False otherwise + """ + # Navigate up to the model root directory: snapshots/hash -> snapshots -> model_root + # HF cache structure: models--org--name/snapshots/hash/ + try: + snapshot_dir = os.path.abspath(model_path) + snapshots_dir = os.path.dirname(snapshot_dir) + repo_folder = os.path.dirname(snapshots_dir) + + # Verify this looks like an HF cache structure + if os.path.basename(snapshots_dir) != "snapshots": + logger.warning( + "Model path %s doesn't appear to be in HF cache structure, skipping cleanup", + model_path, + ) + return False + + logger.warning( + "Removing model cache at %s. Reason: %s", + repo_folder, + reason, + ) + shutil.rmtree(repo_folder) + logger.info("Successfully removed corrupted cache directory") + return True + except Exception as e: + logger.error( + "Failed to remove corrupted cache directory %s: %s. " + "Manual cleanup may be required.", + model_path, + e, + ) + return False + + +def _ci_validate_diffusers_model(model_path: str) -> tuple[bool, bool]: + """ + CI-specific validation for diffusers models. + + Checks all subdirectories (transformer, transformer_2, vae, etc.) for + missing shards based on their index files. If issues are found in CI, + cleans up the cache to force re-download. + + Args: + model_path: Path to the model directory + + Returns: + Tuple of (is_valid, cleanup_performed) + - is_valid: True if the model is valid + - cleanup_performed: True if cleanup was performed (only relevant when is_valid=False) + """ + if not is_in_ci(): + return True, False + + is_valid, missing_files, checked_subdirs = _check_index_files_for_missing_shards( + model_path + ) + + if not is_valid: + logger.error( + "CI validation failed for %s. Missing %d file(s): %s. " + "Checked subdirectories: %s", + model_path, + len(missing_files), + missing_files[:5] if len(missing_files) > 5 else missing_files, + checked_subdirs, + ) + cleanup_performed = _cleanup_model_cache( + model_path, + f"Missing {len(missing_files)} shard file(s): {missing_files[:3]}", + ) + return False, cleanup_performed + + if checked_subdirs: + logger.info( + "CI validation passed for %s. Checked subdirectories: %s", + model_path, + checked_subdirs, + ) + + return True, False + + _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = { # ChatGLMConfig.model_type: ChatGLMConfig, # DbrxConfig.model_type: DbrxConfig, @@ -402,8 +562,33 @@ def maybe_download_model( # 1. Local path check: if path exists locally, verify it's complete (skip for LoRA) if os.path.exists(model_name_or_path): if is_lora or _verify_model_complete(model_name_or_path): - logger.info("Model already exists locally and is complete") - return model_name_or_path + # CI validation: check all subdirectories for missing shards + if not is_lora: + is_valid, cleanup_performed = _ci_validate_diffusers_model( + model_name_or_path + ) + if not is_valid: + if cleanup_performed: + logger.warning( + "CI validation failed for local model at %s, " + "cache has been cleaned up, will re-download", + model_name_or_path, + ) + # Fall through to download + else: + # Local path is not in HF cache structure, can't clean up + # Raise error since we can't fix this automatically + raise ValueError( + f"CI validation failed for local model at {model_name_or_path}. " + "Some safetensors shards are missing. " + "Please manually delete the model directory and retry." + ) + else: + logger.info("Model already exists locally and is complete") + return model_name_or_path + else: + logger.info("Model already exists locally and is complete") + return model_name_or_path else: logger.warning( "Local model at %s appears incomplete (missing transformer/ or vae/), " @@ -427,8 +612,30 @@ def maybe_download_model( etag_timeout=60, ) if is_lora or _verify_model_complete(local_path): - logger.info("Found complete model in cache at %s", local_path) - return str(local_path) + # CI validation: check all subdirectories for missing shards + if not is_lora: + is_valid, cleanup_performed = _ci_validate_diffusers_model(local_path) + if not is_valid: + if cleanup_performed: + logger.warning( + "CI validation failed for cached model at %s, " + "cache has been cleaned up, will re-download", + local_path, + ) + # Fall through to download + else: + # This shouldn't happen for HF cache paths, but handle it + logger.warning( + "CI validation failed for cached model at %s, " + "but cleanup was not performed, will attempt re-download", + local_path, + ) + else: + logger.info("Found complete model in cache at %s", local_path) + return str(local_path) + else: + logger.info("Found complete model in cache at %s", local_path) + return str(local_path) else: # Model found in cache but incomplete if not download: @@ -498,6 +705,17 @@ def maybe_download_model( "The model repository may be missing required components (model_index.json, transformer/, or vae/)." ) + # CI validation: check all subdirectories for missing shards after download + if not is_lora: + is_valid, cleanup_performed = _ci_validate_diffusers_model(local_path) + if not is_valid: + # In CI, if validation fails after download, we have a serious issue + # If cleanup was performed, the next retry should get a fresh download + raise ValueError( + f"CI validation failed for downloaded model at {local_path}. " + f"Some safetensors shards are missing. Cleanup performed: {cleanup_performed}." + ) + logger.info("Downloaded model to %s", local_path) return str(local_path)