ob-metaflow-extensions 1.1.155__py2.py3-none-any.whl → 1.1.156__py2.py3-none-any.whl

metaflow_extensions/outerbounds/plugins/__init__.py

@@ -326,7 +326,7 @@ STEP_DECORATORS_DESC = [
     ("snowpark", ".snowpark.snowpark_decorator.SnowparkDecorator"),
     ("tensorboard", ".tensorboard.TensorboardDecorator"),
     ("gpu_profile", ".profilers.gpu_profile_decorator.GPUProfileDecorator"),
-    ("nim", ".nim.NimDecorator"),
+    ("nim", ".nim.nim_decorator.NimDecorator"),
     ("ollama", ".ollama.OllamaDecorator"),
     ("app_deploy", ".apps.deploy_decorator.WorkstationAppDeployDecorator"),
 ]

metaflow_extensions/outerbounds/plugins/nim/card.py

@@ -1,8 +1,7 @@
-import sqlite3
 from metaflow.cards import Markdown, Table
 from metaflow.metaflow_current import current
 
-from .utilities import get_storage_path
+from .utils import get_storage_path
 from ..card_utilities.async_cards import CardRefresher
 from ..card_utilities.extra_components import BarPlot, ViolinPlot
 
@@ -17,9 +16,7 @@ class NimMetricsRefresher(CardRefresher):
         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"
@@ -85,7 +82,6 @@ class NimMetricsRefresher(CardRefresher):
         current_card.refresh()
 
     def on_error(self, current_card, error_message):
-
         if isinstance(error_message, FileNotFoundError):
             return
 
@@ -99,7 +95,6 @@ class NimMetricsRefresher(CardRefresher):
             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"] = [
             {

metaflow_extensions/outerbounds/plugins/nim/{__init__.py → nim_decorator.py}

@@ -1,64 +1,31 @@
-from functools import partial
-from uuid import uuid4
-import os, time
-from metaflow.decorators import StepDecorator
+import os
+import time
 from metaflow import current
-
+from .utils import get_storage_path, NIM_MONITOR_LOCAL_STORAGE_ROOT
 from .nim_manager import NimManager
+from metaflow.decorators import StepDecorator
 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
+from ..card_utilities.async_cards import AsyncPeriodicRefresher
 
 
 class NimDecorator(StepDecorator, CardDecoratorInjector):
-    """
-    This decorator is used to run NIM containers in Metaflow tasks as sidecars.
-
-    User code call
-    -----------
-    @nim(
-        models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
-        backend='managed'
-    )
-
-    Valid backend options
-    ---------------------
-    - 'managed': Outerbounds selects a compute provider based on the model.
-
-    Valid model options
-    ----------------
-        - 'meta/llama3-8b-instruct': 8B parameter model
-        - 'meta/llama3-70b-instruct': 70B parameter model
-        - any model here: https://nvcf.ngc.nvidia.com/functions?filter=nvidia-functions
-
-    Parameters
-    ----------
-    models: list[NIM]
-        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 step_init(
-        self, flow, graph, step_name, decorators, environment, flow_datastore, logger
-    ):
-
+    # Refer https://github.com/Netflix/metaflow/blob/master/docs/lifecycle.png
+    # to understand where these functions are invoked in the lifecycle of a
+    # Metaflow flow.
+    def step_init(self, flow, graph, step, decos, environment, flow_datastore, logger):
         if self.attributes["monitor"]:
             self.attach_card_decorator(
                 flow,
-                step_name,
+                step,
                 NimMetricsRefresher.CARD_ID,
                 "blank",
                 refresh_interval=4.0,
@@ -68,11 +35,9 @@ class NimDecorator(StepDecorator, CardDecoratorInjector):
             {
                 "nim": NimManager(
                     models=self.attributes["models"],
-                    backend=self.attributes["backend"],
                     flow=flow,
-                    step_name=step_name,
+                    step_name=step,
                     monitor=self.attributes["monitor"],
-                    queue_timeout=self.attributes["queue_timeout"],
                 )
             }
         )
@@ -81,15 +46,14 @@ class NimDecorator(StepDecorator, CardDecoratorInjector):
         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(
                 """

metaflow_extensions/outerbounds/plugins/nim/nim_manager.py

@@ -1,47 +1,163 @@
-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
-from .utilities import get_storage_path
-from ..nvcf.nvcf import retry_on_status
+import sys
+import time
+import requests
+import sqlite3
+from urllib3.util.retry import Retry
+from requests.adapters import HTTPAdapter
+from typing import Dict, Optional, Any
+from .utils import get_ngc_response, get_storage_path
+
+
+def nvcf_submit_helper(
+    url: str,
+    payload: Dict[str, Any],
+    headers: Optional[Dict[str, str]] = None,
+    timeout: int = 30,
+    max_retries: int = 300,
+    backoff_factor: float = 0.3,
+    request_delay: float = 1.1,
+    log_callback: Optional[callable] = None,
+) -> Dict[str, Any]:
+    def _log_error(start_time: float, status_code: int, poll_count: int):
+        if log_callback:
+            end_time = time.time()
+            try:
+                log_callback({}, end_time - start_time, status_code, poll_count)
+            except Exception as log_error:
+                print(f"Warning: Logging callback failed: {log_error}")
+
+    # use default headers
+    if not headers:
+        headers = {"accept": "application/json", "content-type": "application/json"}
+        print(f"Using Default Headers: {headers}")
+
+    # Configure session with retry strategy
+    session = requests.Session()
+    status_forcelist = [429, 500, 502, 503, 504, 404]
+    retry_strategy = Retry(
+        total=max_retries,
+        backoff_factor=backoff_factor,
+        status_forcelist=status_forcelist,
+        allowed_methods=["GET", "POST"],
+    )
+    adapter = HTTPAdapter(max_retries=retry_strategy)
+    session.mount("http://", adapter)
+    session.mount("https://", adapter)
+
+    # Add artificial delay if specified
+    time.sleep(request_delay)
+
+    start_time = time.time()
+    poll_count = 0
+    status_code = 0
+    response_data = {}
+
+    try:
+        # Make initial request
+        response = session.post(url, json=payload, headers=headers, timeout=timeout)
+        time.sleep(request_delay)
+
+        # Handle initial response
+        response.raise_for_status()
+        request_id = response.headers.get("NVCF-REQID")
+        polling_url = f"https://api.nvcf.nvidia.com/v2/nvcf/pexec/status/{request_id}"
+
+        print(f"Polling NVCF Request ID: {request_id}")
+
+        # Initial response status
+        status_code = response.status_code
+        print(f"Initial response status: {status_code}")
+
+        # Create a variable to store the final response
+        final_response = response
+
+        # Continue polling while we get 202 (Accepted/Processing)
+        while status_code == 202:
+            poll_count += 1
+            print(f"Polling attempt #{poll_count} to {polling_url}")
+
+            # Wait before next poll
+            time.sleep(request_delay)
+
+            # Make a new poll request
+            poll_response = session.get(polling_url, headers=headers, timeout=timeout)
+            status_code = poll_response.status_code
+            print(f"Poll #{poll_count} status: {status_code}")
+
+            # Check for errors
+            try:
+                poll_response.raise_for_status()
+            except requests.exceptions.HTTPError as e:
+                print(f"Poll request failed: {str(e)}")
+                poll_response.close()
+                # Log the error before re-raising
+                _log_error(start_time, poll_response.status_code, poll_count)
+                raise
+
+            # If status is 200, the job is complete
+            if status_code == 200:
+                print("Polling complete - job finished successfully")
+                # Update our final response to be this poll response
+                final_response = poll_response
+                break
+
+            # Close this poll response if we're going to loop again
+            if status_code == 202:
+                poll_response.close()
+
+        # If we exited the loop without a 200 status, something went wrong
+        if status_code != 200:
+            print(f"Polling ended with unexpected status: {status_code}")
+            # Log the error before raising
+            _log_error(start_time, status_code, poll_count)
+            raise Exception(f"Unexpected status code after polling: {status_code}")
+
+        # Get the response data for logging
+        response_data = final_response.json()
+
+    except requests.exceptions.HTTPError as e:
+        # Handle HTTP errors (4xx, 5xx status codes)
+        status_code = e.response.status_code if e.response else 0
+        print(f"HTTP Error: {str(e)}", file=sys.stderr)
+        # Log the error
+        _log_error(start_time, status_code, poll_count)
+        raise
+
+    except Exception as e:
+        # Handle other errors (connection errors, timeouts, etc.)
+        print(f"Request Error: {str(e)}", file=sys.stderr)
+        # Log the error with status_code 0 to indicate non-HTTP error
+        _log_error(start_time, 0, poll_count)
+        raise
+
+    # Calculate final duration and log successful requests
+    end_time = time.time()
+    duration = end_time - start_time
+
+    # Call the logging callback if provided
+    if log_callback:
+        try:
+            log_callback(response_data, duration, status_code, poll_count)
+        except Exception as e:
+            print(f"Warning: Logging callback failed: {e}")
 
+    # Log metrics
+    print(
+        f"Request completed: duration={duration:.2f}s, polls={poll_count}, "
+        f"status={status_code}, size={len(final_response.content)} bytes"
+    )
 
-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"
-NVCF_POLL_INTERVAL_SECONDS = 1
-COMMON_HEADERS = {
-    "accept": "application/json",
-    "Content-Type": "application/json",
-    "nvcf-feature-enable-gateway-timeout": "true",
-}
+    return response_data
 
 
 class NimMetadata(object):
     def __init__(self):
         self._nvcf_chat_completion_models = []
-        self._coreweave_chat_completion_models = []
-
-        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]
+        ngc_response = get_ngc_response()
 
-        nim_info_url = "https://" + auth_host + "/generate/nim"
+        self.ngc_api_key = ngc_response["nvcf"]["api_key"]
 
-        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()
-        self._ngc_api_key = res.json()["nvcf"]["api_key"]
-
-        for model in res.json()["nvcf"]["functions"]:
+        for model in ngc_response["nvcf"]["functions"]:
             self._nvcf_chat_completion_models.append(
                 {
                     "name": model["model_key"],
@@ -49,64 +165,48 @@ class NimMetadata(object):
                     "version-id": model["version"],
                 }
             )
-        for model in res.json()["coreweave"]["containers"]:
-            self._coreweave_chat_completion_models.append(
-                {"name": model["nim_name"], "ip-address": model["ip_addr"]}
-            )
 
     def get_nvcf_chat_completion_models(self):
         return self._nvcf_chat_completion_models
 
     def get_headers_for_nvcf_request(self):
-        return {**COMMON_HEADERS, "Authorization": f"Bearer {self._ngc_api_key}"}
+        return {
+            "accept": "application/json",
+            "content-type": "application/json",
+            "Authorization": f"Bearer {self.ngc_api_key}",
+            "NVCF-POLL-SECONDS": "5",
+        }
 
 
 class NimManager(object):
-    def __init__(self, models, backend, flow, step_name, monitor, queue_timeout):
-
+    def __init__(self, models, flow, step_name, monitor):
         nim_metadata = NimMetadata()
-        if backend == "managed":
-            nvcf_models = [
-                m["name"] for m in nim_metadata.get_nvcf_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,
-                        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}\n\n"
-                        "Reach out to Outerbounds if there are other models you'd like supported."
-                    )
-        else:
-            raise ValueError(
-                f"Backend {backend} not supported by the Outerbounds @nim offering. Please reach out to Outerbounds."
-            )
-
-
-class JobStatus(object):
-    SUBMITTED = "SUBMITTED"
-    RUNNING = "RUNNING"
-    SUCCESSFUL = "SUCCESSFUL"
-    FAILED = "FAILED"
+        nvcf_models = [
+            m["name"] for m in nim_metadata.get_nvcf_chat_completion_models()
+        ]
+        self.models = {}
+
+        for each_model in models:
+            if each_model in nvcf_models:
+                self.models[each_model] = NimChatCompletion(
+                    model=each_model,
+                    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."
+                )
 
 
 class NimChatCompletion(object):
     def __init__(
         self,
-        model="meta/llama3-8b-instruct",
-        provider="NVCF",
-        nim_metadata=None,
-        monitor=False,
-        queue_timeout=None,
+        model: str = "meta/llama3-8b-instruct",
+        nim_metadata: NimMetadata = None,
+        monitor: bool = False,
         **kwargs,
     ):
         if nim_metadata is None:
@@ -114,79 +214,54 @@ class NimChatCompletion(object):
                 "NimMetadata object is required to initialize NimChatCompletion object."
             )
 
-        self._nim_metadata = nim_metadata
-        self.compute_provider = provider
-        self.invocations = []
-        self.max_request_retries = int(
-            os.environ.get("METAFLOW_EXT_HTTP_MAX_RETRIES", "10")
-        )
+        self.model_name = model
+        self.nim_metadata = nim_metadata
         self.monitor = monitor
 
-        if self.compute_provider == "NVCF":
-            nvcf_model_names = [
-                m["name"] for m in self._nim_metadata.get_nvcf_chat_completion_models()
-            ]
-            self.model = model
-            self.function_id = self._nim_metadata.get_nvcf_chat_completion_models()[
-                nvcf_model_names.index(model)
-            ]["function-id"]
-            self.version_id = self._nim_metadata.get_nvcf_chat_completion_models()[
-                nvcf_model_names.index(model)
-            ]["version-id"]
-        else:
+        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"Backend compute provider {self.compute_provider} not yet supported for @nim."
+                f"Model {self.model_name} not found in available NVCF models"
             )
 
-        # to know whether to set file_name
+        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
 
-        # 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
+    def log_stats(self, response_data, duration, status_code, poll_count):
+        if not self.monitor:
+            return
+
+        stats = {
+            "status_code": status_code,
+            "success": 1 if status_code == 200 else 0,
+            "error": 0 if status_code == 200 else 1,
+            "e2e_time": duration,
+            "model": self.model_name,
+            "poll_count": poll_count,
+        }
+
+        if status_code == 200 and response_data:
+            try:
+                stats["prompt_tokens"] = response_data["usage"]["prompt_tokens"]
+            except (KeyError, TypeError):
+                stats["prompt_tokens"] = None
+
+            try:
+                stats["completion_tokens"] = response_data["usage"]["completion_tokens"]
+            except (KeyError, TypeError):
+                stats["completion_tokens"] = None
         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
 
         conn = sqlite3.connect(self.file_name)
         cursor = conn.cursor()
+
         try:
             cursor.execute(
                 """
@@ -207,112 +282,37 @@ class NimChatCompletion(object):
         finally:
             conn.close()
 
-    @retry_on_status(status_codes=[500], max_retries=3, delay=5)
-    @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
 
             self.file_name = get_storage_path(current.task_id)
+            self.first_request = False
 
-        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)
-
-        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
-                    poll_response = requests.get(
-                        poll_request_url,
-                        headers=self._nim_metadata.get_headers_for_nvcf_request(),
-                    )
-                    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:
-                    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(NVCF_POLL_INTERVAL_SECONDS)
+        # Create log callback if monitoring is enabled
+        log_callback = self.log_stats if self.monitor else None
+
+        request_data = {"model": self.model_name, **kwargs}
+        request_url = (
+            f"https://api.nvcf.nvidia.com/v2/nvcf/pexec/functions/{self.function_id}"
+        )
+
+        try:
+            response_data = nvcf_submit_helper(
+                url=request_url,
+                payload=request_data,
+                headers=self.nim_metadata.get_headers_for_nvcf_request(),
+                log_callback=log_callback,
+            )
+
+            return response_data
+
+        except requests.exceptions.HTTPError as e:
+            error_msg = f"[@nim ERROR] NVCF API request failed: {str(e)}"
+            print(error_msg, file=sys.stderr)
+            raise
+
+        except Exception as e:
+            error_msg = f"[@nim ERROR] Unexpected error: {str(e)}"
+            print(error_msg, file=sys.stderr)
+            raise

metaflow_extensions/outerbounds/plugins/nim/utils.py

@@ -0,0 +1,36 @@
+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
+
+
+NIM_MONITOR_LOCAL_STORAGE_ROOT = ".nim-monitor"
+
+
+def get_storage_path(task_id):
+    return f"{NIM_MONITOR_LOCAL_STORAGE_ROOT}/" + task_id + ".sqlite"
+
+
+def get_ngc_response():
+    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()

{ob_metaflow_extensions-1.1.155.dist-info → ob_metaflow_extensions-1.1.156.dist-info}/METADATA

@@ -1,13 +1,13 @@
 Metadata-Version: 2.1
 Name: ob-metaflow-extensions
-Version: 1.1.155
+Version: 1.1.156
 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.15.13.1)
+Requires-Dist: ob-metaflow (==2.15.14.1)
 
 # Outerbounds platform package
 

{ob_metaflow_extensions-1.1.155.dist-info → ob_metaflow_extensions-1.1.156.dist-info}/RECORD

@@ -1,7 +1,7 @@
 metaflow_extensions/outerbounds/__init__.py,sha256=Gb8u06s9ClQsA_vzxmkCzuMnigPy7kKcDnLfb7eB-64,514
 metaflow_extensions/outerbounds/remote_config.py,sha256=pEFJuKDYs98eoB_-ryPjVi9b_c4gpHMdBHE14ltoxIU,4672
 metaflow_extensions/outerbounds/config/__init__.py,sha256=JsQGRuGFz28fQWjUvxUgR8EKBLGRdLUIk_buPLJplJY,1225
-metaflow_extensions/outerbounds/plugins/__init__.py,sha256=eHcM_t2Mzlge7B9Dv3VGVM5x8qNZYdLyqBOAC6uRxec,13228
+metaflow_extensions/outerbounds/plugins/__init__.py,sha256=gytuNt3lNabirHLEYzrmHFMp-JWh8dA2AZPK11HmaNw,13242
 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/apps/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -23,10 +23,10 @@ metaflow_extensions/outerbounds/plugins/fast_bakery/fast_bakery_cli.py,sha256=kq
 metaflow_extensions/outerbounds/plugins/fast_bakery/fast_bakery_decorator.py,sha256=MXSIp05-jvt8Q2uGaLKjtuM_ToLeRLxhtMbfHc9Kcko,1515
 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=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/nim/card.py,sha256=dXOJvsZed5NyYyxYLPDvtwg9z_X4azL9HTJGYaiNriY,4690
+metaflow_extensions/outerbounds/plugins/nim/nim_decorator.py,sha256=50YVvC7mcZYlPluM0Wq1UtufhzlQb-RxzZkTOJJ3LkM,3439
+metaflow_extensions/outerbounds/plugins/nim/nim_manager.py,sha256=5YkohM-vfoDHPUMWb19sY0HErORoKOKf4jexERJTO80,10912
+metaflow_extensions/outerbounds/plugins/nim/utils.py,sha256=nU-v1sheBjmITXfHiJx2ucm_Tq_nGb5BcuAm5c235cQ,1164
 metaflow_extensions/outerbounds/plugins/nvcf/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 metaflow_extensions/outerbounds/plugins/nvcf/constants.py,sha256=aGHdNw_hqBu8i0zWXcatQM6e769wUXox0l8g0f6fNZ8,146
 metaflow_extensions/outerbounds/plugins/nvcf/exceptions.py,sha256=-Pm9cOWUzpv94TvVUeq-FenAWdfLBJd5N7WPqIGZVqU,3671
@@ -68,7 +68,7 @@ metaflow_extensions/outerbounds/toplevel/plugins/gcp/__init__.py,sha256=BbZiaH3u
 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.155.dist-info/METADATA,sha256=1SCxq4RN5keRpny8BRcQn9FJ71aB58cUeIkk2yytPI8,521
-ob_metaflow_extensions-1.1.155.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-ob_metaflow_extensions-1.1.155.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
-ob_metaflow_extensions-1.1.155.dist-info/RECORD,,
+ob_metaflow_extensions-1.1.156.dist-info/METADATA,sha256=G9c19j9g0v8dDQU5sP5Zaaub2fot__EMCJ6iBQBb4Qo,521
+ob_metaflow_extensions-1.1.156.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+ob_metaflow_extensions-1.1.156.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
+ob_metaflow_extensions-1.1.156.dist-info/RECORD,,

metaflow_extensions/outerbounds/plugins/nim/utilities.py

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