#!/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 1 if download was required (indicating cache corruption), which causes the CI job to fail and surface cache issues. """ import os import re 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]] = { "8-gpu-h200": ["deepseek-ai/DeepSeek-V3-0324", "moonshotai/Kimi-K2-Thinking"], } 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]]: """ 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) """ # 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 single model file single_files = list(model_path.glob("model.safetensors")) or list( model_path.glob("pytorch_model.bin") ) if single_files: # Validate the single safetensors file if it exists if single_files[0].suffix == ".safetensors": is_valid, error_msg = validate_safetensors_file(single_files[0]) if not is_valid: return False, f"Corrupted file {single_files[0].name}: {error_msg}" return True, None return False, "No model 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) 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...") for shard_file in shard_files: if shard_file.suffix == ".safetensors": is_valid, error_msg = validate_safetensors_file(shard_file) if not is_valid: return False, f"Corrupted shard {shard_file.name}: {error_msg}" return True, None def validate_model(model_id: str, cache_dir: str) -> Tuple[bool, Optional[str]]: """ Validate a model's cache integrity. Args: model_id: Model identifier cache_dir: HuggingFace cache directory Returns: Tuple of (is_valid, error_message) """ 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 = validate_model_shards(model_path) if not is_valid: return False, error_msg print(f" ✓ Model validated successfully") return True, None def download_model(model_id: str) -> bool: """ Download a model from HuggingFace. Args: model_id: Model identifier 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}") print(f" 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 or runner doesn't need validation 1 if models needed to be downloaded or validation failed """ 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 models_needing_download = [] validation_errors = [] # Validate each required model for model_id in required_models: is_valid, error_msg = validate_model(model_id, cache_dir) if not is_valid: print(f" ✗ Validation failed: {error_msg}") models_needing_download.append(model_id) validation_errors.append(f"{model_id}: {error_msg}") 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 error in validation_errors: print(f" - {error}") print("-" * 70) print("Attempting to download missing/corrupted models...") print("-" * 70) download_failed = False for model_id in models_needing_download: if not download_model(model_id): download_failed = True print("-" * 70) if download_failed: print("✗ FAILED: Some models could not be downloaded") return 1 # All downloads succeeded, but we still exit with error to flag cache issues print("✗ FAILED: Models were downloaded due to cache corruption/missing files") print("This indicates the cache was invalid and needed to be repaired.") print("Failing the job to surface this issue for investigation.") return 1 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)