ob-metaflow-extensions 1.1.157__tar.gz → 1.1.159__tar.gz

{ob-metaflow-extensions-1.1.157 → ob-metaflow-extensions-1.1.159}/PKG-INFO

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

{ob-metaflow-extensions-1.1.157 → ob-metaflow-extensions-1.1.159}/metaflow_extensions/outerbounds/plugins/__init__.py

@@ -347,4 +347,4 @@ SECRETS_PROVIDERS_DESC = [
     ("outerbounds", ".secrets.secrets.OuterboundsSecretsProvider"),
 ]
 # Adding an override here so the library can be imported at the metaflow.plugins level
-__mf_promote_submodules__ = ["snowflake", "ollama"]
+__mf_promote_submodules__ = ["snowflake", "ollama", "torchtune"]

{ob-metaflow-extensions-1.1.157 → ob-metaflow-extensions-1.1.159}/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.159/metaflow_extensions/outerbounds/plugins/ollama/constants.py

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

ob-metaflow-extensions-1.1.159/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)