diff --git a/ramalama/model_store/store.py b/ramalama/model_store/store.py index 65382e290..1965fe6f4 100644 --- a/ramalama/model_store/store.py +++ b/ramalama/model_store/store.py @@ -422,7 +422,7 @@ def _get_refcounts(self, snapshot_hash: str) -> tuple[int, Counter[str]]: ] refs = [ref for tag in model_tags if (ref := self.get_ref_file(tag))] - blob_refcounts = Counter(file.name for ref in refs for file in ref.files) + blob_refcounts = Counter(file_hash for ref in refs for file_hash in {file.hash for file in ref.files}) snap_refcount = sum(ref.hash == snapshot_hash for ref in refs) @@ -438,7 +438,7 @@ def remove_snapshot(self, model_tag: str) -> bool: # Remove all blobs first for file in ref_file.files: - blob_refcount = blob_refcounts.get(file.name, 0) + blob_refcount = blob_refcounts.get(file.hash, 0) if blob_refcount <= 1: blob_absolute_path = Path(self.get_blob_file_path(file.hash)) self._remove_blob_path(blob_absolute_path) diff --git a/test/unit/test_model_store.py b/test/unit/test_model_store.py index 81594bc50..6a7fa8231 100644 --- a/test/unit/test_model_store.py +++ b/test/unit/test_model_store.py @@ -1,4 +1,5 @@ import os +from pathlib import Path import pytest @@ -67,7 +68,9 @@ def test_model_factory_create(input: list[SnapshotFile], expect_error: bool): validate_snapshot_files(input) -def test_try_convert_existing_chat_template_converts_flat_jinja(tmp_path, monkeypatch): +def test_try_convert_existing_chat_template_converts_flat_jinja( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: base_path = tmp_path global_store = GlobalModelStore(str(base_path)) model_store = ModelStore(global_store, model_name="sample", model_type="file", model_organization="org") @@ -118,7 +121,7 @@ def fake_update_snapshot(ref_file_arg, snapshot_hash_arg, files): assert converted_file.content == wrap_template_with_messages_loop(original_template).encode("utf-8") -def test_local_snapshot_file_binary_download_and_digest(tmp_path): +def test_local_snapshot_file_binary_download_and_digest(tmp_path: Path) -> None: # Use a payload that includes a null byte to ensure we are truly treating this as binary data content = b"binary-\x00-test-content" expected_digest = generate_sha256_binary(content) @@ -140,3 +143,79 @@ def test_local_snapshot_file_binary_download_and_digest(tmp_path): # Assert: digest matches generate_sha256_binary(content) assert snapshot_file.hash == expected_digest + + +def test_remove_snapshot_refcounts_blobs_by_hash(tmp_path: Path) -> None: + global_store = GlobalModelStore(str(tmp_path)) + model_store = ModelStore(global_store, model_name="sample", model_type="file", model_organization="org") + model_store.ensure_directory_setup() + + shared_hash = "shared-blob" + RefJSONFile( + hash="snap-a", + path=model_store.get_ref_file_path("a"), + files=[StoreFile(shared_hash, "model-a.gguf", StoreFileType.GGUF_MODEL)], + ).write_to_file() + RefJSONFile( + hash="snap-b", + path=model_store.get_ref_file_path("b"), + files=[StoreFile(shared_hash, "renamed-model.gguf", StoreFileType.GGUF_MODEL)], + ).write_to_file() + + shared_blob_path = model_store.get_blob_file_path(shared_hash) + with open(shared_blob_path, "w") as blob_file: + blob_file.write("shared") + + assert model_store.remove_snapshot("a") is True + assert os.path.exists(shared_blob_path) + + +def test_remove_snapshot_does_not_refcount_different_hashes_with_same_name(tmp_path: Path) -> None: + global_store = GlobalModelStore(str(tmp_path)) + model_store = ModelStore(global_store, model_name="sample", model_type="file", model_organization="org") + model_store.ensure_directory_setup() + + RefJSONFile( + hash="snap-a", + path=model_store.get_ref_file_path("a"), + files=[StoreFile("blob-a", "model.gguf", StoreFileType.GGUF_MODEL)], + ).write_to_file() + RefJSONFile( + hash="snap-b", + path=model_store.get_ref_file_path("b"), + files=[StoreFile("blob-b", "model.gguf", StoreFileType.GGUF_MODEL)], + ).write_to_file() + + blob_a_path = model_store.get_blob_file_path("blob-a") + blob_b_path = model_store.get_blob_file_path("blob-b") + with open(blob_a_path, "w") as blob_file: + blob_file.write("a") + with open(blob_b_path, "w") as blob_file: + blob_file.write("b") + + assert model_store.remove_snapshot("a") is True + assert not os.path.exists(blob_a_path) + assert os.path.exists(blob_b_path) + + +def test_remove_snapshot_counts_duplicate_blob_hash_once_per_ref_file(tmp_path: Path) -> None: + global_store = GlobalModelStore(str(tmp_path)) + model_store = ModelStore(global_store, model_name="sample", model_type="file", model_organization="org") + model_store.ensure_directory_setup() + + duplicate_hash = "duplicate-blob" + RefJSONFile( + hash="snap-a", + path=model_store.get_ref_file_path("a"), + files=[ + StoreFile(duplicate_hash, "model.gguf", StoreFileType.GGUF_MODEL), + StoreFile(duplicate_hash, "renamed-model.gguf", StoreFileType.GGUF_MODEL), + ], + ).write_to_file() + + duplicate_blob_path = model_store.get_blob_file_path(duplicate_hash) + with open(duplicate_blob_path, "w") as blob_file: + blob_file.write("duplicate") + + assert model_store.remove_snapshot("a") is True + assert not os.path.exists(duplicate_blob_path)