Support pre-generating and using expected checksums (#16730)
This commit is contained in:
fzyzcjy
@@ -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():
|
||||
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__":
|
||||
|
||||
@@ -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 ========
|
||||
|
||||
Reference in New Issue
Block a user