Merge pull request #1756 from engelmi/fix-big-model-fetching

Fix assembling URLs for big models

Author: rhatdan

ramalama/url.py

@@ -75,6 +75,35 @@ def extract_model_identifiers(self):
 
         return model_name, model_tag, model_organization
 
+    def _assemble_split_file_list(self, snapshot_hash: str) -> list[SnapshotFile]:
+        files: list[SnapshotFile] = []
+
+        # model is split, lets fetch all files based on the name pattern
+        match = re.match(SPLIT_MODEL_PATH_RE, self.model)
+        if match is None:
+            return files
+
+        path_part = match[1]
+        filename_base = match[2]
+        total_parts = int(match[3])
+
+        for i in range(1, total_parts + 1):
+            file_name = f"{filename_base}-{i:05d}-of-{total_parts:05d}.gguf"
+            url = f"{self.type}://{path_part}/{file_name}"
+            files.append(
+                SnapshotFile(
+                    url=url,
+                    header={},
+                    hash=snapshot_hash,
+                    type=SnapshotFileType.Model,
+                    name=file_name,
+                    should_show_progress=True,
+                    required=True,
+                )
+            )
+
+        return files
+
     def pull(self, _):
         name, tag, _ = self.extract_model_identifiers()
         _, _, all_files = self.model_store.get_cached_files(tag)
@@ -96,40 +125,21 @@ def pull(self, _):
             self.model_store.new_snapshot(tag, snapshot_hash, files)
             return
 
-        if not is_split_file_model(self.model):
-            files.append(
-                SnapshotFile(
-                    url=f"{self.type}://{self.model}",
-                    header={},
-                    hash=snapshot_hash,
-                    type=SnapshotFileType.Model,
-                    name=name,
-                    should_show_progress=True,
-                    required=True,
-                )
-            )
+        if is_split_file_model(self.model):
+            files = self._assemble_split_file_list(snapshot_hash)
             self.model_store.new_snapshot(tag, snapshot_hash, files)
             return
 
-        # model is split, lets fetch all files based on the name pattern
-        match = re.match(SPLIT_MODEL_PATH_RE, self.model)
-        path_part = match[1]
-        filename_base = match[2]
-        total_parts = int(match[3])
-
-        for i in range(total_parts - 1):
-            i_off = i + 2
-            url = f"{self.type}://{path_part}/{filename_base}-{i_off:05d}-of-{total_parts:05d}.gguf"
-            files.append(
-                SnapshotFile(
-                    url=url,
-                    header={},
-                    hash=snapshot_hash,
-                    type=SnapshotFileType.Model,
-                    name=name,
-                    should_show_progress=True,
-                    required=True,
-                )
+        files.append(
+            SnapshotFile(
+                url=f"{self.type}://{self.model}",
+                header={},
+                hash=snapshot_hash,
+                type=SnapshotFileType.Model,
+                name=name,
+                should_show_progress=True,
+                required=True,
             )
-
+        )
         self.model_store.new_snapshot(tag, snapshot_hash, files)
+        return

test/unit/test_url.py

@@ -0,0 +1,58 @@
+from dataclasses import dataclass, field
+
+import pytest
+
+from ramalama.model_store.snapshot_file import SnapshotFile
+from ramalama.url import URL
+
+
+@dataclass
+class Input:
+    Model: str
+
+
+@dataclass
+class Expected:
+    URLList: list[str] = field(default_factory=lambda: [])
+    Names: list[str] = field(default_factory=lambda: [])
+
+
+@pytest.mark.parametrize(
+    "input,expected",
+    [
+        (Input(""), Expected()),
+        (Input("file:///tmp/models/granite-3b-code-base.Q4_K_M.gguf"), Expected()),
+        (
+            Input(
+                "huggingface.co/unsloth/Qwen3-Coder-480B-A35B-Instruct-GGUF/resolve/main/Q3_K_M/Qwen3-Coder-480B-A35B-Instruct-Q3_K_M-00001-of-00005.gguf"  # noqa: E501
+            ),
+            Expected(
+                URLList=[
+                    "https://huggingface.co/unsloth/Qwen3-Coder-480B-A35B-Instruct-GGUF/resolve/main/Q3_K_M/Qwen3-Coder-480B-A35B-Instruct-Q3_K_M-00001-of-00005.gguf",  # noqa: E501
+                    "https://huggingface.co/unsloth/Qwen3-Coder-480B-A35B-Instruct-GGUF/resolve/main/Q3_K_M/Qwen3-Coder-480B-A35B-Instruct-Q3_K_M-00002-of-00005.gguf",  # noqa: E501
+                    "https://huggingface.co/unsloth/Qwen3-Coder-480B-A35B-Instruct-GGUF/resolve/main/Q3_K_M/Qwen3-Coder-480B-A35B-Instruct-Q3_K_M-00003-of-00005.gguf",  # noqa: E501
+                    "https://huggingface.co/unsloth/Qwen3-Coder-480B-A35B-Instruct-GGUF/resolve/main/Q3_K_M/Qwen3-Coder-480B-A35B-Instruct-Q3_K_M-00004-of-00005.gguf",  # noqa: E501
+                    "https://huggingface.co/unsloth/Qwen3-Coder-480B-A35B-Instruct-GGUF/resolve/main/Q3_K_M/Qwen3-Coder-480B-A35B-Instruct-Q3_K_M-00005-of-00005.gguf",  # noqa: E501
+                ],
+                Names=[
+                    "Qwen3-Coder-480B-A35B-Instruct-Q3_K_M-00001-of-00005.gguf",
+                    "Qwen3-Coder-480B-A35B-Instruct-Q3_K_M-00002-of-00005.gguf",
+                    "Qwen3-Coder-480B-A35B-Instruct-Q3_K_M-00003-of-00005.gguf",
+                    "Qwen3-Coder-480B-A35B-Instruct-Q3_K_M-00004-of-00005.gguf",
+                    "Qwen3-Coder-480B-A35B-Instruct-Q3_K_M-00005-of-00005.gguf",
+                ],
+            ),
+        ),
+    ],
+)
+def test__assemble_split_file_list(input: Input, expected: Expected):
+    # store path and scheme irrelevant here
+    model = URL(input.Model, "/store", "https")
+    files: list[SnapshotFile] = model._assemble_split_file_list("doesnotmatterhere")
+    file_count = len(files)
+    assert file_count == len(expected.Names)
+    assert file_count == len(expected.URLList)
+
+    for i in range(file_count):
+        assert files[i].url == expected.URLList[i]
+        assert files[i].name == expected.Names[i]