[draft_model] Fix OCI mount options for draft_model

When supplying draft model the old code pretty much tried to bindmount
model with both src and destination in the end translating to same
location. Also there was no support for mulitple file models so when one
tried to use 440B as main model and 30B as draft model it would not be
possible.

To prevent file collisions for example for mmproj and chat_template we
use namespaced directory. As of now draft_models won't work with
multimodal models but we can still do speculative decoding when
disabling vision by passing --no-mmproj.

Signed-off-by: Lukas Bezdicka <[email protected]>

Author: xbezdick

ramalama/command/context.py

@@ -2,7 +2,7 @@
 import os
 from typing import Optional
 
-from ramalama.common import check_metal, check_nvidia
+from ramalama.common import MNT_DIR, check_metal, check_nvidia
 from ramalama.console import should_colorize
 from ramalama.transports.transport_factory import CLASS_MODEL_TYPES, New
 
@@ -123,7 +123,15 @@ def chat_template_path(self) -> Optional[str]:
     def draft_model_path(self) -> str:
         if getattr(self.model, "draft_model", None):
             assert self.model.draft_model
-            return self.model.draft_model._get_entry_model_path(self.is_container, self.should_generate, self.dry_run)
+            path = self.model.draft_model._get_entry_model_path(self.is_container, self.should_generate, self.dry_run)
+            if self.is_container and not path.startswith("oci://"):
+                # Handle container paths by inserting 'drafts' into the MNT_DIR path
+                if path.startswith(MNT_DIR):
+                    rest = path[len(MNT_DIR) :].lstrip('/')
+                    return f"{MNT_DIR}/drafts/{rest}"
+                else:
+                    return path
+            return path
         return ""
 
 

ramalama/transports/base.py

@@ -12,7 +12,6 @@
 import ramalama.chat as chat
 from ramalama.common import (
     MNT_DIR,
-    MNT_FILE_DRAFT,
     accel_image,
     exec_cmd,
     genname,
@@ -158,8 +157,10 @@ def __init__(self, model: str, model_store_path: str):
     def artifact(self) -> bool:
         return self.is_artifact()
 
-    def extract_model_identifiers(self):
-        model_name = self.model
+    def extract_model_identifiers(self, model=None):
+        if model is None:
+            model = self.model
+        model_name = model
         model_tag = "latest"
         model_organization = ""
 
@@ -391,12 +392,20 @@ def setup_mounts(self, args):
             )
 
         if self.draft_model:
-            draft_model = self.draft_model._get_entry_model_path(args.container, args.generate, args.dryrun)
-            # Convert path to container-friendly format (handles Windows path conversion)
-            container_draft_model = get_container_mount_path(draft_model)
-            mount_opts = f"--mount=type=bind,src={container_draft_model},destination={MNT_FILE_DRAFT}"
-            mount_opts += f",ro{self.engine.relabel()}"
-            self.engine.add([mount_opts])
+            # Get the model tag for draft model
+            _, draft_tag, _ = self.extract_model_identifiers(model=self.draft_model.model)
+            ref_file_draft = self.draft_model.model_store.get_ref_file(draft_tag)
+            if ref_file_draft is None:
+                raise NoRefFileFound(self.draft_model.model)
+
+            # Mount the draft model files
+            for file in ref_file_draft.files:
+                blob_path = self.draft_model.model_store.get_blob_file_path(file.hash)
+                container_blob_path = get_container_mount_path(blob_path)
+                mount_path = f"{MNT_DIR}/drafts/{file.name}"
+                self.engine.add(
+                    [f"--mount=type=bind,src={container_blob_path},destination={mount_path},ro{self.engine.relabel()}"]
+                )
 
     def bench(self, args, cmd: list[str]):
         set_accel_env_vars()

ramalama/transports/huggingface.py

@@ -163,8 +163,8 @@ def create_repository(self, name, organization, tag):
     def get_cli_download_args(self, directory_path, model):
         return ["hf", "download", "--local-dir", directory_path, model]
 
-    def extract_model_identifiers(self):
-        model_name, model_tag, model_organization = super().extract_model_identifiers()
+    def extract_model_identifiers(self, model=None):
+        model_name, model_tag, model_organization = super().extract_model_identifiers(model)
         if '/' not in model_organization:
             # if it is a repo then normalize the case insensitive quantization tag
             if model_tag != "latest":

ramalama/transports/ollama.py

@@ -146,8 +146,8 @@ def __init__(self, model, model_store_path) -> None:
 
         self.type = "Ollama"
 
-    def extract_model_identifiers(self) -> tuple[str, str, str]:
-        model_name, model_tag, model_organization = super().extract_model_identifiers()
+    def extract_model_identifiers(self, model=None) -> tuple[str, str, str]:
+        model_name, model_tag, model_organization = super().extract_model_identifiers(model)
 
         # use the ollama default namespace if no model organization has been identified
         if not model_organization:

ramalama/transports/url.py

@@ -51,8 +51,8 @@ def __init__(self, model, model_store_path, scheme):
         split = self.model.rsplit("/", 1)
         self.directory = split[0].removeprefix("/") if len(split) > 1 else ""
 
-    def extract_model_identifiers(self):
-        model_name, model_tag, model_organization = super().extract_model_identifiers()
+    def extract_model_identifiers(self, model=None):
+        model_name, model_tag, model_organization = super().extract_model_identifiers(model)
 
         parts = model_organization.split("/")
         if len(parts) > 2 and parts[-2] == "blob":