From 09a76f0dde65ad371a975c393cafb4225ce48be3 Mon Sep 17 00:00:00 2001 From: Kush Gupta Date: Mon, 10 Feb 2025 16:55:06 -0500 Subject: [PATCH 1/4] working ollama cache implementation Signed-off-by: Kush Gupta --- ramalama/ollama.py | 44 +++++++++++++++++++++++++++++++++----------- 1 file changed, 33 insertions(+), 11 deletions(-) diff --git a/ramalama/ollama.py b/ramalama/ollama.py index e2435124..6e35a3b0 100644 --- a/ramalama/ollama.py +++ b/ramalama/ollama.py @@ -1,7 +1,7 @@ import os import urllib.request import json -from ramalama.common import run_cmd, verify_checksum, download_file +from ramalama.common import run_cmd, verify_checksum, download_file, available from ramalama.model import Model, rm_until_substring @@ -30,15 +30,18 @@ def pull_blob(repos, layer_digest, accept, registry_head, models, model_name, mo layer_blob_path = os.path.join(repos, "blobs", layer_digest) url = f"{registry_head}/blobs/{layer_digest}" headers = {"Accept": accept} - download_file(url, layer_blob_path, headers=headers, show_progress=True) - - # Verify checksum after downloading the blob - if not verify_checksum(layer_blob_path): - print(f"Checksum mismatch for blob {layer_blob_path}, retrying download...") - os.remove(layer_blob_path) + local_blob = in_existing_cache(model_name, model_tag) + if local_blob is not None: + run_cmd(["ln", "-sf", local_blob, layer_blob_path]) + else: download_file(url, layer_blob_path, headers=headers, show_progress=True) + # Verify checksum after downloading the blob if not verify_checksum(layer_blob_path): - raise ValueError(f"Checksum verification failed for blob {layer_blob_path}") + print(f"Checksum mismatch for blob {layer_blob_path}, retrying download...") + os.remove(layer_blob_path) + download_file(url, layer_blob_path, headers=headers, show_progress=True) + if not verify_checksum(layer_blob_path): + raise ValueError(f"Checksum verification failed for blob {layer_blob_path}") os.makedirs(models, exist_ok=True) relative_target_path = os.path.relpath(layer_blob_path, start=os.path.dirname(model_path)) @@ -58,11 +61,30 @@ def init_pull(repos, accept, registry_head, model_name, model_tag, models, model return model_path +def in_existing_cache(model_name, model_tag): + if not available("ollama"): + return None + default_ollama_caches=[ + os.path.join(os.environ['HOME'], '.ollama/models'), + '/usr/share/ollama/.ollama/models' + f'C:\\Users\\{os.getlogin()}\\.ollama\\models' + ] + + for cache_dir in default_ollama_caches: + manifest_path = os.path.join(cache_dir, 'manifests', 'registry.ollama.ai', model_name, model_tag) + if os.path.exists(manifest_path): + with open(manifest_path, 'r') as file: + manifest_data = json.load(file) + for layer in manifest_data["layers"]: + if layer["mediaType"] == "application/vnd.ollama.image.model": + layer_digest = layer["digest"] + ollama_digest_path = os.path.join(cache_dir, 'blobs', layer_digest) + return str(ollama_digest_path).replace(':','-') + return None + class Ollama(Model): def __init__(self, model): - model = rm_until_substring(model, "ollama.com/library/") - model = rm_until_substring(model, "://") - super().__init__(model) + super().__init__(model.removeprefix("ollama://")) self.type = "Ollama" def _local(self, args): From cebe5cb8a63c8524f5a2ead1037821562a77fed8 Mon Sep 17 00:00:00 2001 From: Kush Gupta Date: Mon, 10 Feb 2025 17:05:21 -0500 Subject: [PATCH 2/4] fix merge issues Signed-off-by: Kush Gupta --- ramalama/ollama.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/ramalama/ollama.py b/ramalama/ollama.py index 6e35a3b0..5d2d39a5 100644 --- a/ramalama/ollama.py +++ b/ramalama/ollama.py @@ -84,7 +84,9 @@ def in_existing_cache(model_name, model_tag): class Ollama(Model): def __init__(self, model): - super().__init__(model.removeprefix("ollama://")) + model = rm_until_substring(model, "ollama.com/library/") + model = rm_until_substring(model, "://") + super().__init__(model) self.type = "Ollama" def _local(self, args): From 0971ea8c860e8da6fea18100286aba03f5bd3726 Mon Sep 17 00:00:00 2001 From: Kush Gupta Date: Mon, 10 Feb 2025 17:27:38 -0500 Subject: [PATCH 3/4] Fix missing comma on default ollama cache list Signed-off-by: Kush Gupta --- ramalama/ollama.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ramalama/ollama.py b/ramalama/ollama.py index 5d2d39a5..8c577169 100644 --- a/ramalama/ollama.py +++ b/ramalama/ollama.py @@ -66,7 +66,7 @@ def in_existing_cache(model_name, model_tag): return None default_ollama_caches=[ os.path.join(os.environ['HOME'], '.ollama/models'), - '/usr/share/ollama/.ollama/models' + '/usr/share/ollama/.ollama/models', f'C:\\Users\\{os.getlogin()}\\.ollama\\models' ] From 1428cd809345b5a31e226e1ce11c868b018f60e3 Mon Sep 17 00:00:00 2001 From: Kush Gupta Date: Mon, 10 Feb 2025 21:06:41 -0500 Subject: [PATCH 4/4] add conditional to check if the local cache layer exists Signed-off-by: Kush Gupta --- ramalama/ollama.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/ramalama/ollama.py b/ramalama/ollama.py index 8c577169..b80a8922 100644 --- a/ramalama/ollama.py +++ b/ramalama/ollama.py @@ -79,7 +79,8 @@ def in_existing_cache(model_name, model_tag): if layer["mediaType"] == "application/vnd.ollama.image.model": layer_digest = layer["digest"] ollama_digest_path = os.path.join(cache_dir, 'blobs', layer_digest) - return str(ollama_digest_path).replace(':','-') + if os.path.exists(str(ollama_digest_path).replace(':','-')): + return str(ollama_digest_path).replace(':','-') return None class Ollama(Model):