ob-metaflow-extensions 1.1.71__py2.py3-none-any.whl → 1.1.73__py2.py3-none-any.whl

metaflow_extensions/outerbounds/plugins/auth_server.py

@@ -52,10 +52,17 @@ def get_token(url_path):
             return token_info
         except requests.exceptions.HTTPError as e:
             raise MetaflowException(repr(e))
-        except requests.exceptions.ConnectionError as e:
+        except (
+            requests.exceptions.ConnectionError,
+            requests.exceptions.ReadTimeout,
+        ) as e:
             # ConnectionErrors are generally temporary errors like DNS resolution failures,
             # timeouts etc.
-            print("received connection error. Retrying...", e, file=sys.stderr)
+            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

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,72 @@ 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,
+                    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}"
-                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,
+                        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()

{ob_metaflow_extensions-1.1.71.dist-info → ob_metaflow_extensions-1.1.73.dist-info}/METADATA

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

{ob_metaflow_extensions-1.1.71.dist-info → ob_metaflow_extensions-1.1.73.dist-info}/RECORD

@@ -2,12 +2,12 @@ metaflow_extensions/outerbounds/__init__.py,sha256=TRGvIUMjkfneWtYUFSWoubu_Kf2ek
 metaflow_extensions/outerbounds/remote_config.py,sha256=Zpfpjgz68_ZgxlXezjzlsDLo4840rkWuZgwDB_5H57U,4059
 metaflow_extensions/outerbounds/config/__init__.py,sha256=mYo95obHU1IE1wbPkeVz_pfTzNqlNabp1QBEMTGllbE,112
 metaflow_extensions/outerbounds/plugins/__init__.py,sha256=46NgbJBhVowDR6FyQrZPF2jHHqRTSyCBCYIQAyQ4Ryo,9516
-metaflow_extensions/outerbounds/plugins/auth_server.py,sha256=JhlMFcR7SPSfR1C9w6GlqJq-NYNhOfISmHl2PdkYUok,2212
+metaflow_extensions/outerbounds/plugins/auth_server.py,sha256=1v2GBqoMBxp5E7Lejz139w-jxJtPnLDvvHXP0HhEIHI,2361
 metaflow_extensions/outerbounds/plugins/perimeters.py,sha256=QXh3SFP7GQbS-RAIxUOPbhPzQ7KDFVxZkTdKqFKgXjI,2697
 metaflow_extensions/outerbounds/plugins/kubernetes/__init__.py,sha256=5zG8gShSj8m7rgF4xgWBZFuY3GDP5n1T0ktjRpGJLHA,69
 metaflow_extensions/outerbounds/plugins/kubernetes/kubernetes_client.py,sha256=gj6Iaz26bGbZm3aQuNS18Mqh_80iJp5PgFwFSlJRcn8,1968
 metaflow_extensions/outerbounds/plugins/nim/__init__.py,sha256=GVnvSTjqYVj5oG2yh8KJFt7iZ33cEadDD5HbdmC9hJ0,1457
-metaflow_extensions/outerbounds/plugins/nim/nim_manager.py,sha256=l8WDfVtsMt7aZaOaeIPT5ySidxfxXU8gmwLoKUP3f04,7044
+metaflow_extensions/outerbounds/plugins/nim/nim_manager.py,sha256=SWieODDxtIaeZwdMYtObDi57Kjyfw2DUuE6pJtU750w,9206
 metaflow_extensions/outerbounds/plugins/nvcf/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 metaflow_extensions/outerbounds/plugins/nvcf/nvcf.py,sha256=ftxC5SCo64P5Ycpv5vudluTnQi3-VCZW0umdsPP326A,7926
 metaflow_extensions/outerbounds/plugins/nvcf/nvcf_cli.py,sha256=ow3lonclEDoZEUQCDV_L8lEr6HopXqjNXzubRrfdIm4,7219
@@ -19,7 +19,7 @@ metaflow_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.py,
 metaflow_extensions/outerbounds/toplevel/plugins/azure/__init__.py,sha256=WUuhz2YQfI4fz7nIcipwwWq781eaoHEk7n4GAn1npDg,63
 metaflow_extensions/outerbounds/toplevel/plugins/gcp/__init__.py,sha256=BbZiaH3uILlEZ6ntBLKeNyqn3If8nIXZFq_Apd7Dhco,70
 metaflow_extensions/outerbounds/toplevel/plugins/kubernetes/__init__.py,sha256=5zG8gShSj8m7rgF4xgWBZFuY3GDP5n1T0ktjRpGJLHA,69
-ob_metaflow_extensions-1.1.71.dist-info/METADATA,sha256=GvyeQsYpp05xKshOovytjvOPBo1JbfmfyzhcXjBIAu0,519
-ob_metaflow_extensions-1.1.71.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-ob_metaflow_extensions-1.1.71.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
-ob_metaflow_extensions-1.1.71.dist-info/RECORD,,
+ob_metaflow_extensions-1.1.73.dist-info/METADATA,sha256=d45Pkzj0dMiznAUc91auqJVpgAp7o7QD_R5DHIKqwgg,519
+ob_metaflow_extensions-1.1.73.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+ob_metaflow_extensions-1.1.73.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
+ob_metaflow_extensions-1.1.73.dist-info/RECORD,,