diff --git a/python/sglang/srt/model_loader/weight_utils.py b/python/sglang/srt/model_loader/weight_utils.py index 67156b16d..16babccaf 100644 --- a/python/sglang/srt/model_loader/weight_utils.py +++ b/python/sglang/srt/model_loader/weight_utils.py @@ -686,6 +686,16 @@ def filter_duplicate_safetensors_files( weight_files_in_index = set() for weight_name in weight_map: weight_files_in_index.add(os.path.join(hf_folder, weight_map[weight_name])) + # Fail fast if the index references shard files that are not on disk (e.g. an + # incomplete or interrupted download). Otherwise those shards are silently + # dropped and the model loads with uninitialized weights. + missing_files = sorted(f for f in weight_files_in_index if not os.path.isfile(f)) + if missing_files: + raise RuntimeError( + f"{index_file} references {len(missing_files)} shard file(s) missing " + f"from {hf_folder} (incomplete download?): " + f"{[os.path.basename(f) for f in missing_files]}" + ) # Filter out any fields that are not found in the index file. hf_weights_files = [f for f in hf_weights_files if f in weight_files_in_index] return hf_weights_files diff --git a/test/registered/unit/model_loader/test_weight_utils.py b/test/registered/unit/model_loader/test_weight_utils.py new file mode 100644 index 000000000..0e4ba5ed0 --- /dev/null +++ b/test/registered/unit/model_loader/test_weight_utils.py @@ -0,0 +1,88 @@ +"""Unit tests for srt/model_loader/weight_utils.py shard-index consistency.""" + +import json +import os +import tempfile +import unittest + +from sglang.srt.model_loader.weight_utils import filter_duplicate_safetensors_files +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + +INDEX_NAME = "model.safetensors.index.json" + + +def _write_index(folder, weight_map): + with open(os.path.join(folder, INDEX_NAME), "w") as f: + json.dump({"weight_map": weight_map}, f) + + +def _touch(folder, name): + path = os.path.join(folder, name) + open(path, "w").close() + return path + + +class TestFilterDuplicateSafetensorsFiles(CustomTestCase): + def setUp(self): + self._tmp = tempfile.TemporaryDirectory() + self.folder = self._tmp.name + + def tearDown(self): + self._tmp.cleanup() + + def test_missing_shard_raises(self): + # Index lists two shards, only one on disk (interrupted download). + _write_index( + self.folder, + { + "w1": "model-00001-of-00002.safetensors", + "w2": "model-00002-of-00002.safetensors", + }, + ) + present = _touch(self.folder, "model-00001-of-00002.safetensors") + + with self.assertRaises(RuntimeError) as cm: + filter_duplicate_safetensors_files( + hf_weights_files=[present], + hf_folder=self.folder, + index_file=INDEX_NAME, + ) + self.assertIn("model-00002-of-00002.safetensors", str(cm.exception)) + + def test_complete_checkpoint_filters_non_indexed(self): + # All indexed shards present; a non-indexed duplicate is still filtered out. + _write_index( + self.folder, + { + "w1": "model-00001-of-00002.safetensors", + "w2": "model-00002-of-00002.safetensors", + }, + ) + shard1 = _touch(self.folder, "model-00001-of-00002.safetensors") + shard2 = _touch(self.folder, "model-00002-of-00002.safetensors") + extra = _touch(self.folder, "consolidated.safetensors") + + result = filter_duplicate_safetensors_files( + hf_weights_files=[shard1, shard2, extra], + hf_folder=self.folder, + index_file=INDEX_NAME, + ) + self.assertEqual(sorted(result), sorted([shard1, shard2])) + + def test_single_file_model_no_index_returns_unchanged(self): + # No index on disk (single-file / dummy / object-storage): early return. + single = _touch(self.folder, "model.safetensors") + + result = filter_duplicate_safetensors_files( + hf_weights_files=[single], + hf_folder=self.folder, + index_file=INDEX_NAME, + ) + self.assertEqual(result, [single]) + + +if __name__ == "__main__": + unittest.main()