Merge pull request #1189 from engelmi/set-model-store-as-default

Set model store as default

Author: rhatdan

bin/ramalama

@@ -69,6 +69,13 @@ def main(args):
         ramalama.perror("Error: " + str(e).strip("'\""))
         sys.exit(exit_code)
 
+    try:
+        from ramalama.migrate import ModelStoreImport
+
+        ModelStoreImport(args.store).import_all()
+    except Exception as ex:
+        print(f"Failed to import models to new store: {ex}")
+
     # Process CLI
     try:
         args.func(args)

ramalama/cli.py

@@ -219,7 +219,6 @@ def parse_arguments(parser):
 
 def post_parse_setup(args):
     """Perform additional setup after parsing arguments."""
-    mkdirs(args.store)
     if hasattr(args, "MODEL") and args.subcommand != "rm":
         resolved_model = shortnames.resolve(args.MODEL)
         if resolved_model:
@@ -283,23 +282,6 @@ def logout_cli(args):
     return model.logout(args)
 
 
-def mkdirs(store):
-    # List of directories to create
-    directories = [
-        "models/huggingface",
-        "repos/huggingface",
-        "models/oci",
-        "repos/oci",
-        "models/ollama",
-        "repos/ollama",
-    ]
-
-    # Create each directory
-    for directory in directories:
-        full_path = os.path.join(store, directory)
-        os.makedirs(full_path, exist_ok=True)
-
-
 def human_duration(d):
     if d < 1:
         return "Less than a second"
@@ -624,14 +606,14 @@ def convert_cli(args):
         raise ValueError("convert command cannot be run with the --nocontainer option.")
 
     target = args.TARGET
-    source = _get_source(args)
+    source_model = _get_source_model(args)
 
     tgt = shortnames.resolve(target)
     if not tgt:
         tgt = target
 
     model = ModelFactory(tgt, args).create_oci()
-    model.convert(source, args)
+    model.convert(source_model, args)
 
 
 def push_parser(subparsers):
@@ -668,44 +650,38 @@ def push_parser(subparsers):
     parser.set_defaults(func=push_cli)
 
 
-def _get_source(args):
-    if os.path.exists(args.SOURCE):
-        return args.SOURCE
-
+def _get_source_model(args):
     src = shortnames.resolve(args.SOURCE)
     if not src:
         src = args.SOURCE
     smodel = New(src, args)
     if smodel.type == "OCI":
         raise ValueError("converting from an OCI based image %s is not supported" % src)
     if not smodel.exists(args):
-        return smodel.pull(args)
-    return smodel.model_path(args)
+        smodel.pull(args)
+    return smodel
 
 
 def push_cli(args):
-    if args.TARGET:
-        target = args.TARGET
-        source = _get_source(args)
-    else:
-        target = args.SOURCE
-        source = args.SOURCE
 
-    tgt = shortnames.resolve(target)
-    if not tgt:
-        tgt = target
+    source_model = _get_source_model(args)
+    target = args.SOURCE
+    if args.TARGET:
+        target = shortnames.resolve(args.TARGET)
+        if not target:
+            target = args.TARGET
+    target_model = New(target, args)
 
     try:
-        model = New(tgt, args)
-        model.push(source, args)
+        target_model.push(source_model, args)
     except NotImplementedError as e:
         for mtype in MODEL_TYPES:
-            if tgt.startswith(mtype + "://"):
+            if target.startswith(mtype + "://"):
                 raise e
         try:
             # attempt to push as a container image
-            m = ModelFactory(tgt, args).create_oci()
-            m.push(source, args)
+            m = ModelFactory(target, args).create_oci()
+            m.push(source_model, args)
         except Exception as e1:
             if args.debug:
                 print(e1)

ramalama/config.py

@@ -101,7 +101,7 @@ def load_config_defaults(config: Dict[str, Any]):
     config.setdefault('store', get_store())
     config.setdefault('temp', "0.8")
     config.setdefault('transport', "ollama")
-    config.setdefault('use_model_store', False)
+    config.setdefault('use_model_store', True)
 
 
 class Config(ChainMap):

ramalama/huggingface.py

@@ -313,7 +313,7 @@ def url_pull(self, args, model_path, directory_path):
         os.symlink(relative_target_path, model_path)
         return model_path
 
-    def push(self, source, args):
+    def push(self, _, args):
         if not self.hf_cli_available:
             raise NotImplementedError(missing_huggingface)
         proc = run_cmd(
@@ -356,15 +356,18 @@ def _collect_cli_files(self, tempdir: str) -> tuple[str, list[HuggingfaceCLIFile
             if entry.lower() == "readme.md":
                 snapshot_hash = sha256
                 continue
-            files.append(
-                HuggingfaceCLIFile(
-                    url=entry_path,
-                    header={},
-                    hash=sha256,
-                    type=SnapshotFileType.Other,
-                    name=entry,
-                )
+
+            hf_file = HuggingfaceCLIFile(
+                url=entry_path,
+                header={},
+                hash=sha256,
+                type=SnapshotFileType.Other,
+                name=entry,
             )
+            # try to identify the model file in the pulled repo
+            if entry.endswith(".safetensors") or entry.endswith(".gguf"):
+                hf_file.type = SnapshotFileType.Model
+            files.append(hf_file)
 
         return snapshot_hash, files
 
@@ -376,7 +379,7 @@ def _pull_with_model_store(self, debug: bool = False):
 
         try:
             # Fetch the SHA-256 checksum of model from the API and use as snapshot hash
-            snapshot_hash = f"sha256:{fetch_checksum_from_api(self.organization, self.name)}"
+            snapshot_hash = f"sha256:{fetch_checksum_from_api(organization, name)}"
 
             hf_repo = HuggingfaceRepository(name, organization)
             files = hf_repo.get_file_list(cached_files, snapshot_hash)
@@ -387,7 +390,11 @@ def _pull_with_model_store(self, debug: bool = False):
                 raise KeyError(f"Failed to pull model: {str(e)}")
 
             # Cleanup previously created snapshot
-            self.store.remove_snapshot(tag)
+            try:
+                self.store.remove_snapshot(tag)
+            except Exception:
+                # ignore any error when removing snapshot
+                pass
 
             # Create temporary directory for downloading via huggingface-cli
             with tempfile.TemporaryDirectory() as tempdir:

ramalama/migrate.py

@@ -0,0 +1,113 @@
+#
+# Utility for migrating from the old to the new model store
+#
+import os
+import shutil
+
+from ramalama.common import generate_sha256
+from ramalama.model import MODEL_TYPES
+from ramalama.model_factory import ModelFactory
+from ramalama.model_store import GlobalModelStore, SnapshotFile, SnapshotFileType
+
+
+class dotdict(dict):
+    """dot.notation access to dictionary attributes"""
+
+    __getattr__ = dict.get
+    __setattr__ = dict.__setitem__
+    __delattr__ = dict.__delitem__
+
+
+class ModelStoreImport:
+
+    def __init__(self, store_path: str):
+        self.store_path = store_path
+        self._old_model_path = os.path.join(store_path, "models")
+        self._old_repo_path = os.path.join(store_path, "repos")
+        self._global_store = GlobalModelStore(self.store_path)
+
+    class LocalModelFile(SnapshotFile):
+
+        def __init__(
+            self, url, header, hash, name, should_show_progress=False, should_verify_checksum=False, required=True
+        ):
+            super().__init__(
+                url, header, hash, name, SnapshotFileType.Model, should_show_progress, should_verify_checksum, required
+            )
+
+        def download(self, blob_file_path, snapshot_dir):
+            if not os.path.exists(self.url):
+                raise FileNotFoundError(f"No such file: '{self.url}'")
+            # moving from the local location to blob directory so the model store "owns" the data
+            shutil.copy(self.url, blob_file_path)
+            return os.path.relpath(blob_file_path, start=snapshot_dir)
+
+    def import_all(self):
+        if not os.path.exists(self._old_model_path):
+            return
+
+        print("Starting importing AI models to new store...")
+        for root, _, files in os.walk(self._old_model_path):
+            if not files:
+                continue
+
+            try:
+                # reconstruct the cli model input
+                model = root.replace(self._old_model_path, "")
+                model = model.replace(os.sep, "", 1)
+                if model.startswith("file/"):
+                    model = model.replace("/", ":///", 1)
+                else:
+                    model = model.replace("/", "://", 1)
+
+                if model in MODEL_TYPES:
+                    model = ""
+
+                for file in files:
+                    m = ModelFactory(
+                        os.path.join(model, file),
+                        args=dotdict(
+                            {
+                                "store": self.store_path,
+                                "use_model_store": True,
+                                "engine": "podman",
+                                "container": True,
+                            }
+                        ),
+                    ).create()
+                    _, model_tag, _ = m.extract_model_identifiers()
+                    _, _, all = m.store.get_cached_files(model_tag)
+                    if all:
+                        print(f"Already imported: {root}/{file}")
+                        continue
+
+                    snapshot_hash = generate_sha256(file)
+                    old_model_path = os.path.join(root, file)
+
+                    files: list[SnapshotFile] = []
+                    files.append(
+                        ModelStoreImport.LocalModelFile(
+                            url=old_model_path,
+                            header={},
+                            hash=snapshot_hash,
+                            name=file,
+                            required=True,
+                        )
+                    )
+
+                    m.store.new_snapshot(model_tag, snapshot_hash, files)
+                    print(f"Imported {old_model_path} -> {m.store.get_snapshot_file_path(snapshot_hash, file)}")
+
+            except Exception as ex:
+                print(f"Failed to import {root}: {ex}")
+
+        if os.path.exists(self._old_model_path):
+            try:
+                shutil.rmtree(self._old_model_path)
+            except Exception as ex:
+                print(f"Failed to remove old model directory: {ex}")
+        if os.path.exists(self._old_repo_path):
+            try:
+                shutil.rmtree(self._old_repo_path)
+            except Exception as ex:
+                print(f"Failed to remove old blob directory: {ex}")