From 4a3f34cd8c0f221cf12a2bc594610165dbce18db Mon Sep 17 00:00:00 2001 From: Boaz Shuster Date: Wed, 16 Sep 2026 22:03:18 +0300 Subject: [PATCH] fix: refcount blobs by hash when removing snapshots Snapshot removal used file names when counting blob references. If two refs pointed at the same blob hash under different file names, removing one ref could delete the shared blob while the other ref still referenced it. This patch counts blob references by hash instead, matching how blobs are stored on disk. A script to reproduce the bug: ```bash STORE="$(mktemp -d)" BASE="$STORE/store/huggingface/acme/shared-model" mkdir -p "$BASE/refs" "$BASE/blobs" "$BASE/snapshots" echo "shared model bytes" > "$BASE/blobs/sha256-shared" cat > "$BASE/refs/A.json" < "$BASE/refs/B.json" < --- ramalama/model_store/store.py | 4 +- test/unit/test_model_store.py | 83 ++++++++++++++++++++++++++++++++++- 2 files changed, 83 insertions(+), 4 deletions(-) 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)