diff --git a/.claude/rules/comment-style.md b/.claude/rules/comment-style.md new file mode 100644 index 000000000..a46170ebe --- /dev/null +++ b/.claude/rules/comment-style.md @@ -0,0 +1,164 @@ +--- +paths: + - "**/*.py" + - "**/*.cu" + - "**/*.cuh" + - "**/*.cpp" + - "**/*.h" + - "**/*.rs" +--- + +# Comment Style + +Applies to `#` and `//` comments, Python docstrings, and C/C++ Doxygen blocks +(`///`, `/** ... */`) -- Python, CUDA/C++, and Rust alike. Comments are reviewed +like code, with the burden of proof reversed: the author justifies a comment's +existence, not the reviewer its removal. + +**The test.** Delete the comment. If someone familiar with the repo can recover +the fact from the surrounding code plus a grep, it stays deleted. + +The next three sections are that test applied. A comment either carries a fact +you cannot see from here (**keep**), carries nothing the code does not already +carry (**delete**), or carries a real fact whose home is somewhere else +(**move**). + +## Keep: facts you cannot see from here + +- **Cross-boundary constraints.** Layout, field order, or call order shared with + a CUDA kernel, the Rust router, an IPC schema, or another process. Nothing in + the Python file shows the other side, so nothing else can warn the next editor. + + ```python + # The PD wire schema must match on P and D even when only D runs spec decoding; + # a seedless prefill writes the invalid sentinel. + ``` + +- **Units and layout the name cannot carry.** Tokens vs reqs vs pages vs bytes + vs slots, and tensor shape/dtype/layout. Encode it in the name first; comment + only when the name is fixed by an existing interface. + + ```python + # [num_tokens, num_kv_heads, head_dim], fp8 e4m3, page-major. + ``` + +- **Where a magic number came from** -- not what it means. A hardware + constraint, a measurement, or an admission that it was picked arbitrarily. + The last one is the most valuable: it tells the next person the value is safe + to change. + + ```python + # Hardware constraint: TMA descriptors require 16B alignment. + # Measured on fp8 GEMM; re-tune when the kernel changes. + # Arbitrary; no evidence this is the right threshold. + ``` + +- **Workarounds, anchored to a verifiable reference and a retirement + condition.** An unanchored workaround is immortal -- nobody can prove it is + safe to delete. + + ```python + # Workaround for pytorch/pytorch#12345; drop once we require torch >= 2.9. + # Temporary workaround: Event.wait() regresses TPOT on AMD MI355. + ``` + (second line: `python/sglang/srt/managers/overlap_utils.py:471`) + +- **Contracts that are a decision, not a mechanism** -- a sentinel's meaning, a + deliberate omission from a list, an ordering that looks incidental. + + ```python + # None means the request is excluded from the radix cache, not that lookup failed. + ``` + +## Delete: nothing the code does not already carry + +- **Restating the code.** The function name, the next line, the branch + condition, or the loop body in prose. +- **Step-by-step play-by-play.** `# Step 1: ...` / `# Step 2: ...` over + straight-line code. If the flow needs numbering, extract named helpers. +- **Naming the callers.** That is what grep is for. +- **Section banners on short bodies.** `# ===== Helpers =====` above two + functions. Banners are for genuinely long modules only + (see `python/sglang/srt/environ.py`). +- **Hedging.** "This should probably be revisited." Either establish the fact + and state it, or file it as `TODO()`. + +## Move: real facts that live somewhere else + +Every explanation has exactly one home. A second copy in a comment drifts from +the first. + +- **What the code used to do -> `git log`.** Comments describe the current + state: not what was tried first, not which approach was abandoned. + + The exception is a past failure that is still a live constraint. That is not + history -- but write it as the constraint, not as the story. + + ```python + # Bad: We used to call this before init_new but it broke CUDA graph capture. + # Good: Must run after graph capture; capturing this allocation deadlocks. + ``` + +- **Why the change was made -> the PR body.** Why now, what else was tried, + which benchmark moved. "Now we also handle the case where ..." argues for a + diff, and is written for a reviewer who is long gone. +- **Design rationale -> `docs/` or a module docstring. CLI semantics -> the + `ServerArgs` help text. Test intent -> the test.** +- **Superseded code -> `git show`.** Never leave commented-out code behind. + +Also out: review attribution ("as suggested in review") and our own PR numbers +used as a changelog. An upstream issue URL is different -- it is not a changelog +entry but a workaround's retirement condition, and the Keep section requires it. + +What stays next to the line is why *this line* exists, written for a reader who +has neither the PR nor the discussion. + +## Form + +Governs `#` / `//` comments; the length and placement of documentation blocks +follow the section below. ASCII-and-English applies to both. + +- **One or two lines.** A genuinely intricate invariant may run longer; that is + a rare exception, not a licence. +- **ASCII and English only.** No Unicode arrows, math symbols, or CJK. +- **Break at a clause boundary.** If a comment needs a second line, wrap after a + semicolon or a comma -- never mid-phrase. A sentence that will not split + cleanly is a sentence that should be shortened instead. +- **Attach the comment to the line it constrains**, not to the top of the + function. Comments collected into a preamble are the ones that go stale. +- **You change the line, you own its comment.** Update it or delete it -- never + leave it orphaned. A stale comment is worse than no comment. + +## Tags + +Two spellings only. New code does not use `FIXME`, `XXX`, or `HACK`; existing +occurrences are grandfathered and get folded into these when the line is touched. + +- `# NOTE:` -- a constraint or trap. Most of the time the prefix adds nothing; + drop it and just state the fact. +- `# TODO():` -- planned work, **always** with an owner or an issue + link. A bare `# TODO: fix this` is rejected in review; an unowned TODO is + never retired. `TODO(perf)` and similar topic tags are acceptable when the work + is a standing category rather than one person's task. + +## API documentation blocks + +Python docstrings and Doxygen blocks are part of the product on the surfaces +users, integrators, and tooling read, and noise everywhere else. Where warranted +they may run past two lines; every other rule in this file still applies. + +- **Yes:** `Engine` and entrypoint public methods, `ServerArgs` help text, base + classes third parties subclass (attention backends, quant methods, model + extension points), and `sgl-kernel` op signatures -- where the shape/dtype + contract is the documentation. +- **Yes:** exported C++ / CUDA entities, as Doxygen with `\brief` / `\param` / + `\tparam` / `\return`. clangd renders these on hover, driven by + `CommentFormat: Doxygen`, so the per-parameter enumeration is the + caller-facing contract rather than filler. +- **Yes:** a bug-regression test, stating in one or two lines which black-box + behavior must not come back -- the live constraint, not the incident. Root + cause and repro belong in the issue and the PR + (see `.claude/rules/unit-test-admission.md`). +- **No:** internal helpers, overrides, private methods. Never hand-write or + generate Python `Args:` / `Returns:` blocks -- the signature and its type + hints already carry the names and types. diff --git a/python/sglang/srt/model_loader/ci_weight_validation.py b/python/sglang/srt/model_loader/ci_weight_validation.py index a50b2cf99..98abdd3cb 100644 --- a/python/sglang/srt/model_loader/ci_weight_validation.py +++ b/python/sglang/srt/model_loader/ci_weight_validation.py @@ -1,17 +1,8 @@ -""" -CI-specific weight validation and cache cleanup utilities. +"""CI-only weight validation and cache cleanup. -This module contains validation and cleanup logic that is ONLY used in CI environments. -These functions handle: -- Validating safetensors files for corruption -- Checking for missing shards in sharded models -- Cleaning up corrupted files (selective or full cache deletion) -- Automatic retry logic for corrupted downloads -- Validating config/tokenizer files completeness to enable offline mode - -For regular users, weight_utils.py provides simple download functionality without -the overhead of validation and automatic cleanup. The CI-specific behavior is -gated by is_in_ci() checks in weight_utils.py. +Validates safetensors/bin files and shard completeness, and repairs the HF cache +by deleting what has to be re-downloaded. `weight_utils.py` gates every entry +point here behind `is_in_ci()`; regular users take the plain download path. """ import glob as glob_module @@ -33,17 +24,7 @@ logger = logging.getLogger(__name__) def _get_per_run_marker_dir() -> str: - """ - Get the directory for per-run validation markers. - - These markers are specific to the current CI run and are not shared across - runners. They are stored in a temporary directory that is cleaned up after - the run completes. - - Returns: - Path to per-run marker directory - """ - # Prefer RUNNER_TEMP (GitHub Actions) or TMPDIR, fallback to /tmp + # Markers are per CI run; sharing them across runners leaks cache state. base_dir = os.environ.get("RUNNER_TEMP", os.environ.get("TMPDIR", "/tmp")) marker_dir = os.path.join(base_dir, "sglang_ci_offline_markers") os.makedirs(marker_dir, exist_ok=True) @@ -51,18 +32,6 @@ def _get_per_run_marker_dir() -> str: def _get_per_run_marker_path(snapshot_dir: str) -> Optional[str]: - """ - Get the path to per-run validation marker file for a snapshot. - - Per-run markers are specific to the current CI run and are not shared - across runners. This prevents cross-runner cache state pollution. - - Args: - snapshot_dir: Path to snapshot directory - - Returns: - Path to per-run marker file or None if snapshot_dir is invalid - """ if not snapshot_dir or not os.path.isdir(snapshot_dir): return None @@ -74,15 +43,6 @@ def _get_per_run_marker_path(snapshot_dir: str) -> Optional[str]: def _read_per_run_marker(snapshot_dir: str) -> Optional[dict]: - """ - Read per-run validation marker for a snapshot. - - Args: - snapshot_dir: Path to snapshot directory - - Returns: - Marker dict if exists and valid, None otherwise - """ marker_path = _get_per_run_marker_path(snapshot_dir) if not marker_path or not os.path.exists(marker_path): return None @@ -91,7 +51,6 @@ def _read_per_run_marker(snapshot_dir: str) -> Optional[dict]: with open(marker_path, "r", encoding="utf-8") as f: marker = json.load(f) - # Validate marker structure if not isinstance(marker, dict): return None @@ -112,14 +71,6 @@ def _read_per_run_marker(snapshot_dir: str) -> Optional[dict]: def _write_per_run_marker( snapshot_dir: str, model_id: str, required_files: Optional[list] = None ) -> None: - """ - Write per-run validation marker for a snapshot. - - Args: - snapshot_dir: Path to snapshot directory - model_id: Model identifier - required_files: List of required files that were validated - """ marker_path = _get_per_run_marker_path(snapshot_dir) if not marker_path: logger.debug("Cannot write per-run marker: invalid snapshot_dir") @@ -165,22 +116,7 @@ def _write_per_run_marker( def validate_cache_lightweight( snapshot_dir: str, requires_hf_quant_config: bool = False ) -> bool: - """ - Lightweight runtime validation for cache completeness. - - This is used during test runs to ensure the current runner's cache - is complete before enabling offline mode. Much faster than full validation - as it only checks file existence, not corruption. - - Args: - snapshot_dir: Path to the model snapshot directory - requires_hf_quant_config: If True, hf_quant_config.json must exist - (required for modelopt quantization) - - Returns: - True if cache is complete, False otherwise - """ - # Check required config files + """Existence-only cache check: no corruption reads, cheap enough to run per test.""" required_files = [ "config.json", "tokenizer_config.json", @@ -190,7 +126,6 @@ def validate_cache_lightweight( if not os.path.exists(os.path.join(snapshot_dir, fname)): return False - # Check tokenizer files (at least one must exist) tokenizer_files = [ "tokenizer.json", "tokenizer.model", @@ -203,7 +138,6 @@ def validate_cache_lightweight( if not has_tokenizer: return False - # Check for trust_remote_code dynamic module files if needed # When auto_map exists in config.json, the model requires custom Python files # These files must be present for offline mode to work config_path = os.path.join(snapshot_dir, "config.json") @@ -214,17 +148,13 @@ def validate_cache_lightweight( auto_map = config.get("auto_map", {}) if auto_map and isinstance(auto_map, dict): - # Extract Python module files from auto_map # auto_map format: {"AutoConfig": "configuration_xxx.ConfigClass", ...} - # We need to check if the .py files exist custom_files = set() for key, value in auto_map.items(): if isinstance(value, str) and "." in value: - # Extract module name (e.g., "configuration_xxx" from "configuration_xxx.ConfigClass") module_name = value.split(".")[0] custom_files.add(f"{module_name}.py") - # Check if all custom files exist in snapshot directory for custom_file in custom_files: custom_file_path = os.path.join(snapshot_dir, custom_file) if not os.path.exists(custom_file_path): @@ -249,13 +179,11 @@ def validate_cache_lightweight( has_index = os.path.exists(index_path) if has_index: - # If index exists, validate that all shards listed in it exist try: with open(index_path, "r", encoding="utf-8") as f: index_data = json.load(f) weight_map = index_data.get("weight_map", {}) if weight_map: - # Check that all shard files referenced in index exist required_shards = set(weight_map.values()) for shard_name in required_shards: shard_path = os.path.join(snapshot_dir, shard_name) @@ -270,7 +198,6 @@ def validate_cache_lightweight( logger.debug("Failed to validate index file %s: %s", index_path, e) return False else: - # No index file - check for weight files and validate shard completeness safetensors_files = glob_module.glob( os.path.join(snapshot_dir, "*.safetensors") ) @@ -278,7 +205,6 @@ def validate_cache_lightweight( return False # Check shard completeness for sharded models (e.g., model-00001-of-00047.safetensors) - # Pattern: prefix-NNNNN-of-NNNNN.safetensors shard_pattern = re.compile(r"(.*?)-(\d+)-of-(\d+)\.safetensors$") shard_groups = {} @@ -298,7 +224,6 @@ def validate_cache_lightweight( } shard_groups[group_key]["found_shards"].add(shard_id) - # Validate each shard group has all expected shards for group_key, group_info in shard_groups.items(): total_shards = group_info["total"] found_shards = group_info["found_shards"] @@ -324,20 +249,9 @@ def validate_cache_lightweight( 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 the file is valid, False if corrupted - """ try: - # Attempt to open and read the header - # This will fail if the file is corrupted or incomplete with safetensors.safe_open(file_path, framework="pt", device="cpu") as f: - # Just accessing the keys validates the header is readable + # Listing keys forces the header parse; open() alone does not. _ = list(f.keys()) return True except Exception as e: @@ -351,20 +265,7 @@ def _validate_safetensors_file(file_path: str) -> bool: def _validate_pytorch_bin_file(file_path: str) -> bool: - """ - Validate that a PyTorch .bin file is readable and not corrupted. - - This catches corruption issues like truncated downloads or invalid archives - that would cause errors like: - "RuntimeError: PytorchStreamReader failed reading file data/X: invalid header - or archive is corrupted" - - Args: - file_path: Path to the .bin file - - Returns: - True if the file is valid, False if corrupted - """ + # Truncated archives surface as a PytorchStreamReader "invalid header" error. try: import torch @@ -383,19 +284,6 @@ def _validate_pytorch_bin_file(file_path: str) -> bool: def _check_index_files_exist(snapshot_dir: str) -> Tuple[bool, Optional[str]]: - """ - Check if all files listed in safetensors index files actually exist on disk. - - This catches cases where the snapshot directory exists but files are missing - (e.g., due to incomplete downloads or corrupted cache). - - Args: - snapshot_dir: Path to the model snapshot directory - - Returns: - Tuple of (all_exist, error_message) - """ - # Find all safetensors index files index_files = [ f for f in os.listdir(snapshot_dir) if f.endswith(".safetensors.index.json") ] @@ -416,7 +304,6 @@ def _check_index_files_exist(snapshot_dir: str) -> Tuple[bool, Optional[str]]: logger.warning( "Removed broken index symlink: %s (blob missing)", index_file ) - # Also try to remove dangling blob reference if it somehow exists if os.path.exists(blob_path): os.remove(blob_path) except Exception as e: @@ -434,7 +321,6 @@ def _check_index_files_exist(snapshot_dir: str) -> Tuple[bool, Optional[str]]: if not weight_map: continue - # Check that all files in weight_map exist required_files = set(weight_map.values()) missing_files = [] @@ -467,18 +353,6 @@ def _check_index_files_exist(snapshot_dir: str) -> Tuple[bool, Optional[str]]: def _validate_sharded_model( snapshot_dir: str, weight_files: List[str] ) -> Tuple[bool, Optional[str], List[str]]: - """ - Validate that all model shards are present and not corrupted. - - Args: - snapshot_dir: Path to the model snapshot directory - weight_files: List of weight file paths - - Returns: - Tuple of (is_valid, error_message, corrupted_files) - - corrupted_files: List of file paths that are corrupted (for selective cleanup) - """ - # First, check if all files from the index actually exist # This catches missing files that wouldn't be found by glob index_check_valid, index_error = _check_index_files_exist(snapshot_dir) if not index_check_valid: @@ -487,7 +361,6 @@ def _validate_sharded_model( # Pattern for sharded files: model-00001-of-00009.safetensors shard_pattern = re.compile(r"(.*?)-(\d+)-of-(\d+)\.(safetensors|bin)") - # Group files by shard pattern (prefix-*-of-N) shard_groups = {} for f in weight_files: base_name = os.path.basename(f) @@ -511,10 +384,8 @@ def _validate_sharded_model( shard_groups[group_key]["found_shards"].append(shard_id) shard_groups[group_key]["files"].append(f) - # Track corrupted files for selective cleanup corrupted_files = [] - # Validate each shard group for group_key, group_info in shard_groups.items(): total_shards = group_info["total"] found_shards = set(group_info["found_shards"]) @@ -523,7 +394,6 @@ def _validate_sharded_model( min_idx = min(found_shards) if found_shards else 1 expected_shards = set(range(min_idx, min_idx + total_shards)) - # Check for missing shards missing_shards = expected_shards - found_shards if missing_shards: return ( @@ -532,7 +402,6 @@ def _validate_sharded_model( [], ) - # Validate weight files for corruption if group_info["suffix"] == "safetensors": for f in group_info["files"]: if not _validate_safetensors_file(f): @@ -542,7 +411,6 @@ def _validate_sharded_model( if not _validate_pytorch_bin_file(f): corrupted_files.append(f) - # Check for required index file for safetensors shards if group_info["suffix"] == "safetensors": index_file = os.path.join( snapshot_dir, f"{group_info['prefix']}.safetensors.index.json" @@ -567,19 +435,6 @@ def _validate_sharded_model( def _cleanup_corrupted_files_selective( model_name_or_path: str, corrupted_files: List[str] ) -> int: - """ - Selectively remove corrupted files and their blobs to force re-download. - - This is more efficient than removing the entire model cache as it only - re-downloads corrupted files rather than the entire model. - - Args: - model_name_or_path: Model identifier - corrupted_files: List of corrupted file paths (symlinks in snapshot) - - Returns: - Number of files successfully cleaned up - """ cleaned_count = 0 for file_path in corrupted_files: @@ -588,13 +443,11 @@ def _cleanup_corrupted_files_selective( if os.path.islink(file_path): blob_path = os.path.realpath(file_path) - # Delete the symlink os.remove(file_path) logger.info( "Removed corrupted symlink: %s", os.path.basename(file_path) ) - # Delete the blob (the actual corrupted data) if os.path.exists(blob_path): os.remove(blob_path) logger.info( @@ -603,7 +456,6 @@ def _cleanup_corrupted_files_selective( cleaned_count += 1 elif os.path.exists(file_path): - # Not a symlink, just delete the file os.remove(file_path) logger.info("Removed corrupted file: %s", os.path.basename(file_path)) cleaned_count += 1 @@ -629,17 +481,7 @@ def _cleanup_corrupted_files_selective( def _cleanup_corrupted_model_cache( model_name_or_path: str, snapshot_dir: str, reason: str ) -> None: - """ - Remove entire corrupted model cache directory to force a clean re-download. - - This is used when we cannot selectively clean (e.g., missing shards, incomplete - downloads with unknown affected files). - - Args: - model_name_or_path: Model identifier - snapshot_dir: Path to the snapshot directory - reason: Reason for cleanup - """ + # Full-cache delete: for when the affected files are unknown. # Navigate up to the model root directory: snapshots/hash -> snapshots -> model_root repo_folder = os.path.abspath(os.path.join(snapshot_dir, "..", "..")) @@ -666,27 +508,10 @@ def ci_validate_and_cleanup_local_snapshot( found_local_snapshot_dir: str, local_weight_files: List[str], ) -> bool: - """ - CI-specific validation and cleanup for local model snapshots. - - This function validates the local snapshot and performs automatic cleanup - if corruption or missing files are detected. This behavior is only appropriate - for CI environments where we want automatic recovery. - - Args: - model_name_or_path: Model identifier for logging - found_local_snapshot_dir: Path to the local snapshot directory - local_weight_files: List of weight file paths found in the snapshot - - Returns: - True if the snapshot is valid and can be used, False if it was invalid - and cleanup was performed (caller should re-download) - """ - # Check for incomplete files and clean up if found + """Validate a local snapshot, cleaning it up on failure; False means re-download.""" 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_module.glob(os.path.join(blobs_dir, "*.incomplete")) @@ -704,14 +529,12 @@ def ci_validate_and_cleanup_local_snapshot( ) return False - # 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 " @@ -722,8 +545,8 @@ def ci_validate_and_cleanup_local_snapshot( return False else: # Missing shards (not corruption) - let snapshot_download handle it. - # IMPORTANT: Do NOT delete the entire cache here, as other processes - # (TP/EP ranks) may already be loading weights from these files. + # Other processes (TP/EP ranks) may already be loading these + # files, so the whole cache must not be deleted here. log_info_on_rank0( logger, f"Validation failed for {model_name_or_path}: {error_msg}. " @@ -731,10 +554,8 @@ def ci_validate_and_cleanup_local_snapshot( ) return False - # Also validate single (non-sharded) weight 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", @@ -747,10 +568,8 @@ def ci_validate_and_cleanup_local_snapshot( 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 False - # Also validate single PyTorch .bin files elif base_name in [ "pytorch_model.bin", "model.bin", @@ -774,23 +593,6 @@ def _validate_weights_after_download( allow_patterns: List[str], model_name_or_path: str, ) -> bool: - """ - Validate downloaded weight files to catch corruption early. - - This function validates safetensors files after download to catch - corruption issues (truncated downloads, network errors, etc.) before - model loading fails with cryptic errors. If corruption is found, - the corrupted files are automatically cleaned up. - - Args: - hf_folder: Path to the downloaded model folder - allow_patterns: Patterns used to match weight files - model_name_or_path: Model identifier for error messages - - Returns: - True if all files are valid, False if corrupted files were found and cleaned up - """ - # Find all weight files that were downloaded weight_files: List[str] = [] for pattern in allow_patterns: weight_files.extend(glob_module.glob(os.path.join(hf_folder, pattern))) @@ -798,7 +600,6 @@ def _validate_weights_after_download( if not weight_files: return True # No weight files to validate - # Validate weight files (safetensors and .bin) corrupted_files = [] for f in weight_files: if f.endswith(".safetensors") and os.path.exists(f): @@ -828,24 +629,6 @@ def _validate_weights_after_download( def _get_lock_file_path( model_name_or_path: str, cache_dir: Optional[str] = None ) -> str: - """ - Generate a unique lock file path for download coordination. - - In CI environments where multiple containers share an NFS-mounted HF cache, - the lock file is placed on the shared cache directory so ALL containers - coordinate on the same lock. This prevents cross-container .incomplete - file race conditions. - - Falls back to /dev/shm (container-local) for non-CI or when the cache - dir is not accessible. - - Args: - model_name_or_path: Model identifier - cache_dir: HF cache directory (None to use default) - - Returns: - Path to the lock file - """ key_hash = hashlib.sha256(model_name_or_path.encode()).hexdigest()[:16] # In CI, place lock on the shared HF cache directory so that ALL containers @@ -862,27 +645,13 @@ def _get_lock_file_path( except Exception: pass - # Fallback to container-local lock if os.path.isdir("/dev/shm"): return f"/dev/shm/sglang_download_lock_{key_hash}" return f"/tmp/sglang_download_lock_{key_hash}" def _cleanup_incomplete_blobs(model_name_or_path: str, cache_dir: Optional[str]) -> int: - """ - Remove stale .incomplete files from the model's blobs directory. - - This is lighter than _cleanup_corrupted_model_cache (which deletes the - entire cache). We only remove .incomplete files so snapshot_download - starts fresh on retry, preserving any successfully downloaded blobs. - - Args: - model_name_or_path: Model identifier (e.g., "meta-llama/Llama-2-7b-hf") - cache_dir: HF cache directory (None to use default) - - Returns: - Number of .incomplete files removed - """ + # Only .incomplete files go, so retries keep the blobs already downloaded. try: import huggingface_hub.constants @@ -929,30 +698,10 @@ def ci_download_with_validation_and_retry( revision: Optional[str], max_retries: int = 3, ) -> str: - """ - CI-specific download with validation and automatic retry on corruption. + """Download weights, validating each attempt and retrying on corruption. - This function handles the download of model weights in CI environments, - with automatic validation and retry logic for handling corrupted downloads. - - Uses filelock.FileLock on the shared HF cache directory to coordinate - downloads across all processes AND all containers sharing the same - NFS-mounted cache. Only one process downloads at a time; others wait - for the lock then use the cached result. - - Args: - model_name_or_path: The model name or path - allow_patterns: The allowed patterns for weight files - ignore_patterns: The patterns to filter out weight files - cache_dir: The cache directory to store model weights - revision: The revision of the model - max_retries: Maximum number of download retries if corruption is detected - - Returns: - str: The path to the downloaded model weights - - Raises: - RuntimeError: If download fails after max_retries attempts + Holds a filelock on the shared HF cache so that processes and containers on + the same NFS mount take turns; the rest wait and reuse the cached result. """ import filelock import huggingface_hub.constants @@ -1104,18 +853,10 @@ def ci_download_with_validation_and_retry( def ci_validate_and_clean_hf_cache(model_path: str) -> None: - """ - Validate and clean corrupted safetensors files in HF cache before loading. + """Drop corrupted safetensors from the HF cache before a non-SGLang load. - This function is needed because HFRunner (used in tests) calls transformers' - from_pretrained() directly, which bypasses SGLang's weight validation. - Corrupted cached files can cause cryptic errors like "EOF while parsing" - from safetensors. - - Only runs in CI to avoid overhead for regular users. - - Args: - model_path: Model identifier (e.g., "meta-llama/Llama-2-7b") + HFRunner calls transformers' from_pretrained() directly, which bypasses the + validation in this module; a corrupted cache surfaces as "EOF while parsing". """ from sglang.utils import is_in_ci