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

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

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

{ob-metaflow-extensions-1.1.70 → 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.70 → ob-metaflow-extensions-1.1.72}/metaflow_extensions/outerbounds/plugins/perimeters.py

@@ -7,6 +7,9 @@ from typing import Union
 
 CURRENT_PERIMETER_KEY = "OB_CURRENT_PERIMETER"
 CURRENT_PERIMETER_URL = "OB_CURRENT_PERIMETER_MF_CONFIG_URL"
+CURRENT_PERIMETER_URL_LEGACY_KEY = (
+    "OB_CURRENT_PERIMETER_URL"  # For backwards compatibility with workstations.
+)
 
 
 def get_perimeter_config_url_if_set_in_ob_config() -> Union[str, None]:
@@ -37,10 +40,18 @@ def get_perimeter_config_url_if_set_in_ob_config() -> Union[str, None]:
         with open(file_path, "r") as f:
             ob_config = json.loads(f.read())
 
-        if CURRENT_PERIMETER_KEY in ob_config and CURRENT_PERIMETER_URL in ob_config:
+        if CURRENT_PERIMETER_KEY in ob_config and (
+            CURRENT_PERIMETER_URL in ob_config
+            or CURRENT_PERIMETER_URL_LEGACY_KEY in ob_config
+        ):
             os.environ[CURRENT_PERIMETER_KEY] = ob_config[CURRENT_PERIMETER_KEY]
-            os.environ[CURRENT_PERIMETER_URL] = ob_config[CURRENT_PERIMETER_URL]
-            return ob_config[CURRENT_PERIMETER_URL]
+            if CURRENT_PERIMETER_URL in ob_config:
+                os.environ[CURRENT_PERIMETER_URL] = ob_config[CURRENT_PERIMETER_URL]
+            elif CURRENT_PERIMETER_URL_LEGACY_KEY in ob_config:
+                os.environ[CURRENT_PERIMETER_URL] = ob_config[
+                    CURRENT_PERIMETER_URL_LEGACY_KEY
+                ]
+            return os.environ[CURRENT_PERIMETER_URL]
         else:
             raise MetaflowException(
                 "{} does not contain the key {}".format(

{ob-metaflow-extensions-1.1.70 → 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.70
+Version: 1.1.72
 Summary: Outerbounds Platform Extensions for Metaflow
 Author: Outerbounds, Inc.
 License: Commercial

{ob-metaflow-extensions-1.1.70 → 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.70"
+version = "1.1.72"
 this_directory = Path(__file__).parent
 long_description = (this_directory / "README.md").read_text()
 

{ob-metaflow-extensions-1.1.70 → ob-metaflow-extensions-1.1.72}/metaflow_extensions/outerbounds/remote_config.py

@@ -86,14 +86,14 @@ def init_config() -> Dict[str, str]:
                                  command with the string provided in the Outerbounds dashboard"
         )
 
-    # users still have a legacy format and that's ok.
-    if OBP_REMOTE_CONFIG_KEY not in remote_config:
-        return remote_config
-
     perimeter_config_url = get_perimeter_config_url_if_set_in_ob_config()
     if perimeter_config_url:
         remote_config[OBP_REMOTE_CONFIG_KEY] = perimeter_config_url
 
+    # users still have a legacy format and that's ok.
+    if OBP_REMOTE_CONFIG_KEY not in remote_config:
+        return remote_config
+
     metaflow_config = resolve_config_from_remote(
         remote_url=remote_config[OBP_REMOTE_CONFIG_KEY],
         auth_token=remote_config[AUTH_KEY],