ob-metaflow-extensions 1.1.156__tar.gz → 1.1.158__tar.gz

{ob-metaflow-extensions-1.1.156 → ob-metaflow-extensions-1.1.158}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ob-metaflow-extensions
-Version: 1.1.156
+Version: 1.1.158
 Summary: Outerbounds Platform Extensions for Metaflow
 Author: Outerbounds, Inc.
 License: Commercial

{ob-metaflow-extensions-1.1.156 → ob-metaflow-extensions-1.1.158}/metaflow_extensions/outerbounds/plugins/nim/nim_manager.py

@@ -186,18 +186,61 @@ class NimManager(object):
         ]
         self.models = {}
 
-        for each_model in models:
-            if each_model in nvcf_models:
-                self.models[each_model] = NimChatCompletion(
-                    model=each_model,
+        # Convert models to a standard format
+        standardized_models = []
+        # If models is a single string, convert it to a list with a dict
+        if isinstance(models, str):
+            standardized_models = [{"name": models}]
+        # If models is a list, process each item
+        elif isinstance(models, list):
+            for model_item in models:
+                # If the item is a string, convert it to a dict
+                if isinstance(model_item, str):
+                    standardized_models.append({"name": model_item})
+                # If it's already a dict, use it as is
+                elif isinstance(model_item, dict):
+                    standardized_models.append(model_item)
+                else:
+                    raise ValueError(
+                        f"Model specification must be a string or dictionary, got {type(model_item)}"
+                    )
+        else:
+            raise ValueError(
+                f"Models must be a string or a list of strings/dictionaries, got {type(models)}"
+            )
+
+        # Process each standardized model
+        for each_model_dict in standardized_models:
+            model_name = each_model_dict.get("name", "")
+            nvcf_id = each_model_dict.get("nvcf_id", "")
+            nvcf_version = each_model_dict.get("nvcf_version", "")
+
+            if model_name and not (nvcf_id and nvcf_version):
+                if model_name in nvcf_models:
+                    self.models[model_name] = NimChatCompletion(
+                        model=model_name,
+                        nvcf_id=nvcf_id,
+                        nvcf_version=nvcf_version,
+                        nim_metadata=nim_metadata,
+                        monitor=monitor,
+                    )
+                else:
+                    raise ValueError(
+                        f"Model {model_name} not supported by the Outerbounds @nim offering."
+                        f"\nYou can choose from these options: {nvcf_models}\n\n"
+                        "Reach out to Outerbounds if there are other models you'd like supported."
+                    )
+            elif nvcf_id and nvcf_version:
+                self.models[model_name] = NimChatCompletion(
+                    model=model_name,
+                    nvcf_id=nvcf_id,
+                    nvcf_version=nvcf_version,
                     nim_metadata=nim_metadata,
                     monitor=monitor,
                 )
             else:
                 raise ValueError(
-                    f"Model {each_model} not supported by the Outerbounds @nim offering."
-                    f"\nYou can choose from these options: {nvcf_models}\n\n"
-                    "Reach out to Outerbounds if there are other models you'd like supported."
+                    "You must provide either a valid 'name' or a custom 'name' along with both 'nvcf_id' and 'nvcf_version'."
                 )
 
 
@@ -205,6 +248,8 @@ class NimChatCompletion(object):
     def __init__(
         self,
         model: str = "meta/llama3-8b-instruct",
+        nvcf_id: str = "",
+        nvcf_version: str = "",
         nim_metadata: NimMetadata = None,
         monitor: bool = False,
         **kwargs,
@@ -217,18 +262,34 @@ class NimChatCompletion(object):
         self.model_name = model
         self.nim_metadata = nim_metadata
         self.monitor = monitor
-
         all_nvcf_models = self.nim_metadata.get_nvcf_chat_completion_models()
-        all_nvcf_model_names = [m["name"] for m in all_nvcf_models]
 
-        if self.model_name not in all_nvcf_model_names:
-            raise ValueError(
-                f"Model {self.model_name} not found in available NVCF models"
-            )
+        if nvcf_id and nvcf_version:
+            matching_models = [
+                m
+                for m in all_nvcf_models
+                if m["function-id"] == nvcf_id and m["version-id"] == nvcf_version
+            ]
+            if matching_models:
+                self.model = matching_models[0]
+                self.function_id = self.model["function-id"]
+                self.version_id = self.model["version-id"]
+                self.model_name = self.model["name"]
+            else:
+                raise ValueError(
+                    f"Function {self.function_id} with version {self.version_id} not found on NVCF"
+                )
+        else:
+            all_nvcf_model_names = [m["name"] for m in all_nvcf_models]
+
+            if self.model_name not in all_nvcf_model_names:
+                raise ValueError(
+                    f"Model {self.model_name} not found in available NVCF models"
+                )
 
-        self.model = all_nvcf_models[all_nvcf_model_names.index(self.model_name)]
-        self.function_id = self.model["function-id"]
-        self.version_id = self.model["version-id"]
+            self.model = all_nvcf_models[all_nvcf_model_names.index(self.model_name)]
+            self.function_id = self.model["function-id"]
+            self.version_id = self.model["version-id"]
 
         self.first_request = True
 

{ob-metaflow-extensions-1.1.156 → ob-metaflow-extensions-1.1.158}/metaflow_extensions/outerbounds/plugins/ollama/__init__.py

@@ -1,6 +1,7 @@
 from metaflow.decorators import StepDecorator
 from metaflow import current
 import functools
+import os
 
 from .ollama import OllamaManager
 from ..card_utilities.injector import CardDecoratorInjector
@@ -13,10 +14,10 @@ class OllamaDecorator(StepDecorator, CardDecoratorInjector):
     This decorator is used to run Ollama APIs as Metaflow task sidecars.
 
     User code call
-    -----------
+    --------------
     @ollama(
-        models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
-        backend='local'
+        models=[...],
+        ...
     )
 
     Valid backend options
@@ -26,21 +27,39 @@ class OllamaDecorator(StepDecorator, CardDecoratorInjector):
     - (TODO) 'remote': Spin up separate instance to serve Ollama models.
 
     Valid model options
-    ----------------
-        - 'llama3.2'
-        - 'llama3.3'
-        - any model here https://ollama.com/search
+    -------------------
+    Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
 
     Parameters
     ----------
-    models: list[Ollama]
+    models: list[str]
         List of Ollama containers running models in sidecars.
     backend: str
         Determines where and how to run the Ollama process.
+    force_pull: bool
+        Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
+    skip_push_check: bool
+        Whether to skip the check that populates/overwrites remote cache on terminating an ollama model.
+    debug: bool
+        Whether to turn on verbose debugging logs.
     """
 
     name = "ollama"
-    defaults = {"models": [], "backend": "local", "debug": False}
+    defaults = {
+        "models": [],
+        "backend": "local",
+        "force_pull": False,
+        "skip_push_check": False,
+        "debug": False,
+    }
+
+    def step_init(
+        self, flow, graph, step_name, decorators, environment, flow_datastore, logger
+    ):
+        super().step_init(
+            flow, graph, step_name, decorators, environment, flow_datastore, logger
+        )
+        self.flow_datastore_backend = flow_datastore._storage_impl
 
     def task_decorate(
         self, step_func, flow, graph, retry_count, max_user_code_retries, ubf_context
@@ -51,6 +70,9 @@ class OllamaDecorator(StepDecorator, CardDecoratorInjector):
                 self.ollama_manager = OllamaManager(
                     models=self.attributes["models"],
                     backend=self.attributes["backend"],
+                    flow_datastore_backend=self.flow_datastore_backend,
+                    force_pull=self.attributes["force_pull"],
+                    skip_push_check=self.attributes["skip_push_check"],
                     debug=self.attributes["debug"],
                 )
             except Exception as e:
@@ -59,10 +81,7 @@ class OllamaDecorator(StepDecorator, CardDecoratorInjector):
             try:
                 step_func()
             finally:
-                try:
-                    self.ollama_manager.terminate_models()
-                except Exception as term_e:
-                    print(f"[@ollama] Error during sidecar termination: {term_e}")
+                self.ollama_manager.terminate_models()
             if self.attributes["debug"]:
                 print(f"[@ollama] process statuses: {self.ollama_manager.processes}")
                 print(f"[@ollama] process runtime stats: {self.ollama_manager.stats}")

ob-metaflow-extensions-1.1.158/metaflow_extensions/outerbounds/plugins/ollama/constants.py

@@ -0,0 +1 @@
+OLLAMA_SUFFIX = "mf.ollama"

ob-metaflow-extensions-1.1.158/metaflow_extensions/outerbounds/plugins/ollama/exceptions.py

@@ -0,0 +1,22 @@
+from metaflow.exception import MetaflowException
+
+
+class UnspecifiedRemoteStorageRootException(MetaflowException):
+    headline = "Storage root not specified."
+
+    def __init__(self, message):
+        super(UnspecifiedRemoteStorageRootException, self).__init__(message)
+
+
+class EmptyOllamaManifestCacheException(MetaflowException):
+    headline = "Model not found."
+
+    def __init__(self, message):
+        super(EmptyOllamaManifestCacheException, self).__init__(message)
+
+
+class EmptyOllamaBlobCacheException(MetaflowException):
+    headline = "Blob not found."
+
+    def __init__(self, message):
+        super(EmptyOllamaBlobCacheException, self).__init__(message)