diff --git a/tests/model_executor/model_loader/test_filter_duplicate_safetensors.py b/tests/model_executor/model_loader/test_filter_duplicate_safetensors.py new file mode 100644 index 0000000000..babc6c061a --- /dev/null +++ b/tests/model_executor/model_loader/test_filter_duplicate_safetensors.py @@ -0,0 +1,79 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import json +import os +import tempfile + +import pytest + +from vllm.model_executor.model_loader.weight_utils import ( + filter_duplicate_safetensors_files, +) + + +def test_filter_duplicate_safetensors_files_missing_weight(): + with tempfile.TemporaryDirectory() as tmpdir: + existing_file = os.path.join(tmpdir, "model-00001-of-00002.safetensors") + with open(existing_file, "wb") as f: + f.write(b"") + + existing_file2 = os.path.join(tmpdir, "model-00002-of-00002.safetensors") + with open(existing_file2, "wb") as f: + f.write(b"") + + index_file = os.path.join(tmpdir, "model.safetensors.index.json") + index_content = { + "weight_map": { + "layer.0.weight": "model-00001-of-00002.safetensors", + "layer.1.weight": "model-00002-of-00002.safetensors", + "layer.2.weight": "model-00003-of-00002.safetensors", + } + } + with open(index_file, "w") as f: + json.dump(index_content, f) + + hf_weights_files = [ + os.path.join(tmpdir, "model-00001-of-00002.safetensors"), + os.path.join(tmpdir, "model-00002-of-00002.safetensors"), + ] + + with pytest.raises(FileNotFoundError) as exc_info: + filter_duplicate_safetensors_files( + hf_weights_files=hf_weights_files, + hf_folder=tmpdir, + index_file="model.safetensors.index.json", + ) + + assert "model-00003-of-00002.safetensors" in str(exc_info.value) + + +def test_filter_duplicate_safetensors_files_all_exist(): + with tempfile.TemporaryDirectory() as tmpdir: + existing_files = [] + for i in range(1, 3): + file_path = os.path.join(tmpdir, f"model-0000{i}-of-00002.safetensors") + with open(file_path, "wb") as f: + f.write(b"") + existing_files.append(file_path) + + index_file = os.path.join(tmpdir, "model.safetensors.index.json") + index_content = { + "weight_map": { + "layer.0.weight": "model-00001-of-00002.safetensors", + "layer.1.weight": "model-00002-of-00002.safetensors", + } + } + with open(index_file, "w") as f: + json.dump(index_content, f) + + filter_duplicate_safetensors_files( + hf_weights_files=existing_files, + hf_folder=tmpdir, + index_file="model.safetensors.index.json", + ) + + +if __name__ == "__main__": + test_filter_duplicate_safetensors_files_missing_weight() + test_filter_duplicate_safetensors_files_all_exist() diff --git a/vllm/model_executor/model_loader/weight_utils.py b/vllm/model_executor/model_loader/weight_utils.py index 47c6c02be6..ef25a156d9 100644 --- a/vllm/model_executor/model_loader/weight_utils.py +++ b/vllm/model_executor/model_loader/weight_utils.py @@ -595,6 +595,14 @@ 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])) + # Check if files referenced in model.safetensors.index.json actually exist. + # Raise error if any file is missing. + hf_weights_files_set = set(hf_weights_files) + missing_files = weight_files_in_index - hf_weights_files_set + if missing_files: + raise FileNotFoundError( + f"Weight files referenced in index but missing: {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