From 83abecd0c27cb9695a4df074705872e781681aaf Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Thu, 8 Jan 2026 20:26:13 +0800 Subject: [PATCH] Support pre-generating and using expected checksums (#16730) --- .../sglang/srt/utils/model_file_verifier.py | 102 ++++++++++-- .../utils/test_model_file_verifier.py | 156 +++++++++++++++++- 2 files changed, 238 insertions(+), 20 deletions(-) diff --git a/python/sglang/srt/utils/model_file_verifier.py b/python/sglang/srt/utils/model_file_verifier.py index 24fbc999d..535d2a941 100644 --- a/python/sglang/srt/utils/model_file_verifier.py +++ b/python/sglang/srt/utils/model_file_verifier.py @@ -1,12 +1,19 @@ """ Model File Verifier - Verify model file integrity using SHA256 checksums. -Example command: +Example commands: + # Verify using HuggingFace model online metadata python -m sglang.srt.utils.model_file_verifier verify --model-path /path/to/model --model-checksum Qwen/Qwen3-0.6B + + # Verify using locally generated checksum + python -m sglang.srt.utils.model_file_verifier generate --model-path --model-checksum checksums.json + python -m sglang.srt.utils.model_file_verifier verify --model-path /path/to/model --model-checksum checksums.json """ import argparse +import fnmatch import hashlib +import json from concurrent.futures import ThreadPoolExecutor from pathlib import Path from typing import Dict, List, Optional, Tuple @@ -29,7 +36,7 @@ IGNORE_PATTERNS = [ def verify(*, model_path: str, checksums_source: str, max_workers: int = 4) -> None: model_path = Path(model_path).resolve() - expected = _load_checksums_from_hf(repo_id=checksums_source) + expected = _load_checksums(checksums_source) actual = _compute_checksums_from_folder( model_path=model_path, filenames=list(expected.keys()), max_workers=max_workers ) @@ -51,9 +58,52 @@ def _compare_checksums(*, expected: Dict[str, str], actual: Dict[str, str]) -> N raise IntegrityError("Integrity check failed: " + "; ".join(errors)) +# ======== Generate ======== + + +def generate_checksums( + *, source: str, output_path: str, max_workers: int = 4 +) -> Dict[str, str]: + if Path(source).is_dir(): + model_path = Path(source).resolve() + files = _discover_files(model_path) + if not files: + raise IntegrityError(f"No model files found in {model_path}") + checksums = _compute_checksums_from_folder( + model_path=model_path, filenames=files, max_workers=max_workers + ) + else: + checksums = _load_checksums_from_hf(repo_id=source) + + output = {"checksums": checksums} + Path(output_path).write_text(json.dumps(output, indent=2, sort_keys=True)) + + print( + f"[ModelFileVerifier] Generated checksums for {len(checksums)} files -> {output_path}" + ) + return checksums + + +def _discover_files(model_path: Path) -> List[str]: + return sorted( + e.name + for e in model_path.iterdir() + if e.is_file() + and not e.name.startswith(".") + and not any(fnmatch.fnmatch(e.name, p) for p in IGNORE_PATTERNS) + ) + + # ======== Load Checksums ======== +def _load_checksums(source: str) -> Dict[str, str]: + if Path(source).is_file(): + data = json.loads(Path(source).read_text()) + return data["checksums"] + return _load_checksums_from_hf(repo_id=source) + + def _load_checksums_from_hf(*, repo_id: str) -> Dict[str, str]: from huggingface_hub import HfFileSystem @@ -72,8 +122,6 @@ def _load_checksums_from_hf(*, repo_id: str) -> Dict[str, str]: def _get_filename_and_checksum_from_hf_file(fs, file_info): - import fnmatch - if file_info.get("type") != "file": return None @@ -137,30 +185,54 @@ class IntegrityError(Exception): # ======== CLI ======== -def main(): - parser = argparse.ArgumentParser( - description="Model File Verifier - Verify model file integrity using checksums" - ) +def _add_common_args(parser): parser.add_argument( "--model-path", required=True, - help="Local model directory", + help="Local model directory or HuggingFace repo ID", ) parser.add_argument( "--model-checksum", required=True, - help="HuggingFace repo ID for checksums", + help="Checksums JSON file path", ) parser.add_argument( "--workers", type=int, default=4, help="Number of parallel workers" ) - args = parser.parse_args() - verify( - model_path=args.model_path, - checksums_source=args.model_checksum, - max_workers=args.workers, + +def main(): + parser = argparse.ArgumentParser( + description="Model File Verifier - Verify model file integrity using checksums" ) + subparsers = parser.add_subparsers(dest="command", required=True) + + gen_parser = subparsers.add_parser( + "generate", help="Generate checksums.json for a model" + ) + _add_common_args(gen_parser) + gen_parser.set_defaults( + func=lambda args: generate_checksums( + source=args.model_path, + output_path=args.model_checksum, + max_workers=args.workers, + ) + ) + + verify_parser = subparsers.add_parser( + "verify", help="Verify model files against checksums" + ) + _add_common_args(verify_parser) + verify_parser.set_defaults( + func=lambda args: verify( + model_path=args.model_path, + checksums_source=args.model_checksum, + max_workers=args.workers, + ) + ) + + args = parser.parse_args() + args.func(args) if __name__ == "__main__": diff --git a/test/registered/utils/test_model_file_verifier.py b/test/registered/utils/test_model_file_verifier.py index f8ed0747a..9d7aa678b 100644 --- a/test/registered/utils/test_model_file_verifier.py +++ b/test/registered/utils/test_model_file_verifier.py @@ -1,6 +1,9 @@ import hashlib +import json import os import shutil +import subprocess +import sys import tempfile import unittest from contextlib import nullcontext @@ -10,7 +13,12 @@ import requests from huggingface_hub import snapshot_download from sglang.srt.utils import kill_process_tree -from sglang.srt.utils.model_file_verifier import compute_sha256, verify +from sglang.srt.utils.model_file_verifier import ( + IntegrityError, + compute_sha256, + generate_checksums, + verify, +) from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.test_utils import ( DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, @@ -62,6 +70,30 @@ class _RealModelTestCase(unittest.TestCase): class TestModelFileVerifier(_FakeModelTestCase): + def test_detect_bit_rot(self): + checksums_file = os.path.join(self.test_dir, "checksums.json") + generate_checksums(source=self.test_dir, output_path=checksums_file) + + target_file = os.path.join(self.test_dir, "model.safetensors") + _flip_bit_in_file(target_file, byte_offset=50, bit_position=3) + + with self.assertRaises(IntegrityError) as ctx: + verify(model_path=self.test_dir, checksums_source=checksums_file) + + self.assertIn("model.safetensors", str(ctx.exception)) + self.assertIn("mismatch", str(ctx.exception).lower()) + + def test_detect_missing_file(self): + checksums_file = os.path.join(self.test_dir, "checksums.json") + generate_checksums(source=self.test_dir, output_path=checksums_file) + + os.remove(os.path.join(self.test_dir, "config.json")) + + with self.assertRaises(IntegrityError) as ctx: + verify(model_path=self.test_dir, checksums_source=checksums_file) + + self.assertIn("config.json", str(ctx.exception)) + def test_compute_sha256(self): test_file = os.path.join(self.test_dir, "test.bin") content = b"hello world" @@ -72,12 +104,113 @@ class TestModelFileVerifier(_FakeModelTestCase): expected = hashlib.sha256(content).hexdigest() self.assertEqual(result, expected) + def test_parallel_checksum_computation(self): + for i in range(10): + _create_test_file( + self.test_dir, f"shard_{i}.safetensors", f"content_{i}".encode() * 1000 + ) + + checksums_file = os.path.join(self.test_dir, "checksums.json") + checksums = generate_checksums( + source=self.test_dir, output_path=checksums_file, max_workers=4 + ) + + self.assertGreaterEqual(len(checksums), 10) + + +# ======== CLI Tests ======== + + +class TestModelFileVerifierCLI(_FakeModelTestCase): + + def test_cli_generate(self): + checksums_file = os.path.join(self.test_dir, "checksums.json") + result = subprocess.run( + [ + sys.executable, + "-m", + "sglang.srt.utils.model_file_verifier", + "generate", + "--model-path", + self.test_dir, + "--model-checksum", + checksums_file, + ], + capture_output=True, + text=True, + ) + self.assertEqual(result.returncode, 0, f"stderr: {result.stderr}") + self.assertTrue(os.path.exists(checksums_file)) + + with open(checksums_file) as f: + data = json.load(f) + self.assertIn("checksums", data) + self.assertEqual(len(data["checksums"]), 3) + + def test_cli_verify_success(self): + checksums_file = os.path.join(self.test_dir, "checksums.json") + generate_checksums(source=self.test_dir, output_path=checksums_file) + + result = subprocess.run( + [ + sys.executable, + "-m", + "sglang.srt.utils.model_file_verifier", + "verify", + "--model-path", + self.test_dir, + "--model-checksum", + checksums_file, + ], + capture_output=True, + text=True, + ) + self.assertEqual(result.returncode, 0, f"stderr: {result.stderr}") + self.assertIn("verified successfully", result.stdout) + + def test_cli_verify_fails_on_corruption(self): + checksums_file = os.path.join(self.test_dir, "checksums.json") + generate_checksums(source=self.test_dir, output_path=checksums_file) + + target_file = os.path.join(self.test_dir, "model.safetensors") + _flip_bit_in_file(target_file, byte_offset=50, bit_position=3) + + result = subprocess.run( + [ + sys.executable, + "-m", + "sglang.srt.utils.model_file_verifier", + "verify", + "--model-path", + self.test_dir, + "--model-checksum", + checksums_file, + ], + capture_output=True, + text=True, + ) + self.assertNotEqual(result.returncode, 0) + combined = result.stdout + result.stderr + self.assertTrue( + "IntegrityError" in combined or "mismatch" in combined.lower(), + f"Expected integrity error, got: {combined}", + ) + # ======== HuggingFace Tests ======== class TestModelFileVerifierHF(_RealModelTestCase): + def test_generate_checksums_from_hf(self): + checksums_file = os.path.join(self.test_dir, "checksums.json") + checksums = generate_checksums(source=MODEL_NAME, output_path=checksums_file) + + self.assertTrue(os.path.exists(checksums_file)) + self.assertGreater(len(checksums), 0) + for filename, sha256 in checksums.items(): + self.assertEqual(len(sha256), 64) + def test_verify_with_hf_checksums_source(self): verify(model_path=self.test_dir, checksums_source=MODEL_NAME) @@ -87,7 +220,14 @@ class TestModelFileVerifierHF(_RealModelTestCase): class TestModelFileVerifierWithRealModel(_RealModelTestCase): - def _run_server_test(self, *, corrupt_weights: bool): + def _run_server_test(self, *, corrupt_weights: bool, use_hf_checksum: bool): + if use_hf_checksum: + checksum_arg = MODEL_NAME + else: + checksums_file = os.path.join(self.test_dir, "checksums.json") + generate_checksums(source=self.test_dir, output_path=checksums_file) + checksum_arg = checksums_file + corrupted_file = None if corrupt_weights: safetensors_files = [ @@ -104,7 +244,7 @@ class TestModelFileVerifierWithRealModel(_RealModelTestCase): model=self.test_dir, base_url=DEFAULT_URL_FOR_TEST, timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - other_args=["--model-checksum", MODEL_NAME], + other_args=["--model-checksum", checksum_arg], return_stdout_stderr=(stdout_io, stderr_io), ) @@ -124,10 +264,16 @@ class TestModelFileVerifierWithRealModel(_RealModelTestCase): kill_process_tree(process.pid) def test_server_launch_with_checksum_intact(self): - self._run_server_test(corrupt_weights=False) + self._run_server_test(corrupt_weights=False, use_hf_checksum=False) def test_server_launch_fails_with_corrupted_weights(self): - self._run_server_test(corrupt_weights=True) + self._run_server_test(corrupt_weights=True, use_hf_checksum=False) + + def test_server_launch_with_hf_checksum_intact(self): + self._run_server_test(corrupt_weights=False, use_hf_checksum=True) + + def test_server_launch_with_hf_checksum_corrupted(self): + self._run_server_test(corrupt_weights=True, use_hf_checksum=True) # ======== Test Utilities ========