ob-metaflow-extensions 1.1.128__py2.py3-none-any.whl → 1.1.129__py2.py3-none-any.whl

metaflow_extensions/outerbounds/plugins/nim/__init__.py

@@ -1,10 +1,17 @@
 from functools import partial
-from metaflow.decorators import FlowDecorator
+from uuid import uuid4
+import os, time
+from metaflow.decorators import StepDecorator
 from metaflow import current
+
 from .nim_manager import NimManager
+from .card import NimMetricsRefresher
+from .utilities import get_storage_path, NIM_MONITOR_LOCAL_STORAGE_ROOT
+from ..card_utilities.async_cards import AsyncPeriodicRefresher
+from ..card_utilities.injector import CardDecoratorInjector
 
 
-class NimDecorator(FlowDecorator):
+class NimDecorator(StepDecorator, CardDecoratorInjector):
     """
     This decorator is used to run NIM containers in Metaflow tasks as sidecars.
 
@@ -18,13 +25,12 @@ class NimDecorator(FlowDecorator):
     Valid backend options
     ---------------------
     - 'managed': Outerbounds selects a compute provider based on the model.
-    - 🚧 'dataplane': Run in your account.
 
     Valid model options
     ----------------
         - 'meta/llama3-8b-instruct': 8B parameter model
         - 'meta/llama3-70b-instruct': 70B parameter model
-        - Upon request, any model here: https://nvcf.ngc.nvidia.com/functions?filter=nvidia-functions
+        - any model here: https://nvcf.ngc.nvidia.com/functions?filter=nvidia-functions
 
     Parameters
     ----------
@@ -32,21 +38,100 @@ class NimDecorator(FlowDecorator):
         List of NIM containers running models in sidecars.
     backend: str
         Compute provider to run the NIM container.
+    queue_timeout : int
+        Time to keep the job in NVCF's queue.
     """
 
     name = "nim"
     defaults = {
         "models": [],
         "backend": "managed",
+        "monitor": True,
+        "persist_db": False,
+        "queue_timeout": 5 * 24 * 3600,  # Default 5 days in seconds
     }
 
-    def flow_init(
-        self, flow, graph, environment, flow_datastore, metadata, logger, echo, options
+    def step_init(
+        self, flow, graph, step_name, decorators, environment, flow_datastore, logger
     ):
+
+        if self.attributes["monitor"]:
+            self.attach_card_decorator(
+                flow,
+                step_name,
+                NimMetricsRefresher.CARD_ID,
+                "blank",
+                refresh_interval=4.0,
+            )
+
         current._update_env(
             {
                 "nim": NimManager(
-                    models=self.attributes["models"], backend=self.attributes["backend"]
+                    models=self.attributes["models"],
+                    backend=self.attributes["backend"],
+                    flow=flow,
+                    step_name=step_name,
+                    monitor=self.attributes["monitor"],
+                    queue_timeout=self.attributes["queue_timeout"],
                 )
             }
         )
+
+    def task_decorate(
+        self, step_func, flow, graph, retry_count, max_user_code_retries, ubf_context
+    ):
+        if self.attributes["monitor"]:
+
+            import sqlite3
+            from metaflow import current
+
+            file_path = get_storage_path(current.task_id)
+            if os.path.exists(file_path):
+                os.remove(file_path)
+            os.makedirs(NIM_MONITOR_LOCAL_STORAGE_ROOT, exist_ok=True)
+            conn = sqlite3.connect(file_path)
+            cursor = conn.cursor()
+            cursor.execute(
+                """
+                CREATE TABLE metrics (
+                    error INTEGER,
+                    success INTEGER,
+                    status_code INTEGER,
+                    prompt_tokens INTEGER,
+                    completion_tokens INTEGER,
+                    e2e_time NUMERIC,
+                    model TEXT
+                )
+                """
+            )
+
+            def _wrapped_step_func(*args, **kwargs):
+                async_refresher_metrics = AsyncPeriodicRefresher(
+                    NimMetricsRefresher(),
+                    updater_interval=4.0,
+                    collector_interval=2.0,
+                    file_name=file_path,
+                )
+                try:
+                    async_refresher_metrics.start()
+                    return step_func(*args, **kwargs)
+                finally:
+                    time.sleep(5.0)  # buffer for the last update to synchronize
+                    async_refresher_metrics.stop()
+
+            return _wrapped_step_func
+        else:
+            return step_func
+
+    def task_post_step(
+        self, step_name, flow, graph, retry_count, max_user_code_retries
+    ):
+        if not self.attributes["persist_db"]:
+            import shutil
+
+            file_path = get_storage_path(current.task_id)
+            if os.path.exists(file_path):
+                os.remove(file_path)
+            # if this task is the last one, delete the whole enchilada.
+            if not os.listdir(NIM_MONITOR_LOCAL_STORAGE_ROOT):
+                shutil.rmtree(NIM_MONITOR_LOCAL_STORAGE_ROOT, ignore_errors=True)

metaflow_extensions/outerbounds/plugins/nim/card.py

@@ -0,0 +1,154 @@
+import os, sqlite3
+from metaflow.cards import (
+    Markdown,
+    Table,
+    ProgressBar,
+)
+from metaflow.decorators import StepDecorator
+from metaflow.metaflow_current import current
+
+from .utilities import get_storage_path
+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"
+
+    def __init__(self) -> None:
+        self._metrics_charts = {}
+        self._last_updated_on = None
+        self._already_rendered = False
+        self._file_name = get_storage_path(current.task_id)
+
+    def sqlite_fetch_func(self, conn):
+        cursor = conn.cursor()
+        try:
+            conn = sqlite3.connect(self._file_name)
+            cursor = conn.cursor()
+            cursor.execute(
+                "SELECT error, success, status_code, prompt_tokens, completion_tokens, e2e_time, model FROM metrics"
+            )
+            rows = cursor.fetchall()
+            data = {
+                "error": 0,
+                "success": 0,
+                "status_code": [],
+                "prompt_tokens": [],
+                "completion_tokens": [],
+                "e2e_time": [],
+                "model": [],
+            }
+            for row in rows:
+                data["error"] += row[0]
+                data["success"] += row[1]
+                data["status_code"].append(row[2])
+                data["prompt_tokens"].append(row[3])
+                data["completion_tokens"].append(row[4])
+                data["e2e_time"].append(row[5])
+                data["model"].append(row[6])
+            return data
+        finally:
+            conn.close()
+
+    def render_card_fresh(self, current_card, data):
+        self._already_rendered = True
+        current_card.clear()
+        current_card.append(Markdown("## Metrics"))
+
+        self._metrics_charts["request_success"] = BarPlot(
+            title="Request success",
+            category_name="category",
+            value_name="amount",
+            orientation="horizontal",
+        )
+        self._metrics_charts["latency_distribution"] = ViolinPlot(
+            title="Latency distribution (s)",
+            category_col_name="model",
+            value_col_name="e2e_time",
+        )
+
+        current_card.append(
+            Table(
+                data=[
+                    [
+                        self._metrics_charts["request_success"],
+                    ],
+                    [self._metrics_charts["latency_distribution"]],
+                ]
+            )
+        )
+        current_card.refresh()
+
+    def on_startup(self, current_card):
+        current_card.append(Markdown("# Task-level NIM API metrics"))
+        current_card.append(
+            Markdown(
+                "_waiting for data to appear_",
+            )
+        )
+        current_card.refresh()
+
+    def on_error(self, current_card, error_message):
+
+        if isinstance(error_message, FileNotFoundError):
+            return
+
+        if not self._already_rendered:
+            current_card.clear()
+            current_card.append(
+                Markdown(
+                    f"## Error: {str(error_message)}",
+                )
+            )
+            current_card.refresh()
+
+    def update_only_components(self, current_card, data_object):
+
+        # update request success data
+        self._metrics_charts["request_success"].spec["data"][0]["values"] = [
+            {
+                "category": "Successful requests",
+                "amount": data_object["metrics"]["success"],
+            },
+            {"category": "Errors", "amount": data_object["metrics"]["error"]},
+        ]
+
+        latency_data = []
+        times = []
+        for m, e in zip(
+            data_object["metrics"]["model"], data_object["metrics"]["e2e_time"]
+        ):
+            latency_data.append({"model": m, "e2e_time": e})
+            times.append(e)
+
+        # update latency data
+        self._metrics_charts["latency_distribution"].spec["data"][0][
+            "values"
+        ] = latency_data
+
+        # update domain for latency plot
+        min_time = min(times)
+        max_time = max(times)
+        for scale in self._metrics_charts["latency_distribution"].spec["scales"]:
+            if scale["name"] == "xscale":
+                scale["domain"] = [min_time - max_time * 0.1, max_time + max_time * 0.1]
+
+        current_card.refresh()
+
+    def on_update(self, current_card, data_object):
+        data_object_keys = set(data_object.keys())
+        if len(data_object_keys) == 0:
+            return
+        if len(self._metrics_charts) == 0:
+            self.render_card_fresh(current_card, data_object)
+            return
+        elif len(data_object["metrics"]["status_code"]) == 0:
+            return
+        else:
+            self.update_only_components(current_card, data_object)
+            return

metaflow_extensions/outerbounds/plugins/nim/nim_manager.py

@@ -1,19 +1,20 @@
-import os
-import time
-import json
-import requests
+import os, sys, time, json, random, requests, sqlite3
 from urllib.parse import urlparse
 from metaflow.metaflow_config import SERVICE_URL
 from metaflow.metaflow_config_funcs import init_config
-import sys
-import random
+from .utilities import get_storage_path
+from ..nvcf.nvcf import retry_on_status
+
 
 NVCF_URL = "https://api.nvcf.nvidia.com"
 NVCF_SUBMIT_ENDPOINT = f"{NVCF_URL}/v2/nvcf/pexec/functions"
 NVCF_RESULT_ENDPOINT = f"{NVCF_URL}/v2/nvcf/pexec/status"
-
-COMMON_HEADERS = {"accept": "application/json", "Content-Type": "application/json"}
-POLL_INTERVAL = 1
+NVCF_POLL_INTERVAL_SECONDS = 1
+COMMON_HEADERS = {
+    "accept": "application/json",
+    "Content-Type": "application/json",
+    "nvcf-feature-enable-gateway-timeout": "true",
+}
 
 
 class NimMetadata(object):
@@ -56,41 +57,33 @@ class NimMetadata(object):
     def get_nvcf_chat_completion_models(self):
         return self._nvcf_chat_completion_models
 
-    def get_coreweave_chat_completion_models(self):
-        return self._coreweave_chat_completion_models
-
     def get_headers_for_nvcf_request(self):
         return {**COMMON_HEADERS, "Authorization": f"Bearer {self._ngc_api_key}"}
 
-    def get_headers_for_coreweave_request(self):
-        return COMMON_HEADERS
-
 
 class NimManager(object):
-    def __init__(self, models, backend):
+    def __init__(self, models, backend, flow, step_name, monitor, queue_timeout):
+
         nim_metadata = NimMetadata()
         if backend == "managed":
             nvcf_models = [
                 m["name"] for m in nim_metadata.get_nvcf_chat_completion_models()
             ]
-            cw_models = [
-                m["name"] for m in nim_metadata.get_coreweave_chat_completion_models()
-            ]
 
             self.models = {}
             for m in models:
                 if m in nvcf_models:
                     self.models[m] = NimChatCompletion(
-                        model=m, provider="NVCF", nim_metadata=nim_metadata
-                    )
-                elif m in cw_models:
-                    self.models[m] = NimChatCompletion(
-                        model=m, provider="CoreWeave", nim_metadata=nim_metadata
+                        model=m,
+                        provider="NVCF",
+                        nim_metadata=nim_metadata,
+                        monitor=monitor,
+                        queue_timeout=queue_timeout,
                     )
                 else:
                     raise ValueError(
                         f"Model {m} not supported by the Outerbounds @nim offering."
-                        f"\nYou can choose from these options: {nvcf_models + cw_models}\n\n"
+                        f"\nYou can choose from these options: {nvcf_models}\n\n"
                         "Reach out to Outerbounds if there are other models you'd like supported."
                     )
         else:
@@ -99,12 +92,21 @@ class NimManager(object):
             )
 
 
+class JobStatus(object):
+    SUBMITTED = "SUBMITTED"
+    RUNNING = "RUNNING"
+    SUCCESSFUL = "SUCCESSFUL"
+    FAILED = "FAILED"
+
+
 class NimChatCompletion(object):
     def __init__(
         self,
         model="meta/llama3-8b-instruct",
-        provider="CoreWeave",
+        provider="NVCF",
         nim_metadata=None,
+        monitor=False,
+        queue_timeout=None,
         **kwargs,
     ):
         if nim_metadata is None:
@@ -118,19 +120,9 @@ class NimChatCompletion(object):
         self.max_request_retries = int(
             os.environ.get("METAFLOW_EXT_HTTP_MAX_RETRIES", "10")
         )
+        self.monitor = monitor
 
-        if self.compute_provider == "CoreWeave":
-            cw_model_names = [
-                m["name"]
-                for m in self._nim_metadata.get_coreweave_chat_completion_models()
-            ]
-            self.model = model
-            self.ip_address = self._nim_metadata.get_coreweave_chat_completion_models()[
-                cw_model_names.index(model)
-            ]["ip-address"]
-            self.endpoint = f"http://{self.ip_address}:8000/v1/chat/completions"
-
-        elif self.compute_provider == "NVCF":
+        if self.compute_provider == "NVCF":
             nvcf_model_names = [
                 m["name"] for m in self._nim_metadata.get_nvcf_chat_completion_models()
             ]
@@ -141,45 +133,175 @@ class NimChatCompletion(object):
             self.version_id = self._nim_metadata.get_nvcf_chat_completion_models()[
                 nvcf_model_names.index(model)
             ]["version-id"]
+        else:
+            raise ValueError(
+                f"Backend compute provider {self.compute_provider} not yet supported for @nim."
+            )
 
-    def __call__(self, **kwargs):
+        # to know whether to set file_name
+        self.first_request = True
+
+        # TODO (Eddie) - this may make more sense in a base class.
+        # @nim arch needs redesign if customers start using it in more creative ways.
+        self._poll_seconds = "3600"
+        self._queue_timeout = queue_timeout
+        self._status = None
+        self._result = {}
+
+    @property
+    def status(self):
+        return self._status
+
+    @property
+    def has_failed(self):
+        return self._status == JobStatus.FAILED
+
+    @property
+    def is_running(self):
+        return self._status == JobStatus.SUBMITTED
+
+    @property
+    def result(self):
+        return self._result
+
+    def _log_stats(self, response, e2e_time):
+        stats = {}
+        if response.status_code == 200:
+            stats["success"] = 1
+            stats["error"] = 0
+        else:
+            stats["success"] = 0
+            stats["error"] = 1
+        stats["status_code"] = response.status_code
+        try:
+            stats["prompt_tokens"] = response.json()["usage"]["prompt_tokens"]
+        except KeyError:
+            stats["prompt_tokens"] = None
+        try:
+            stats["completion_tokens"] = response.json()["usage"]["completion_tokens"]
+        except KeyError:
+            stats["completion_tokens"] = None
+        stats["e2e_time"] = e2e_time
+        stats["provider"] = self.compute_provider
+        stats["model"] = self.model
 
-        if self.compute_provider == "CoreWeave":
-            request_data = {"model": self.model, **kwargs}
-            response = requests.post(
-                self.endpoint,
-                headers=self._nim_metadata.get_headers_for_coreweave_request(),
-                json=request_data,
+        conn = sqlite3.connect(self.file_name)
+        cursor = conn.cursor()
+        try:
+            cursor.execute(
+                """
+                INSERT INTO metrics (error, success, status_code, prompt_tokens, completion_tokens, e2e_time, model)
+                VALUES (?, ?, ?, ?, ?, ?, ?)
+            """,
+                (
+                    stats["error"],
+                    stats["success"],
+                    stats["status_code"],
+                    stats["prompt_tokens"],
+                    stats["completion_tokens"],
+                    stats["e2e_time"],
+                    stats["model"],
+                ),
             )
-            response.raise_for_status()
-            return response.json()
+            conn.commit()
+        finally:
+            conn.close()
 
-        elif self.compute_provider == "NVCF":
+    @retry_on_status(status_codes=[504])
+    def __call__(self, **kwargs):
+
+        if self.first_request:
+            # Put here to guarantee self.file_name is set after task_id exists.
+            from metaflow import current
 
-            request_data = {"model": self.model, **kwargs}
-            request_url = f"{NVCF_SUBMIT_ENDPOINT}/{self.function_id}"
+            self.file_name = get_storage_path(current.task_id)
+
+        request_data = {"model": self.model, **kwargs}
+        request_url = f"{NVCF_SUBMIT_ENDPOINT}/{self.function_id}"
+        retry_delay = 1
+        attempts = 0
+        t0 = time.time()
+        while attempts < self.max_request_retries:
+            try:
+                attempts += 1
+                response = requests.post(
+                    request_url,
+                    headers=self._nim_metadata.get_headers_for_nvcf_request(),
+                    json=request_data,
+                )
+                if response.status_code == 202:
+                    invocation_id = response.headers.get("NVCF-REQID")
+                    self.invocations.append(invocation_id)
+                    self._status = JobStatus.SUBMITTED
+                elif response.status_code == 200:
+                    tf = time.time()
+                    if self.monitor:
+                        self._log_stats(response, tf - t0)
+                    self._status = JobStatus.SUCCESSFUL
+                    self._result = response.json()
+                    return self._result
+                elif response.status_code == 400:
+                    self._status = JobStatus.FAILED
+                    msg = (
+                        "[@nim ERROR] The OpenAI-compatible returned a 400 status code. "
+                        + "Known causes include improper requests or prompts with too many tokens for the selected model. "
+                        + "Please contact Outerbounds if you need assistance resolving the issue."
+                    )
+                    print(msg, file=sys.stderr)
+                    self._result = {"ERROR": msg}
+                    return self._result
+            except (
+                requests.exceptions.ConnectionError,
+                requests.exceptions.ReadTimeout,
+            ) as e:
+                # ConnectionErrors are generally temporary errors like DNS resolution failures,
+                # timeouts etc.
+                print(
+                    "received error of type {}. Retrying...".format(type(e)),
+                    e,
+                    file=sys.stderr,
+                )
+                time.sleep(retry_delay)
+                retry_delay *= 2  # Double the delay for the next attempt
+                retry_delay += random.uniform(0, 1)  # Add jitter
+                retry_delay = min(retry_delay, 10)
 
+        @retry_on_status(status_codes=[500], max_retries=3, delay=5)
+        @retry_on_status(status_codes=[504])
+        def _poll():
+            poll_request_url = f"{NVCF_RESULT_ENDPOINT}/{invocation_id}"
             attempts = 0
+            retry_delay = 1
             while attempts < self.max_request_retries:
                 try:
                     attempts += 1
-                    response = requests.post(
-                        request_url,
+                    poll_response = requests.get(
+                        poll_request_url,
                         headers=self._nim_metadata.get_headers_for_nvcf_request(),
-                        json=request_data,
                     )
-                    response.raise_for_status()
-                    if response.status_code == 202:
-                        invocation_id = response.headers.get("NVCF-REQID")
-                        self.invocations.append(invocation_id)
-                    elif response.status_code == 200:
-                        return response.json()
+                    if poll_response.status_code == 200:
+                        tf = time.time()
+                        self._log_stats(response, tf - t0)
+                        self._status = JobStatus.SUCCESSFUL
+                        self._result = poll_response.json()
+                        return self._result
+                    elif poll_response.status_code == 202:
+                        self._status = JobStatus.SUBMITTED
+                        return 202
+                    elif poll_response.status_code == 400:
+                        self._status = JobStatus.FAILED
+                        msg = (
+                            "[@nim ERROR] The OpenAI-compatible API returned a 400 status code. "
+                            + "Known causes include improper requests or prompts with too many tokens for the selected model. "
+                            + "Please contact Outerbounds if you need assistance resolving the issue."
+                        )
+                        print(msg, file=sys.stderr)
+                        self._result = {"@nim ERROR": msg}
+                        return self._result
                 except (
                     requests.exceptions.ConnectionError,
                     requests.exceptions.ReadTimeout,
                 ) as e:
-                    # ConnectionErrors are generally temporary errors like DNS resolution failures,
-                    # timeouts etc.
                     print(
                         "received error of type {}. Retrying...".format(type(e)),
                         e,
@@ -190,44 +312,8 @@ class NimChatCompletion(object):
                     retry_delay += random.uniform(0, 1)  # Add jitter
                     retry_delay = min(retry_delay, 10)
 
-            def _poll():
-                poll_request_url = f"{NVCF_RESULT_ENDPOINT}/{invocation_id}"
-                attempts = 0
-
-                while attempts < self.max_request_retries:
-                    try:
-                        attempts += 1
-                        poll_response = requests.get(
-                            poll_request_url,
-                            headers=self._nim_metadata.get_headers_for_nvcf_request(),
-                        )
-                        poll_response.raise_for_status()
-                        if poll_response.status_code == 200:
-                            return poll_response.json()
-                        elif poll_response.status_code == 202:
-                            return 202
-                        else:
-                            raise Exception(
-                                f"NVCF returned {poll_response.status_code} status code. Please contact Outerbounds."
-                            )
-                    except (
-                        requests.exceptions.ConnectionError,
-                        requests.exceptions.ReadTimeout,
-                    ) as e:
-                        # ConnectionErrors are generally temporary errors like DNS resolution failures,
-                        # timeouts etc.
-                        print(
-                            "received error of type {}. Retrying...".format(type(e)),
-                            e,
-                            file=sys.stderr,
-                        )
-                        time.sleep(retry_delay)
-                        retry_delay *= 2  # Double the delay for the next attempt
-                        retry_delay += random.uniform(0, 1)  # Add jitter
-                        retry_delay = min(retry_delay, 10)
-
-            while True:
-                data = _poll()
-                if data and data != 202:
-                    return data
-                time.sleep(POLL_INTERVAL)
+        while True:
+            data = _poll()
+            if data and data != 202:
+                return data
+            time.sleep(NVCF_POLL_INTERVAL_SECONDS)

metaflow_extensions/outerbounds/plugins/nim/utilities.py

@@ -0,0 +1,5 @@
+NIM_MONITOR_LOCAL_STORAGE_ROOT = ".nim-monitor"
+
+
+def get_storage_path(task_id):
+    return f"{NIM_MONITOR_LOCAL_STORAGE_ROOT}/" + task_id + ".sqlite"

{ob_metaflow_extensions-1.1.128.dist-info → ob_metaflow_extensions-1.1.129.dist-info}/METADATA

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