Fix modelstore deleting logic when multiple reference refer to the same blob/snapshot

ramalama/model_store.py

@@ -1,6 +1,7 @@
 import os
 import shutil
 import urllib.error
+from collections import Counter
 from dataclasses import dataclass
 from datetime import datetime
 from enum import IntEnum
@@ -50,12 +51,15 @@ def __init__(
         self.required: bool = required
 
     def download(self, blob_file_path: str, snapshot_dir: str) -> str:
-        download_file(
-            url=self.url,
-            headers=self.header,
-            dest_path=blob_file_path,
-            show_progress=self.should_show_progress,
-        )
+        if not os.path.exists(blob_file_path):
+            download_file(
+                url=self.url,
+                headers=self.header,
+                dest_path=blob_file_path,
+                show_progress=self.should_show_progress,
+            )
+        else:
+            logger.debug(f"Using cached blob for {self.name} ({os.path.basename(blob_file_path)})")
         return os.path.relpath(blob_file_path, start=snapshot_dir)
 
 
@@ -443,7 +447,12 @@ def _download_snapshot_files(self, model_tag: str, snapshot_hash: str, snapshot_
                     if not verify_checksum(dest_path):
                         raise ValueError(f"Checksum verification failed for blob {dest_path}")
 
-            os.symlink(blob_relative_path, self.get_snapshot_file_path(snapshot_hash, file.name))
+            link_path = self.get_snapshot_file_path(snapshot_hash, file.name)
+            try:
+                os.symlink(blob_relative_path, link_path)
+            except FileExistsError:
+                os.unlink(link_path)
+                os.symlink(blob_relative_path, link_path)
 
         # save updated ref file
         ref_file.write_to_file()
@@ -589,18 +598,41 @@ def _remove_blob_file(self, snapshot_file_path: str):
         except Exception as ex:
             logger.error(f"Failed to remove blob file '{blob_path}': {ex}")
 
+    def _get_refcounts(self, snapshot_hash: str) -> tuple[int, Counter[str]]:
+        ref_paths = Path(self.refs_directory).iterdir()
+        refs = [RefFile.from_path(str(p)) for p in ref_paths]
+
+        blob_refcounts = Counter(filename for ref in refs for filename in ref.filenames)
+
+        snap_refcount = sum(ref.hash == snapshot_hash for ref in refs)
+
+        return snap_refcount, blob_refcounts
+
     def remove_snapshot(self, model_tag: str):
         ref_file = self.get_ref_file(model_tag)
 
+        if ref_file is None:
+            return
+
+        snapshot_refcount, blob_refcounts = self._get_refcounts(ref_file.hash)
+
         # Remove all blobs first
-        if ref_file is not None:
-            for file in ref_file.filenames:
+        for file in ref_file.filenames:
+            blob_refcount = blob_refcounts.get(file, 0)
+            if blob_refcount <= 1:
                 self._remove_blob_file(self.get_snapshot_file_path(ref_file.hash, file))
-                self._remove_blob_file(self.get_partial_blob_file_path(ref_file.hash))
+            else:
+                logger.debug(f"Not removing blob {file} refcount={blob_refcount}")
 
         # Remove snapshot directory
-        snapshot_directory = self.get_snapshot_directory_from_tag(model_tag)
-        shutil.rmtree(snapshot_directory, ignore_errors=True)
+        if snapshot_refcount <= 1:
+            # FIXME: this only cleans up .partial files where the blob hash equals the snapshot hash
+            self._remove_blob_file(self.get_partial_blob_file_path(ref_file.hash))
+            snapshot_directory = self.get_snapshot_directory_from_tag(model_tag)
+            shutil.rmtree(snapshot_directory, ignore_errors=True)
+            logger.debug(f"Snapshot removed {ref_file.hash}")
+        else:
+            logger.debug(f"Not removing snapshot {ref_file.hash} refcount={snapshot_refcount}")
 
         # Remove ref file, ignore if file is not found
         ref_file_path = self.get_ref_file_path(model_tag)

test/system/050-pull.bats

@@ -82,6 +82,31 @@ load setup_suite
     run_ramalama list
     is "$output" ".*TinyLlama/TinyLlama-1.1B-Chat-v1.0" "image was actually pulled locally"
     run_ramalama rm huggingface://TinyLlama/TinyLlama-1.1B-Chat-v1.0
+
+    run_ramalama pull hf://ggml-org/SmolVLM-256M-Instruct-GGUF
+    run_ramalama list
+    is "$output" ".*ggml-org/SmolVLM-256M-Instruct-GGUF" "image was actually pulled locally"
+    run_ramalama rm huggingface://ggml-org/SmolVLM-256M-Instruct-GGUF
+
+    run_ramalama pull hf://ggml-org/SmolVLM-256M-Instruct-GGUF:Q8_0
+    run_ramalama list
+    is "$output" ".*ggml-org/SmolVLM-256M-Instruct-GGUF:Q8_0" "image was actually pulled locally"
+    run_ramalama rm huggingface://ggml-org/SmolVLM-256M-Instruct-GGUF:Q8_0
+}
+
+# bats test_tags=distro-integration
+@test "ramalama pull huggingface tag multiple references" {
+    run_ramalama pull hf://ggml-org/SmolVLM-256M-Instruct-GGUF
+    run_ramalama list
+    is "$output" ".*ggml-org/SmolVLM-256M-Instruct-GGUF" "image was actually pulled locally"
+    run_ramalama --debug pull hf://ggml-org/SmolVLM-256M-Instruct-GGUF:Q8_0
+    is "$output" ".*Using cached blob" "cached blob was used"
+    run_ramalama list
+    is "$output" ".*ggml-org/SmolVLM-256M-Instruct-GGUF:Q8_0" "reference was created to existing image"
+    run_ramalama --debug rm huggingface://ggml-org/SmolVLM-256M-Instruct-GGUF
+    is "$output" ".*Not removing snapshot" "snapshot with remaining reference was not deleted"
+    run_ramalama --debug rm huggingface://ggml-org/SmolVLM-256M-Instruct-GGUF:Q8_0
+    is "$output" ".*Snapshot removed" "snapshot with no remaining references was deleted"
 }
 
 # bats test_tags=distro-integration