ob-metaflow-extensions 1.1.142__py2.py3-none-any.whl → 1.4.33__py2.py3-none-any.whl

metaflow_extensions/outerbounds/plugins/nvct/nvct_runner.py

@@ -0,0 +1,218 @@
+import os
+import re
+import time
+import math
+import shlex
+import atexit
+
+from metaflow import util
+from metaflow.mflog import (
+    BASH_SAVE_LOGS,
+    bash_capture_logs,
+    export_mflog_env_vars,
+    tail_logs,
+    get_log_tailer,
+)
+
+from .nvct import NVCTClient, NVCTTask, NVCTRequest
+from .exceptions import (
+    NvctKilledException,
+    NvctExecutionException,
+    NvctTaskFailedException,
+)
+
+# Constants for Metaflow logs
+LOGS_DIR = "$PWD/.logs"
+STDOUT_FILE = "mflog_stdout"
+STDERR_FILE = "mflog_stderr"
+STDOUT_PATH = os.path.join(LOGS_DIR, STDOUT_FILE)
+STDERR_PATH = os.path.join(LOGS_DIR, STDERR_FILE)
+NVCT_WRAPPER = "/usr/local/bin/nvct-wrapper.sh"
+
+
+class NvctRunner:
+    def __init__(
+        self,
+        metadata,
+        datastore,
+        environment,
+        gpu_type,
+        instance_type,
+        backend,
+        ngc_api_key,
+    ):
+        self.metadata = metadata
+        self.datastore = datastore
+        self.environment = environment
+        self.gpu_type = gpu_type
+        self.instance_type = instance_type
+        self.backend = backend
+        self._ngc_api_key = ngc_api_key
+        self.client = None
+        self.task = None
+        atexit.register(lambda: self.task.cancel() if hasattr(self, "task") else None)
+
+    def launch_task(
+        self,
+        step_name,
+        step_cli,
+        task_spec,
+        code_package_sha,
+        code_package_url,
+        code_package_ds,
+        env={},
+        max_runtime="PT7H",  # <8H allowed for GFN backend
+        max_queued="PT120H",  # 5 days
+    ):
+        mflog_expr = export_mflog_env_vars(
+            datastore_type=code_package_ds,
+            stdout_path=STDOUT_PATH,
+            stderr_path=STDERR_PATH,
+            **task_spec,
+        )
+        init_cmds = self.environment.get_package_commands(
+            code_package_url, code_package_ds
+        )
+        init_expr = " && ".join(init_cmds)
+        step_expr = bash_capture_logs(
+            " && ".join(
+                self.environment.bootstrap_commands(step_name, code_package_ds)
+                + [step_cli]
+            )
+        )
+        cmd_str = "mkdir -p %s && %s && %s && %s; c=$?; %s; exit $c" % (
+            LOGS_DIR,
+            mflog_expr,
+            init_expr,
+            step_expr,
+            BASH_SAVE_LOGS,
+        )
+
+        # Add optional initialization script execution
+        cmd_str = (
+            '${METAFLOW_INIT_SCRIPT:+eval \\"${METAFLOW_INIT_SCRIPT}\\"} && %s'
+            % cmd_str
+        )
+
+        cmd_str = shlex.split('bash -c "%s"' % cmd_str)[-1]
+
+        def modify_python_c(match):
+            content = match.group(1)
+            # Escape double quotes within the python -c command
+            content = content.replace('"', r"\"")
+            # Replace outermost double quotes with single quotes
+            return 'python -c "%s"' % content
+
+        # Convert python -c single quotes to double quotes
+        cmd_str = re.sub(r"python -c '(.*?)'", modify_python_c, cmd_str)
+        cmd_str = cmd_str.replace("'", '"')
+        # Create the final command with outer single quotes to pass to NVCT wrapper
+        nvct_cmd = f"{NVCT_WRAPPER} bash -c '{cmd_str}'"
+
+        flow_name = task_spec.get("flow_name")
+        run_id = task_spec.get("run_id")
+        task_id = task_spec.get("task_id")
+        retry_count = task_spec.get("retry_count")
+        task_name = f"{flow_name}-{run_id}-{step_name}-{task_id}-{retry_count}"
+
+        if self.backend != "GFN":
+            # if maxRuntimeDuration exceeds 8 hours for a Task on the GFN backend,
+            # the request will be rejected.
+            # (https://docs.nvidia.com/cloud-functions/user-guide/latest/cloud-function/tasks.html#create-task)
+            ## thus, if it is non GFN backend, we increase it to 3 days
+            max_runtime = "PT72H"
+
+        request = (
+            NVCTRequest(task_name)
+            .container_image("nvcr.io/zhxkmsaasxhw/nvct-base:2.0-jovyan")
+            .container_args(nvct_cmd)
+            .gpu(
+                gpu=self.gpu_type,
+                instance_type=self.instance_type,
+                backend=self.backend,
+            )
+            .max_runtime(max_runtime)
+            .max_queued(max_queued)
+        )
+
+        for k, v in env.items():
+            if v is not None:
+                request.env(k, str(v))
+
+        self.client = NVCTClient(self._ngc_api_key)
+        self.task = NVCTTask(self.client, request.to_dict())
+
+        self.task.submit()
+        return self.task.id
+
+    def wait_for_completion(self, stdout_location, stderr_location, echo=None):
+        if not self.task:
+            raise NvctExecutionException("No task has been launched")
+
+        def update_delay(secs_since_start):
+            # this sigmoid function reaches
+            # - 0.1 after 11 minutes
+            # - 0.5 after 15 minutes
+            # - 1.0 after 23 minutes
+            # in other words, the user will see very frequent updates
+            # during the first 10 minutes
+            sigmoid = 1.0 / (1.0 + math.exp(-0.01 * secs_since_start + 9.0))
+            return 0.5 + sigmoid * 30.0
+
+        def wait_for_launch(task):
+            status = task.status
+            echo(
+                "Task is starting (%s)..." % status,
+                "stderr",
+                _id=task.id,
+            )
+
+            t = time.time()
+            start_time = time.time()
+            while task.is_waiting:
+                new_status = task.status
+                if status != new_status or (time.time() - t) > 30:
+                    status = new_status
+                    echo(
+                        "Task is starting (%s)..." % status,
+                        "stderr",
+                        _id=task.id,
+                    )
+                    t = time.time()
+                time.sleep(update_delay(time.time() - start_time))
+
+        _make_prefix = lambda: b"[%s] " % util.to_bytes(self.task.id)
+        stdout_tail = get_log_tailer(stdout_location, self.datastore.TYPE)
+        stderr_tail = get_log_tailer(stderr_location, self.datastore.TYPE)
+
+        # 1) Loop until the job has started
+        wait_for_launch(self.task)
+
+        echo(
+            "Task is starting (%s)..." % self.task.status,
+            "stderr",
+            _id=self.task.id,
+        )
+
+        # 2) Tail logs until the job has finished
+        tail_logs(
+            prefix=_make_prefix(),
+            stdout_tail=stdout_tail,
+            stderr_tail=stderr_tail,
+            echo=echo,
+            has_log_updates=lambda: self.task.is_running,
+        )
+
+        if self.task.has_failed:
+            raise NvctTaskFailedException(
+                f"Task failed with status: {self.task.status}. This could be a transient error. Use @retry to retry."
+            )
+        else:
+            if self.task.is_running:
+                # Kill the job if it is still running by throwing an exception.
+                raise NvctKilledException("Task failed!")
+            echo(
+                f"Task finished with status: {self.task.status}",
+                "stderr",
+                _id=self.task.id,
+            )

metaflow_extensions/outerbounds/plugins/nvct/utils.py

@@ -0,0 +1,29 @@
+import os
+import json
+import requests
+from urllib.parse import urlparse
+from metaflow.metaflow_config import SERVICE_URL
+from metaflow.metaflow_config_funcs import init_config
+
+
+def get_ngc_api_key():
+    conf = init_config()
+    if "OBP_AUTH_SERVER" in conf:
+        auth_host = conf["OBP_AUTH_SERVER"]
+    else:
+        auth_host = "auth." + urlparse(SERVICE_URL).hostname.split(".", 1)[1]
+
+    # NOTE: reusing the same auth_host as the one used in NimMetadata,
+    # however, user should not need to use nim container to use @nvct.
+    # May want to refactor this to a common endpoint.
+    nim_info_url = "https://" + auth_host + "/generate/nim"
+
+    if "METAFLOW_SERVICE_AUTH_KEY" in conf:
+        headers = {"x-api-key": conf["METAFLOW_SERVICE_AUTH_KEY"]}
+        res = requests.get(nim_info_url, headers=headers)
+    else:
+        headers = json.loads(os.environ.get("METAFLOW_SERVICE_HEADERS"))
+        res = requests.get(nim_info_url, headers=headers)
+
+    res.raise_for_status()
+    return res.json()["nvcf"]["api_key"]

metaflow_extensions/outerbounds/plugins/ollama/__init__.py

@@ -1,8 +1,11 @@
 from metaflow.decorators import StepDecorator
 from metaflow import current
 import functools
+import os
+import threading
 
-from .ollama import OllamaManager
+from .ollama import OllamaManager, OllamaRequestInterceptor
+from .status_card import OllamaStatusCard
 from ..card_utilities.injector import CardDecoratorInjector
 
 __mf_promote_submodules__ = ["plugins.ollama"]
@@ -13,10 +16,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,45 +29,197 @@ 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.
+    cache_update_policy: str
+        Cache update policy: "auto", "force", or "never".
+    force_cache_update: bool
+        Simple override for "force" cache update policy.
+    debug: bool
+        Whether to turn on verbose debugging logs.
+    circuit_breaker_config: dict
+        Configuration for circuit breaker protection. Keys: failure_threshold, recovery_timeout, reset_timeout.
+    timeout_config: dict
+        Configuration for various operation timeouts. Keys: pull, stop, health_check, install, server_startup.
     """
 
     name = "ollama"
-    defaults = {"models": [], "backend": "local", "debug": False}
+    defaults = {
+        "models": [],
+        "backend": "local",
+        "force_pull": False,
+        "cache_update_policy": "auto",  # "auto", "force", "never"
+        "force_cache_update": False,  # Simple override for "force"
+        "debug": False,
+        "circuit_breaker_config": {
+            "failure_threshold": 3,
+            "recovery_timeout": 60,
+            "reset_timeout": 30,
+        },
+        "timeout_config": {
+            "pull": 600,  # 10 minutes for model pulls
+            "stop": 30,  # 30 seconds for model stops
+            "health_check": 5,  # 5 seconds for health checks
+            "install": 60,  # 1 minute for Ollama installation
+            "server_startup": 300,  # 5 minutes for server startup
+        },
+        "card_refresh_interval": 10,  # seconds - how often to update the status card
+    }
+
+    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
+
+        # Attach the ollama status card
+        self.attach_card_decorator(
+            flow,
+            step_name,
+            "ollama_status",
+            "blank",
+            refresh_interval=self.attributes["card_refresh_interval"],
+        )
 
     def task_decorate(
         self, step_func, flow, graph, retry_count, max_user_code_retries, ubf_context
     ):
         @functools.wraps(step_func)
         def ollama_wrapper():
+            self.ollama_manager = None
+            self.request_interceptor = None
+            self.status_card = None
+            self.card_monitor_thread = None
+
             try:
+                # Initialize status card and monitoring
+                self.status_card = OllamaStatusCard(
+                    refresh_interval=self.attributes["card_refresh_interval"]
+                )
+
+                # Start card monitoring in background
+                def monitor_card():
+                    try:
+                        self.status_card.on_startup(current.card["ollama_status"])
+
+                        while not getattr(
+                            self.card_monitor_thread, "_stop_event", False
+                        ):
+                            try:
+                                # Trigger card update with current data
+                                self.status_card.on_update(
+                                    current.card["ollama_status"], None
+                                )
+                                import time
+
+                                time.sleep(self.attributes["card_refresh_interval"])
+                            except Exception as e:
+                                if self.attributes["debug"]:
+                                    print(f"[@ollama] Card monitoring error: {e}")
+                                break
+                    except Exception as e:
+                        if self.attributes["debug"]:
+                            print(f"[@ollama] Card monitor thread error: {e}")
+                        self.status_card.on_error(current.card["ollama_status"], str(e))
+
+                self.card_monitor_thread = threading.Thread(
+                    target=monitor_card, daemon=True
+                )
+                self.card_monitor_thread._stop_event = False
+                self.card_monitor_thread.start()
+
+                # Initialize OllamaManager with status card
                 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"],
+                    cache_update_policy=self.attributes["cache_update_policy"],
+                    force_cache_update=self.attributes["force_cache_update"],
                     debug=self.attributes["debug"],
+                    circuit_breaker_config=self.attributes["circuit_breaker_config"],
+                    timeout_config=self.attributes["timeout_config"],
+                    status_card=self.status_card,
                 )
+
+                # Install request protection by monkey-patching ollama package
+                self.request_interceptor = OllamaRequestInterceptor(
+                    self.ollama_manager.circuit_breaker, self.attributes["debug"]
+                )
+                self.request_interceptor.install_protection()
+
+                if self.attributes["debug"]:
+                    print(
+                        "[@ollama] OllamaManager initialized and request protection installed"
+                    )
+
             except Exception as e:
+                if self.status_card:
+                    self.status_card.add_event(
+                        "error", f"Initialization failed: {str(e)}"
+                    )
+                    try:
+                        self.status_card.on_error(current.card["ollama_status"], str(e))
+                    except:
+                        pass
                 print(f"[@ollama] Error initializing OllamaManager: {e}")
                 raise
+
             try:
+                if self.status_card:
+                    self.status_card.add_event("info", "Starting user step function")
                 step_func()
+                if self.status_card:
+                    self.status_card.add_event(
+                        "success", "User step function completed successfully"
+                    )
             finally:
-                try:
+                # Remove request protection first (before terminating models)
+                if self.request_interceptor:
+                    self.request_interceptor.remove_protection()
+                    if self.attributes["debug"]:
+                        print("[@ollama] Request protection removed")
+
+                # Then cleanup ollama manager (while card monitoring is still active)
+                if self.ollama_manager:
                     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}")
+
+                # Give the card a moment to render the final shutdown events
+                if self.card_monitor_thread and self.status_card:
+                    import time
+
+                    # Trigger one final card update to capture all shutdown events
+                    try:
+                        self.status_card.on_update(current.card["ollama_status"], None)
+                    except Exception as e:
+                        if self.attributes["debug"]:
+                            print(f"[@ollama] Final card update error: {e}")
+                    time.sleep(2)  # Allow final events to be rendered
+
+                # Now stop card monitoring
+                if self.card_monitor_thread:
+                    self.card_monitor_thread._stop_event = True
+
+                if self.ollama_manager and self.attributes["debug"]:
+                    print(
+                        f"[@ollama] process statuses: {self.ollama_manager.processes}"
+                    )
+                    print(
+                        f"[@ollama] process runtime stats: {self.ollama_manager.stats}"
+                    )
+                    print(
+                        f"[@ollama] Circuit Breaker status: {self.ollama_manager.circuit_breaker.get_status()}"
+                    )
 
         return ollama_wrapper

metaflow_extensions/outerbounds/plugins/ollama/constants.py

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

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)