ob-metaflow-extensions 1.1.45rc3__py2.py3-none-any.whl → 1.5.1__py2.py3-none-any.whl

metaflow_extensions/outerbounds/plugins/nim/nim_decorator.py

@@ -0,0 +1,101 @@
+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 ..card_utilities.injector import CardDecoratorInjector
+from ..card_utilities.async_cards import AsyncPeriodicRefresher
+
+
+class NimDecorator(StepDecorator, CardDecoratorInjector):
+    name = "nim"
+
+    defaults = {
+        "models": [],
+        "monitor": True,
+        "persist_db": False,
+    }
+
+    # 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,
+                NimMetricsRefresher.CARD_ID,
+                "blank",
+                refresh_interval=4.0,
+            )
+
+        current._update_env(
+            {
+                "nim": NimManager(
+                    models=self.attributes["models"],
+                    flow=flow,
+                    step_name=step,
+                    monitor=self.attributes["monitor"],
+                )
+            }
+        )
+
+    def task_decorate(
+        self, step_func, flow, graph, retry_count, max_user_code_retries, ubf_context
+    ):
+        if self.attributes["monitor"]:
+            import sqlite3
+
+            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/nim_manager.py

@@ -0,0 +1,379 @@
+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"
+    )
+
+    return response_data
+
+
+class NimMetadata(object):
+    def __init__(self):
+        self._nvcf_chat_completion_models = []
+        ngc_response = get_ngc_response()
+
+        self.ngc_api_key = ngc_response["nvcf"]["api_key"]
+
+        for model in ngc_response["nvcf"]["functions"]:
+            self._nvcf_chat_completion_models.append(
+                {
+                    "name": model["model_key"],
+                    "function-id": model["id"],
+                    "version-id": model["version"],
+                }
+            )
+
+    def get_nvcf_chat_completion_models(self):
+        return self._nvcf_chat_completion_models
+
+    def get_headers_for_nvcf_request(self):
+        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, flow, step_name, monitor):
+        nim_metadata = NimMetadata()
+        nvcf_models = [
+            m["name"] for m in nim_metadata.get_nvcf_chat_completion_models()
+        ]
+        self.models = {}
+
+        # Convert models to a standard format
+        standardized_models = []
+        # If models is a single string, convert it to a list with a dict
+        if isinstance(models, str):
+            standardized_models = [{"name": models}]
+        # If models is a list, process each item
+        elif isinstance(models, list):
+            for model_item in models:
+                # If the item is a string, convert it to a dict
+                if isinstance(model_item, str):
+                    standardized_models.append({"name": model_item})
+                # If it's already a dict, use it as is
+                elif isinstance(model_item, dict):
+                    standardized_models.append(model_item)
+                else:
+                    raise ValueError(
+                        f"Model specification must be a string or dictionary, got {type(model_item)}"
+                    )
+        else:
+            raise ValueError(
+                f"Models must be a string or a list of strings/dictionaries, got {type(models)}"
+            )
+
+        # Process each standardized model
+        for each_model_dict in standardized_models:
+            model_name = each_model_dict.get("name", "")
+            nvcf_id = each_model_dict.get("nvcf_id", "")
+            nvcf_version = each_model_dict.get("nvcf_version", "")
+
+            if model_name and not (nvcf_id and nvcf_version):
+                if model_name in nvcf_models:
+                    self.models[model_name] = NimChatCompletion(
+                        model=model_name,
+                        nvcf_id=nvcf_id,
+                        nvcf_version=nvcf_version,
+                        nim_metadata=nim_metadata,
+                        monitor=monitor,
+                    )
+                else:
+                    raise ValueError(
+                        f"Model {model_name} 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."
+                    )
+            elif nvcf_id and nvcf_version:
+                self.models[model_name] = NimChatCompletion(
+                    model=model_name,
+                    nvcf_id=nvcf_id,
+                    nvcf_version=nvcf_version,
+                    nim_metadata=nim_metadata,
+                    monitor=monitor,
+                )
+            else:
+                raise ValueError(
+                    "You must provide either a valid 'name' or a custom 'name' along with both 'nvcf_id' and 'nvcf_version'."
+                )
+
+
+class NimChatCompletion(object):
+    def __init__(
+        self,
+        model: str = "meta/llama3-8b-instruct",
+        nvcf_id: str = "",
+        nvcf_version: str = "",
+        nim_metadata: NimMetadata = None,
+        monitor: bool = False,
+        **kwargs,
+    ):
+        if nim_metadata is None:
+            raise ValueError(
+                "NimMetadata object is required to initialize NimChatCompletion object."
+            )
+
+        self.model_name = model
+        self.nim_metadata = nim_metadata
+        self.monitor = monitor
+        all_nvcf_models = self.nim_metadata.get_nvcf_chat_completion_models()
+
+        if nvcf_id and nvcf_version:
+            matching_models = [
+                m
+                for m in all_nvcf_models
+                if m["function-id"] == nvcf_id and m["version-id"] == nvcf_version
+            ]
+            if matching_models:
+                self.model = matching_models[0]
+                self.function_id = self.model["function-id"]
+                self.version_id = self.model["version-id"]
+                self.model_name = self.model["name"]
+            else:
+                raise ValueError(
+                    f"Function {self.function_id} with version {self.version_id} not found on NVCF"
+                )
+        else:
+            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"Model {self.model_name} not found in available NVCF models"
+                )
+
+            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
+
+    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["prompt_tokens"] = None
+            stats["completion_tokens"] = None
+
+        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"],
+                ),
+            )
+            conn.commit()
+        finally:
+            conn.close()
+
+    def __call__(self, **kwargs):
+        if self.first_request:
+            from metaflow import current
+
+            self.file_name = get_storage_path(current.task_id)
+            self.first_request = False
+
+        # 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()

metaflow_extensions/outerbounds/plugins/nvcf/constants.py

@@ -0,0 +1,3 @@
+SUPPORTABLE_GPU_TYPES = ["L40", "L40S", "L40G", "H100", "NEBIUS_H100"]
+DEFAULT_GPU_TYPE = "H100"
+MAX_N_GPU_BY_TYPE = {"L40": 1, "L40S": 1, "L40G": 1, "H100": 4, "NEBIUS_H100": 8}

metaflow_extensions/outerbounds/plugins/nvcf/exceptions.py

@@ -0,0 +1,94 @@
+from metaflow.exception import MetaflowException
+from .constants import SUPPORTABLE_GPU_TYPES
+
+
+class NvcfJobFailedException(MetaflowException):
+    headline = "[@nvidia] error"
+
+    def __init__(self, msg):
+        super(NvcfJobFailedException, self).__init__(msg)
+
+
+class NvcfPollingConnectionError(MetaflowException):
+    headline = "[@nvidia] polling error."
+
+    def __init__(self, og_error_msg):
+        msg = (
+            "An error occurred while polling the job status. "
+            "\n\nOriginal error message: %s" % (og_error_msg)
+        )
+
+        super(NvcfPollingConnectionError, self).__init__(msg)
+
+
+class RequestedGPUTypeUnavailableException(MetaflowException):
+    headline = "[@nvidia RequestedGPUTypeUnavailableException] GPU type unavailable."
+
+    def __init__(self, requested_gpu_type):
+        msg = (
+            f"The requested GPU type @nvidia(..., gpu_type='{requested_gpu_type}') is not available. "
+            f"Please choose from the following supported GPU types when using @nvidia: {SUPPORTABLE_GPU_TYPES}"
+        )
+        super(RequestedGPUTypeUnavailableException, self).__init__(msg)
+
+
+class UnsupportedNvcfConfigurationException(MetaflowException):
+    headline = (
+        "[@nvidia UnsupportedNvcfConfigurationException] Unsupported GPU configuration"
+    )
+
+    def __init__(self, n_gpu, gpu_type, available_configurations, step):
+        msg = f"The requested configuration of @nvidia(gpu={n_gpu}, gpu_type='{gpu_type}') for @step {step} is not available."
+        if len(available_configurations) == 0:
+            msg += (
+                "\n\nNo configurations are available in your Outerbounds deployment."
+                " Please contact Outerbounds support if you wish to use @nvidia."
+            )
+        else:
+            msg += f"\n\nAvailable configurations for your deployment include: \n\t- {self._display(available_configurations)}"
+            msg += "\n\nPlease contact Outerbounds support if you wish to use a configuration not listed above."
+        super(UnsupportedNvcfConfigurationException, self).__init__(msg)
+
+    def _display(self, configs):
+        _available_decos = []
+        for cfg in configs:
+            n_gpu, gpu_type = cfg[0], cfg[1]
+            _available_decos.append(f"@nvidia(gpu={n_gpu}, gpu_type='{gpu_type}')")
+        return "\n\t- ".join(_available_decos)
+
+
+class UnsupportedNvcfDatastoreException(MetaflowException):
+    headline = "[@nvidia UnsupportedNvcfDatastoreException] Unsupported datastore"
+
+    def __init__(self, ds_type):
+        msg = (
+            "The *@nvidia* decorator requires --datastore=s3 or --datastore=azure or --datastore=gs at the moment."
+            f"Current datastore type: {ds_type}."
+        )
+        super(UnsupportedNvcfDatastoreException, self).__init__(msg)
+
+
+class NvcfTimeoutTooShortException(MetaflowException):
+    headline = "[@nvidia NvcfTimeoutTooShortException] Timeout too short"
+
+    def __init__(self, step):
+        msg = (
+            "The timeout for step *{step}* should be at least 60 seconds for "
+            "execution with @nvidia".format(step=step)
+        )
+        super(NvcfTimeoutTooShortException, self).__init__(msg)
+
+
+class NvcfQueueTimeoutTooShortException(MetaflowException):
+    headline = "[@nvidia NvcfQueueTimeoutTooShortException] Queue Timeout too short"
+
+    def __init__(self, step):
+        msg = (
+            "The queue timeout for step *{step}* should be at least 60 seconds for "
+            "execution with @nvidia".format(step=step)
+        )
+        super(NvcfQueueTimeoutTooShortException, self).__init__(msg)
+
+
+class NvcfKilledException(MetaflowException):
+    headline = "Nvidia job killed"