Support pre-generating and using expected checksums (#16730)

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 <hf-id-or-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():
+def _add_common_args(parser):
-    parser = argparse.ArgumentParser(
-        description="Model File Verifier - Verify model file integrity using checksums"
-    )
     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(
+def main():
-        model_path=args.model_path,
+    parser = argparse.ArgumentParser(
-        checksums_source=args.model_checksum,
+        description="Model File Verifier - Verify model file integrity using checksums"
-        max_workers=args.workers,
     )
+    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__":

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 ========