ob-metaflow-extensions 1.1.131rc0__py2.py3-none-any.whl → 1.1.133__py2.py3-none-any.whl

metaflow_extensions/outerbounds/plugins/__init__.py

@@ -320,6 +320,7 @@ STEP_DECORATORS_DESC = [
     ("tensorboard", ".tensorboard.TensorboardDecorator"),
     ("gpu_profile", ".profilers.gpu_profile_decorator.GPUProfileDecorator"),
     ("nim", ".nim.NimDecorator"),
+    ("ollama", ".ollama.OllamaDecorator"),
 ]
 
 TOGGLE_STEP_DECORATOR = [
@@ -338,4 +339,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"]
+__mf_promote_submodules__ = ["snowflake", "ollama"]

metaflow_extensions/outerbounds/plugins/nim/card.py

@@ -1,10 +1,5 @@
-import os, sqlite3
-from metaflow.cards import (
-    Markdown,
-    Table,
-    ProgressBar,
-)
-from metaflow.decorators import StepDecorator
+import sqlite3
+from metaflow.cards import Markdown, Table
 from metaflow.metaflow_current import current
 
 from .utilities import get_storage_path
@@ -12,10 +7,6 @@ from ..card_utilities.async_cards import CardRefresher
 from ..card_utilities.extra_components import BarPlot, ViolinPlot
 
 
-def json_to_artifact_table(data):
-    return ArtifactTable(data)
-
-
 class NimMetricsRefresher(CardRefresher):
     CARD_ID = "nim_metrics"
 

metaflow_extensions/outerbounds/plugins/ollama/__init__.py

@@ -0,0 +1,70 @@
+from metaflow.decorators import StepDecorator
+from metaflow import current
+import functools
+
+from .ollama import OllamaManager
+from ..card_utilities.injector import CardDecoratorInjector
+
+__mf_promote_submodules__ = ["plugins.ollama"]
+
+
+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'
+    )
+
+    Valid backend options
+    ---------------------
+    - 'local': Run as a separate process on the local task machine.
+    - (TODO) 'managed': Outerbounds hosts and selects compute provider.
+    - (TODO) 'remote': Spin up separate instance to serve Ollama models.
+
+    Valid model options
+    ----------------
+        - 'llama3.2'
+        - 'llama3.3'
+        - any model here https://ollama.com/search
+
+    Parameters
+    ----------
+    models: list[Ollama]
+        List of Ollama containers running models in sidecars.
+    backend: str
+        Determines where and how to run the Ollama process.
+    """
+
+    name = "ollama"
+    defaults = {"models": [], "backend": "local", "debug": False}
+
+    def task_decorate(
+        self, step_func, flow, graph, retry_count, max_user_code_retries, ubf_context
+    ):
+        @functools.wraps(step_func)
+        def ollama_wrapper():
+            try:
+                self.ollama_manager = OllamaManager(
+                    models=self.attributes["models"],
+                    backend=self.attributes["backend"],
+                    debug=self.attributes["debug"],
+                )
+            except Exception as e:
+                print(f"[@ollama] Error initializing OllamaManager: {e}")
+                raise
+            try:
+                step_func()
+            finally:
+                try:
+                    self.ollama_manager.terminate_models()
+                except Exception as term_e:
+                    print(f"[@ollama] Error during sidecar termination: {term_e}")
+            if self.attributes["debug"]:
+                print(f"[@ollama] process statuses: {self.ollama_manager.processes}")
+                print(f"[@ollama] process runtime stats: {self.ollama_manager.stats}")
+
+        return ollama_wrapper

metaflow_extensions/outerbounds/plugins/ollama/ollama.py

@@ -0,0 +1,328 @@
+import subprocess
+from concurrent.futures import ThreadPoolExecutor, as_completed
+import time
+import socket
+import sys
+import os
+import functools
+
+
+class ProcessStatus:
+    RUNNING = "RUNNING"
+    FAILED = "FAILED"
+    SUCCESSFUL = "SUCCESSFUL"
+
+
+class OllamaManager:
+
+    """
+    A process manager for Ollama runtimes.
+    This is run locally, e.g., whether @ollama has a local, remote, or managed backend.
+    """
+
+    def __init__(self, models, backend="local", debug=False):
+        self.models = {}
+        self.processes = {}
+        self.debug = debug
+        self.stats = {}
+
+        if backend != "local":
+            raise ValueError(
+                "OllamaManager only supports the 'local' backend at this time."
+            )
+
+        self._timeit(self._install_ollama, "install_ollama")
+        self._timeit(self._launch_server, "launch_server")
+
+        # Pull models concurrently
+        with ThreadPoolExecutor() as executor:
+            futures = [executor.submit(self._pull_model, m) for m in models]
+            for future in as_completed(futures):
+                try:
+                    future.result()
+                except Exception as e:
+                    raise RuntimeError(f"Error pulling one or more models: {e}") from e
+
+        # Run models as background processes.
+        for m in models:
+            f = functools.partial(self._run_model, m)
+            self._timeit(f, f"model_{m.lower()}")
+
+    def _timeit(self, f, name):
+        t0 = time.time()
+        f()
+        tf = time.time()
+        self.stats[name] = {"process_runtime": tf - t0}
+
+    def _install_ollama(self, max_retries=3):
+
+        try:
+            result = subprocess.run(["which", "ollama"], capture_output=True, text=True)
+            if result.returncode == 0:
+                if self.debug:
+                    print("[@ollama] is already installed.")
+                return
+        except Exception as e:
+            print("[@ollama] Did not find Ollama installation: %s" % e)
+            if sys.platform == "darwin":
+                raise RuntimeError(
+                    "On macOS, please install Ollama manually from https://ollama.com/download"
+                )
+
+        env = os.environ.copy()
+        env["CURL_IPRESOLVE"] = "4"
+        for attempt in range(max_retries):
+            try:
+                install_cmd = ["curl", "-fsSL", "https://ollama.com/install.sh"]
+                curl_proc = subprocess.run(
+                    install_cmd, capture_output=True, text=True, env=env
+                )
+                if curl_proc.returncode != 0:
+                    raise RuntimeError(
+                        f"Failed to download Ollama install script: stdout: {curl_proc.stdout}, stderr: {curl_proc.stderr}"
+                    )
+                sh_proc = subprocess.run(
+                    ["sh"],
+                    input=curl_proc.stdout,
+                    capture_output=True,
+                    text=True,
+                    env=env,
+                )
+                if sh_proc.returncode != 0:
+                    raise RuntimeError(
+                        f"Ollama installation script failed: stdout: {sh_proc.stdout}, stderr: {sh_proc.stderr}"
+                    )
+                if self.debug:
+                    print("[@ollama] Installed successfully.")
+                    break
+            except Exception as e:
+                print(f"Installation attempt {attempt+1} failed: {e}")
+                if attempt < max_retries - 1:
+                    time.sleep(5)
+                else:
+                    raise RuntimeError(
+                        f"Error installing Ollama after {max_retries} attempts: {e}"
+                    ) from e
+
+    def _is_port_open(self, host, port, timeout=1):
+        """Check if a TCP port is open on a given host."""
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+            sock.settimeout(timeout)
+            try:
+                sock.connect((host, port))
+                return True
+            except socket.error:
+                return False
+
+    def _launch_server(self):
+        """
+        Start the Ollama server process and ensure it's running.
+        This version waits until the server is listening on port 11434.
+        """
+        try:
+            if self.debug:
+                print("[@ollama] Starting Ollama server...")
+            process = subprocess.Popen(
+                ["ollama", "serve"],
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                text=True,
+            )
+            self.processes[process.pid] = {
+                "p": process,
+                "properties": {"type": "api-server", "error_details": None},
+                "status": ProcessStatus.RUNNING,
+            }
+            if self.debug:
+                print(
+                    "[@ollama] Started Ollama server process with PID %s" % process.pid
+                )
+
+            # Wait until the server is ready (listening on 127.0.0.1:11434)
+            host, port = "127.0.0.1", 11434
+            retries = 0
+            max_retries = 10
+            while (
+                not self._is_port_open(host, port, timeout=1) and retries < max_retries
+            ):
+                print(
+                    "[@ollama] Waiting for server to be ready... (%d/%d)"
+                    % (retries + 1, max_retries)
+                )
+                time.sleep(1)
+                retries += 1
+
+            if not self._is_port_open(host, port, timeout=1):
+                error_details = (
+                    f"Ollama server did not start listening on {host}:{port}"
+                )
+                self.processes[process.pid]["properties"][
+                    "error_details"
+                ] = error_details
+                self.processes[process.pid]["status"] = ProcessStatus.FAILED
+                raise RuntimeError(f"Ollama server failed to start. {error_details}")
+
+            # Check if the process has unexpectedly terminated
+            returncode = process.poll()
+            if returncode is not None:
+                stdout, stderr = process.communicate()
+                error_details = f"Return code: {returncode}, Error: {stderr}"
+                self.processes[process.pid]["properties"][
+                    "error_details"
+                ] = error_details
+                self.processes[process.pid]["status"] = ProcessStatus.FAILED
+                raise RuntimeError(f"Ollama server failed to start. {error_details}")
+
+            print("[@ollama] Server is ready.")
+
+        except Exception as e:
+            if "process" in locals() and process.pid in self.processes:
+                self.processes[process.pid]["status"] = ProcessStatus.FAILED
+                self.processes[process.pid]["properties"]["error_details"] = str(e)
+            raise RuntimeError(f"Error starting Ollama server: {e}") from e
+
+    def _pull_model(self, m):
+        try:
+            if self.debug:
+                print("[@ollama] Pulling model: %s" % m)
+            result = subprocess.run(
+                ["ollama", "pull", m], capture_output=True, text=True
+            )
+            if result.returncode != 0:
+                raise RuntimeError(
+                    f"Failed to pull model {m}: stdout: {result.stdout}, stderr: {result.stderr}"
+                )
+            if self.debug:
+                print("[@ollama] Model %s pulled successfully." % m)
+        except Exception as e:
+            raise RuntimeError(f"Error pulling Ollama model {m}: {e}") from e
+
+    def _run_model(self, m):
+        """
+        Start the Ollama model as a subprocess and record its status.
+        """
+        process = None
+        try:
+            if self.debug:
+                print("[@ollama] Running model: %s" % m)
+            process = subprocess.Popen(
+                ["ollama", "run", m],
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                text=True,
+            )
+            self.processes[process.pid] = {
+                "p": process,
+                "properties": {"type": "model", "model": m, "error_details": None},
+                "status": ProcessStatus.RUNNING,
+            }
+            if self.debug:
+                print("[@ollama] Stored process %s for model %s." % (process.pid, m))
+
+            try:
+                process.wait(timeout=1)
+            except subprocess.TimeoutExpired:
+                pass
+
+            returncode = process.poll()
+            if returncode is not None:
+                stdout, stderr = process.communicate()
+                if returncode == 0:
+                    self.processes[process.pid]["status"] = ProcessStatus.SUCCESSFUL
+                    if self.debug:
+                        print(
+                            "[@ollama] Process %s for model %s exited successfully."
+                            % (process.pid, m)
+                        )
+                else:
+                    error_details = f"Return code: {returncode}, Error: {stderr}"
+                    self.processes[process.pid]["properties"][
+                        "error_details"
+                    ] = error_details
+                    self.processes[process.pid]["status"] = ProcessStatus.FAILED
+                    if self.debug:
+                        print(
+                            "[@ollama] Process %s for model %s failed: %s"
+                            % (process.pid, m, error_details)
+                        )
+        except Exception as e:
+            if process and process.pid in self.processes:
+                self.processes[process.pid]["status"] = ProcessStatus.FAILED
+                self.processes[process.pid]["properties"]["error_details"] = str(e)
+            raise RuntimeError(f"Error running Ollama model {m}: {e}") from e
+
+    def terminate_models(self):
+        """
+        Terminate all processes gracefully.
+        First, stop model processes using 'ollama stop <model>'.
+        Then, shut down the API server process.
+        """
+
+        for pid, process_info in list(self.processes.items()):
+            if process_info["properties"].get("type") == "model":
+                model_name = process_info["properties"].get("model")
+                if self.debug:
+                    print(
+                        "[@ollama] Stopping model %s using 'ollama stop'" % model_name
+                    )
+                try:
+                    result = subprocess.run(
+                        ["ollama", "stop", model_name], capture_output=True, text=True
+                    )
+                    if result.returncode == 0:
+                        process_info["status"] = ProcessStatus.SUCCESSFUL
+                        if self.debug:
+                            print(
+                                "[@ollama] Model %s stopped successfully." % model_name
+                            )
+                    else:
+                        process_info["status"] = ProcessStatus.FAILED
+                        if self.debug:
+                            print(
+                                "[@ollama] Model %s failed to stop gracefully. Return code: %s, Error: %s"
+                                % (model_name, result.returncode, result.stderr)
+                            )
+                except Exception as e:
+                    process_info["status"] = ProcessStatus.FAILED
+                    print("[@ollama] Error stopping model %s: %s" % (model_name, e))
+
+        # Then, stop the API server
+        for pid, process_info in list(self.processes.items()):
+            if process_info["properties"].get("type") == "api-server":
+                if self.debug:
+                    print(
+                        "[@ollama] Stopping API server process with PID %s using process.terminate()"
+                        % pid
+                    )
+                process = process_info["p"]
+                try:
+                    process.terminate()
+                    try:
+                        process.wait(timeout=5)
+                    except subprocess.TimeoutExpired:
+                        print(
+                            "[@ollama] API server process %s did not terminate in time; killing it."
+                            % pid
+                        )
+                        process.kill()
+                        process.wait()
+                    returncode = process.poll()
+                    if returncode is None or returncode != 0:
+                        process_info["status"] = ProcessStatus.FAILED
+                        print(
+                            "[@ollama] API server process %s terminated with error code %s."
+                            % (pid, returncode)
+                        )
+                    else:
+                        process_info["status"] = ProcessStatus.SUCCESSFUL
+                        if self.debug:
+                            print(
+                                "[@ollama] API server process %s terminated successfully."
+                                % pid
+                            )
+                except Exception as e:
+                    process_info["status"] = ProcessStatus.FAILED
+                    print(
+                        "[@ollama] Warning: Error while terminating API server process %s: %s"
+                        % (pid, e)
+                    )

metaflow_extensions/outerbounds/toplevel/plugins/ollama/__init__.py

@@ -0,0 +1 @@
+__mf_promote_submodules__ = ["plugins.ollama"]

{ob_metaflow_extensions-1.1.131rc0.dist-info → ob_metaflow_extensions-1.1.133.dist-info}/METADATA

@@ -1,13 +1,13 @@
 Metadata-Version: 2.1
 Name: ob-metaflow-extensions
-Version: 1.1.131rc0
+Version: 1.1.133
 Summary: Outerbounds Platform Extensions for Metaflow
 Author: Outerbounds, Inc.
 License: Commercial
 Description-Content-Type: text/markdown
 Requires-Dist: boto3
 Requires-Dist: kubernetes
-Requires-Dist: ob-metaflow (==2.14.0.2rc0)
+Requires-Dist: ob-metaflow (==2.14.3.1)
 
 # Outerbounds platform package
 

{ob_metaflow_extensions-1.1.131rc0.dist-info → ob_metaflow_extensions-1.1.133.dist-info}/RECORD

@@ -1,7 +1,7 @@
 metaflow_extensions/outerbounds/__init__.py,sha256=TRGvIUMjkfneWtYUFSWoubu_Kf2ekAL4WLbV3IxOj9k,499
 metaflow_extensions/outerbounds/remote_config.py,sha256=Zpfpjgz68_ZgxlXezjzlsDLo4840rkWuZgwDB_5H57U,4059
 metaflow_extensions/outerbounds/config/__init__.py,sha256=JsQGRuGFz28fQWjUvxUgR8EKBLGRdLUIk_buPLJplJY,1225
-metaflow_extensions/outerbounds/plugins/__init__.py,sha256=5BFqGhlS9M23WUCoTNQfnOu-mB4iO9z7yK0586IRHMI,12749
+metaflow_extensions/outerbounds/plugins/__init__.py,sha256=nYs0AgKs58TYtRhm6luTp72UFj8-skdO5QPi-sIC3AQ,12802
 metaflow_extensions/outerbounds/plugins/auth_server.py,sha256=_Q9_2EL0Xy77bCRphkwT1aSu8gQXRDOH-Z-RxTUO8N4,2202
 metaflow_extensions/outerbounds/plugins/perimeters.py,sha256=QXh3SFP7GQbS-RAIxUOPbhPzQ7KDFVxZkTdKqFKgXjI,2697
 metaflow_extensions/outerbounds/plugins/card_utilities/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -16,7 +16,7 @@ metaflow_extensions/outerbounds/plugins/fast_bakery/fast_bakery_decorator.py,sha
 metaflow_extensions/outerbounds/plugins/kubernetes/__init__.py,sha256=5zG8gShSj8m7rgF4xgWBZFuY3GDP5n1T0ktjRpGJLHA,69
 metaflow_extensions/outerbounds/plugins/kubernetes/kubernetes_client.py,sha256=fx_XUkgR4r6hF2ilDfT5LubRyVrYMVIv5f6clHkCaEk,5988
 metaflow_extensions/outerbounds/plugins/nim/__init__.py,sha256=MEdX6TPdY9POflCiaYbVmwT-nUNeqwregZBzBZ5CNz0,4487
-metaflow_extensions/outerbounds/plugins/nim/card.py,sha256=EfV8x2XNZkhdOqlNPAGRY_YhahPYxcbFKQAjrXwZTvI,4941
+metaflow_extensions/outerbounds/plugins/nim/card.py,sha256=EM6QtevpJmXpeCesKDk2L6ts6M2qLSYUbajaEEU_yys,4794
 metaflow_extensions/outerbounds/plugins/nim/nim_manager.py,sha256=fCFdRuvagzgPSVQfhy5bzbxfVBeO562h-cDz6CCLQLw,12118
 metaflow_extensions/outerbounds/plugins/nim/utilities.py,sha256=jSdNP3tSCrDjxD2E9bIzxVqDDu6S14femlxSjsMv57o,151
 metaflow_extensions/outerbounds/plugins/nvcf/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -27,6 +27,8 @@ metaflow_extensions/outerbounds/plugins/nvcf/nvcf.py,sha256=zhHgoE76WOpCJnoMvSNe
 metaflow_extensions/outerbounds/plugins/nvcf/nvcf_cli.py,sha256=3D-r5XO88Yh2k1EAZFJTe_PwdbhWp5qXflG8AgE4ZIU,9500
 metaflow_extensions/outerbounds/plugins/nvcf/nvcf_decorator.py,sha256=pxxNxW4bW3kbB6ybRam657GyKIhvIkMuidB94iFfCD8,9116
 metaflow_extensions/outerbounds/plugins/nvcf/utils.py,sha256=DxWSCayfa95e0HJkWacey1s1nxoTpaunGhrb_0Ayv28,133
+metaflow_extensions/outerbounds/plugins/ollama/__init__.py,sha256=HEsI5U4ckQby7K2NsGBOdizhPY3WWqXSnXx_IHL7_No,2307
+metaflow_extensions/outerbounds/plugins/ollama/ollama.py,sha256=KlP8_EmnUoi8-PidyU0IDuENYxKjQaHFC33yGsvaeic,13320
 metaflow_extensions/outerbounds/plugins/profilers/deco_injector.py,sha256=oI_C3c64XBm7n88FILqHwn-Nnc5DeT_68I67lM9rXaI,2434
 metaflow_extensions/outerbounds/plugins/profilers/gpu_profile_decorator.py,sha256=gDHQ2sMIp4NuZSzUspbSd8RGdFAoO5mgZAyFcZ2a51Y,2619
 metaflow_extensions/outerbounds/plugins/secrets/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -49,8 +51,9 @@ metaflow_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.py,
 metaflow_extensions/outerbounds/toplevel/plugins/azure/__init__.py,sha256=WUuhz2YQfI4fz7nIcipwwWq781eaoHEk7n4GAn1npDg,63
 metaflow_extensions/outerbounds/toplevel/plugins/gcp/__init__.py,sha256=BbZiaH3uILlEZ6ntBLKeNyqn3If8nIXZFq_Apd7Dhco,70
 metaflow_extensions/outerbounds/toplevel/plugins/kubernetes/__init__.py,sha256=5zG8gShSj8m7rgF4xgWBZFuY3GDP5n1T0ktjRpGJLHA,69
+metaflow_extensions/outerbounds/toplevel/plugins/ollama/__init__.py,sha256=GRSz2zwqkvlmFS6bcfYD_CX6CMko9DHQokMaH1iBshA,47
 metaflow_extensions/outerbounds/toplevel/plugins/snowflake/__init__.py,sha256=LptpH-ziXHrednMYUjIaosS1SXD3sOtF_9_eRqd8SJw,50
-ob_metaflow_extensions-1.1.131rc0.dist-info/METADATA,sha256=X_n4lJEbpxO3uOnW2OdurRIzuTfHcfHmQecNspjx_uU,526
-ob_metaflow_extensions-1.1.131rc0.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-ob_metaflow_extensions-1.1.131rc0.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
-ob_metaflow_extensions-1.1.131rc0.dist-info/RECORD,,
+ob_metaflow_extensions-1.1.133.dist-info/METADATA,sha256=OA0TksO1c611FUdupK69MqlXGzthWPwYthaHjPqC4dA,520
+ob_metaflow_extensions-1.1.133.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+ob_metaflow_extensions-1.1.133.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
+ob_metaflow_extensions-1.1.133.dist-info/RECORD,,