ob-metaflow-extensions 1.1.71__tar.gz → 1.1.72__tar.gz

{ob-metaflow-extensions-1.1.71 → ob-metaflow-extensions-1.1.72}/PKG-INFO

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

{ob-metaflow-extensions-1.1.71 → ob-metaflow-extensions-1.1.72}/metaflow_extensions/outerbounds/plugins/nim/nim_manager.py

@@ -5,6 +5,8 @@ import requests
 from urllib.parse import urlparse
 from metaflow.metaflow_config import SERVICE_URL
 from metaflow.metaflow_config_funcs import init_config
+import sys
+import random
 
 NVCF_URL = "https://api.nvcf.nvidia.com"
 NVCF_SUBMIT_ENDPOINT = f"{NVCF_URL}/v2/nvcf/pexec/functions"
@@ -113,6 +115,9 @@ class 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")
+        )
 
         if self.compute_provider == "CoreWeave":
             cw_model_names = [
@@ -154,33 +159,64 @@ class NimChatCompletion(object):
             request_data = {"model": self.model, **kwargs}
             request_url = f"{NVCF_SUBMIT_ENDPOINT}/{self.function_id}"
 
-            response = requests.post(
-                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()
+            attempts = 0
+            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,
+                    )
+                    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()
+                except requests.exceptions.ConnectionError as e:
+                    # ConnectionErrors are generally temporary errors like DNS resolution failures,
+                    # timeouts etc.
+                    print(
+                        "Encountered connection error. Retrying...", 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}"
-                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."
-                    )
+                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 as e:
+                        # ConnectionErrors are generally temporary errors like DNS resolution failures,
+                        # timeouts etc.
+                        print(
+                            "Encountered connection error. Retrying...",
+                            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()

{ob-metaflow-extensions-1.1.71 → ob-metaflow-extensions-1.1.72}/ob_metaflow_extensions.egg-info/PKG-INFO

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

{ob-metaflow-extensions-1.1.71 → ob-metaflow-extensions-1.1.72}/setup.py

@@ -2,7 +2,7 @@ from setuptools import setup, find_namespace_packages
 from pathlib import Path
 
 
-version = "1.1.71"
+version = "1.1.72"
 this_directory = Path(__file__).parent
 long_description = (this_directory / "README.md").read_text()