#!/usr/bin/env python3 """ Validate model integrity for CI runners and download if needed. This script checks HuggingFace cache for model completeness and downloads missing models. It exits with code 0 if models are present or successfully downloaded (emitting a warning annotation if repairs were needed), and exits with code 1 only if download attempts fail. """ import os import re import shutil import sys from pathlib import Path from typing import Dict, List, Optional, Tuple try: from huggingface_hub import constants, snapshot_download HF_HUB_AVAILABLE = True except ImportError: print( "Warning: huggingface_hub not available. Install with: pip install huggingface_hub" ) HF_HUB_AVAILABLE = False try: from safetensors import safe_open SAFETENSORS_AVAILABLE = True except ImportError: print("Warning: safetensors not available. Install with: pip install safetensors") SAFETENSORS_AVAILABLE = False # Mapping of runner labels to their required models # Add new runner labels and models here as needed RUNNER_LABEL_MODEL_MAP: Dict[str, List[str]] = { "1-gpu-runner": [ "deepseek-ai/DeepSeek-OCR", "google/gemma-3-4b-it", "lmms-lab/llava-onevision-qwen2-0.5b-ov", "lmsys/sglang-ci-dsv3-test", "lmsys/sglang-EAGLE-llama2-chat-7B", "lmsys/sglang-EAGLE3-LLaMA3.1-Instruct-8B", "meta-llama/Llama-2-7b-chat-hf", "meta-llama/Llama-3.2-1B-Instruct", "meta-llama/Llama-3.1-8B-Instruct", "mistralai/Mixtral-8x7B-Instruct-v0.1", "moonshotai/Kimi-VL-A3B-Instruct", "nvidia/NVIDIA-Nemotron-Nano-9B-v2", "nvidia/NVIDIA-Nemotron-Nano-9B-v2-FP8", "OpenGVLab/InternVL2_5-2B", "Qwen/Qwen2.5-7B-Instruct", "Qwen/Qwen3-8B", "Qwen/Qwen3-Coder-30B-A3B-Instruct", "Qwen/QwQ-32B-AWQ", "Qwen/Qwen3-30B-A3B", ], "2-gpu-runner": [ "mistralai/Mixtral-8x7B-Instruct-v0.1", "moonshotai/Kimi-Linear-48B-A3B-Instruct", "Qwen/Qwen2-57B-A14B-Instruct", "Qwen/Qwen2.5-VL-7B-Instruct", "Qwen/Qwen3-VL-30B-A3B-Instruct", "neuralmagic/Qwen2-72B-Instruct-FP8", "zai-org/GLM-4.5-Air-FP8", ], "8-gpu-h200": [ "deepseek-ai/DeepSeek-V3-0324", "deepseek-ai/DeepSeek-V3.2-Exp", "moonshotai/Kimi-K2-Thinking", ], "8-gpu-b200": ["deepseek-ai/DeepSeek-V3.1", "deepseek-ai/DeepSeek-V3.2-Exp"], "4-gpu-b200": ["nvidia/DeepSeek-V3-0324-FP4"], "4-gpu-gb200": ["nvidia/DeepSeek-V3-0324-FP4"], "4-gpu-h100": ["lmsys/sglang-ci-dsv3-test", "lmsys/sglang-ci-dsv3-test-NextN"], } def get_hf_cache_dir() -> str: """Get the HuggingFace cache directory.""" if HF_HUB_AVAILABLE: return constants.HF_HUB_CACHE # Fallback to environment variable or default hf_home = os.environ.get("HF_HOME", os.path.expanduser("~/.cache/huggingface")) return os.path.join(hf_home, "hub") def get_model_cache_path(model_id: str, cache_dir: str) -> Optional[Path]: """ Find the model's cache directory in HuggingFace hub cache. Args: model_id: Model identifier (e.g., "deepseek-ai/DeepSeek-V3-0324") cache_dir: HuggingFace cache directory Returns: Path to model's snapshot directory, or None if not found """ # Convert model_id to cache directory name format # "deepseek-ai/DeepSeek-V3-0324" -> "models--deepseek-ai--DeepSeek-V3-0324" cache_model_name = "models--" + model_id.replace("/", "--") model_path = Path(cache_dir) / cache_model_name if not model_path.exists(): return None # Find the most recent snapshot directory snapshots_dir = model_path / "snapshots" if not snapshots_dir.exists(): return None # Get all snapshot directories (sorted by modification time, most recent first) snapshot_dirs = sorted( [d for d in snapshots_dir.iterdir() if d.is_dir()], key=lambda x: x.stat().st_mtime, reverse=True, ) if not snapshot_dirs: return None return snapshot_dirs[0] def check_incomplete_files(model_path: Path, cache_dir: str) -> List[str]: """ Check for incomplete download marker files specific to this model. Args: model_path: Path to model's snapshot directory cache_dir: HuggingFace cache directory Returns: List of incomplete files found for this specific model """ incomplete_in_snapshot = [] # Check if any files in the snapshot are symlinks to .incomplete blobs # This ensures we only flag incomplete files for THIS specific model, # not other models that might be downloading concurrently for file_path in model_path.glob("*"): if file_path.is_symlink(): try: target = file_path.resolve() # Check if the symlink target has .incomplete suffix if str(target).endswith(".incomplete"): incomplete_in_snapshot.append(str(target)) except (OSError, RuntimeError): # Broken symlink - also indicates incomplete download incomplete_in_snapshot.append(str(file_path)) return incomplete_in_snapshot def validate_safetensors_file(file_path: Path) -> Tuple[bool, Optional[str]]: """ Validate that a safetensors file is readable and not corrupted. Args: file_path: Path to the safetensors file Returns: Tuple of (is_valid, error_message) """ if not SAFETENSORS_AVAILABLE: # Skip validation if safetensors library is not available return True, None try: # Attempt to open and read the header # This will fail if the file is corrupted or incomplete with safe_open(file_path, framework="pt", device="cpu") as f: # Just accessing the keys validates the header is readable _ = f.keys() return True, None except Exception as e: error_type = type(e).__name__ error_msg = str(e) # Return detailed error for debugging return False, f"{error_type}: {error_msg}" def validate_model_shards(model_path: Path) -> Tuple[bool, Optional[str], List[Path]]: """ Validate that all model shards are present and complete. Args: model_path: Path to model's snapshot directory Returns: Tuple of (is_valid, error_message, corrupted_files) - corrupted_files: List of paths to corrupted shard files that should be removed """ # Pattern for sharded files: model-00001-of-00009.safetensors or pytorch_model-00001-of-00009.bin shard_pattern = re.compile( r"(?:model|pytorch_model)-(\d+)-of-(\d+)\.(safetensors|bin)" ) # Find all shard files (both .safetensors and .bin) shard_files = ( list(model_path.glob("model-*-of-*.safetensors")) + list(model_path.glob("model-*-of-*.bin")) + list(model_path.glob("pytorch_model-*-of-*.bin")) ) if not shard_files: # No sharded files - check for any safetensors or bin files # Exclude non-model files like tokenizer, config, optimizer, etc. all_safetensors = list(model_path.glob("*.safetensors")) all_bins = list(model_path.glob("*.bin")) # Filter out non-model files excluded_prefixes = ["tokenizer", "optimizer", "training_", "config"] single_files = [ f for f in (all_safetensors or all_bins) if not any(f.name.startswith(prefix) for prefix in excluded_prefixes) and not f.name.endswith(".index.json") ] if single_files: # Validate all safetensors files, not just the first one for model_file in single_files: if model_file.suffix == ".safetensors": is_valid, error_msg = validate_safetensors_file(model_file) if not is_valid: return ( False, f"Corrupted file {model_file.name}: {error_msg}", [model_file], ) return True, None, [] return False, "No model weight files found (safetensors or bin)", [] # Extract total shard count from any shard filename total_shards = None for shard_file in shard_files: match = shard_pattern.search(shard_file.name) if match: total_shards = int(match.group(2)) break if total_shards is None: return False, "Could not determine total shard count from filenames", [] # Check that all shards exist expected_shards = set(range(1, total_shards + 1)) found_shards = set() for shard_file in shard_files: match = shard_pattern.search(shard_file.name) if match: shard_num = int(match.group(1)) found_shards.add(shard_num) missing_shards = expected_shards - found_shards if missing_shards: missing_list = sorted(missing_shards) # Missing shards - nothing to remove, let download handle it return ( False, f"Missing shards: {missing_list} (expected {total_shards} total)", [], ) # Check for index file index_file = model_path / "model.safetensors.index.json" if not index_file.exists(): return False, "Missing model.safetensors.index.json", [] # Validate each safetensors shard file for corruption print(f" Validating {len(shard_files)} shard file(s) for corruption...") corrupted_files = [] for shard_file in shard_files: if shard_file.suffix == ".safetensors": is_valid, error_msg = validate_safetensors_file(shard_file) if not is_valid: corrupted_files.append(shard_file) print(f" ✗ Corrupted: {shard_file.name} - {error_msg}") if corrupted_files: return ( False, f"Corrupted shards: {[f.name for f in corrupted_files]}", corrupted_files, ) return True, None, [] def validate_model( model_id: str, cache_dir: str ) -> Tuple[bool, Optional[str], List[Path]]: """ Validate a model's cache integrity. Args: model_id: Model identifier cache_dir: HuggingFace cache directory Returns: Tuple of (is_valid, error_message, corrupted_files) - corrupted_files: List of paths to corrupted files that should be removed """ print(f"Validating model: {model_id}") # Find model in cache model_path = get_model_cache_path(model_id, cache_dir) if model_path is None: return False, "Model not found in cache", [] print(f" Found in cache: {model_path}") # Check for incomplete files incomplete_files = check_incomplete_files(model_path, cache_dir) if incomplete_files: return ( False, f"Found incomplete download files: {len(incomplete_files)} files", [], ) # Validate shards is_valid, error_msg, corrupted_files = validate_model_shards(model_path) if not is_valid: return False, error_msg, corrupted_files print(f" ✓ Model validated successfully") return True, None, [] def download_model(model_id: str, cache_dir: str, corrupted_files: List[Path]) -> bool: """ Download a model from HuggingFace. Completely removes the model cache directory before downloading to ensure a clean download. Args: model_id: Model identifier cache_dir: HuggingFace cache directory corrupted_files: List of specific file paths that are corrupted (unused, kept for compatibility) Returns: True if download succeeded, False otherwise """ if not HF_HUB_AVAILABLE: print(f"ERROR: Cannot download model - huggingface_hub not available") return False print(f"Downloading model: {model_id}") # Completely remove the model directory from cache cache_model_name = "models--" + model_id.replace("/", "--") model_cache_path = Path(cache_dir) / cache_model_name if model_cache_path.exists(): print(f" Removing entire model directory: {model_cache_path}") try: shutil.rmtree(model_cache_path) print(f" ✓ Successfully removed model directory") except Exception as e: print(f" ✗ Failed to remove model directory: {e}") print(f" Attempting download anyway...") else: print(f" Model directory not found in cache (will download fresh)") print(f" Downloading from HuggingFace (this may take a while for large models)...") try: snapshot_download( repo_id=model_id, allow_patterns=["*.safetensors", "*.bin", "*.json", "*.txt", "*.model"], ignore_patterns=["*.msgpack", "*.h5", "*.ot"], # codespell:ignore ot ) print(f" ✓ Download completed: {model_id}") return True except Exception as e: print(f" ✗ Download failed: {e}") return False def get_runner_labels() -> List[str]: """ Get the runner labels from environment variables. GitHub Actions doesn't expose runner labels directly as environment variables. Workflows should set the RUNNER_LABELS environment variable with a comma-separated list of labels (e.g., "self-hosted,8-gpu-h200,linux"). Returns: List of runner labels, empty list if not set """ labels_str = os.environ.get("RUNNER_LABELS", "") if not labels_str: return [] # Split by comma and strip whitespace return [label.strip() for label in labels_str.split(",") if label.strip()] def should_validate_runner(runner_labels: List[str]) -> bool: """ Check if the runner should have model validation based on its labels. Args: runner_labels: List of runner labels Returns: True if any label matches a configured label in RUNNER_LABEL_MODEL_MAP """ if not runner_labels: return False # Check if any label is in the configured map return any(label in RUNNER_LABEL_MODEL_MAP for label in runner_labels) def get_required_models(runner_labels: List[str]) -> List[str]: """ Get list of models required based on runner labels. Args: runner_labels: List of runner labels (e.g., ["self-hosted", "8-gpu-h200", "linux"]) Returns: List of model identifiers to validate (deduplicated) """ all_models = [] for label in runner_labels: if label in RUNNER_LABEL_MODEL_MAP: models = RUNNER_LABEL_MODEL_MAP[label] print( f" ✓ Matched label configuration: '{label}' -> {len(models)} model(s)" ) all_models.extend(models) if not all_models: print(f" ⚠ No configuration found for any label in: {runner_labels}") # Remove duplicates while preserving order seen = set() unique_models = [] for model in all_models: if model not in seen: seen.add(model) unique_models.append(model) return unique_models def main() -> int: """ Main validation logic. Returns: 0 if all models are valid, successfully downloaded, or runner doesn't need validation 1 only if download attempts fail """ print("=" * 70) print("Model Validation for CI Runners") print("=" * 70) runner_labels = get_runner_labels() print(f"Runner labels: {', '.join(runner_labels) if runner_labels else 'NOT SET'}") # Check if this runner needs validation if not should_validate_runner(runner_labels): print( "Skipping validation: No runner labels match configured model requirements" ) return 0 print(f"Proceeding with model validation for this runner") # Get required models for these runner labels required_models = get_required_models(runner_labels) if not required_models: print(f"Warning: No models configured for labels: {runner_labels}") return 0 print(f"Models to validate: {required_models}") print("-" * 70) # Get cache directory cache_dir = get_hf_cache_dir() print(f"HuggingFace cache: {cache_dir}") print("-" * 70) # Track validation results # Maps model_id -> (error_msg, corrupted_files) models_needing_download: Dict[str, Tuple[str, List[Path]]] = {} # Validate each required model for model_id in required_models: is_valid, error_msg, corrupted_files = validate_model(model_id, cache_dir) if not is_valid: print(f" ✗ Validation failed: {error_msg}") models_needing_download[model_id] = (error_msg, corrupted_files) print("-" * 70) # If all models are valid, exit successfully if not models_needing_download: print("✓ All models validated successfully!") return 0 # Models need to be downloaded print(f"⚠ Cache validation failed for {len(models_needing_download)} model(s)") for model_id, (error_msg, _) in models_needing_download.items(): print(f" - {model_id}: {error_msg}") print("-" * 70) print("Attempting to download missing/corrupted models...") print("-" * 70) download_failed = False for model_id, (error_msg, corrupted_files) in models_needing_download.items(): if not download_model(model_id, cache_dir, corrupted_files): download_failed = True print("-" * 70) if download_failed: print("✗ FAILED: Some models could not be downloaded") return 1 # All downloads succeeded - now validate them again print("✓ All models downloaded successfully!") print("-" * 70) print("Validating downloaded models...") print("-" * 70) validation_failed = False for model_id in models_needing_download.keys(): is_valid, error_msg, _ = validate_model(model_id, cache_dir) if not is_valid: print(f" ✗ Post-download validation failed for {model_id}: {error_msg}") validation_failed = True print("-" * 70) if validation_failed: print("✗ FAILED: Some models failed validation after download") return 1 # All validations passed - emit warning but exit successfully print("✓ All downloaded models validated successfully!") print("⚠ WARNING: Models were missing/corrupted in cache and have been repaired.") print(f" Repaired models: {', '.join(models_needing_download.keys())}") # Emit GitHub Actions warning annotation for visibility print( f"::warning file=scripts/ci/validate_and_download_models.py::" f"Cache validation failed for {len(models_needing_download)} model(s). " f"Models were re-downloaded and validated successfully. " f"This may indicate cache corruption or infrastructure issues." ) return 0 if __name__ == "__main__": try: exit_code = main() sys.exit(exit_code) except KeyboardInterrupt: print("\nInterrupted by user") sys.exit(1) except Exception as e: print(f"ERROR: Unexpected error: {e}") import traceback traceback.print_exc() sys.exit(1)