ob-metaflow-extensions 1.1.121__tar.gz → 1.1.123__tar.gz

{ob-metaflow-extensions-1.1.121 → ob-metaflow-extensions-1.1.123}/PKG-INFO

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

{ob-metaflow-extensions-1.1.121 → ob-metaflow-extensions-1.1.123}/metaflow_extensions/outerbounds/plugins/nvcf/nvcf.py

@@ -1,8 +1,9 @@
 import json
 import os
+import time
 import threading
-from urllib.parse import urlparse
 from urllib.request import HTTPError, Request, URLError, urlopen
+from functools import wraps
 
 from metaflow import util
 from metaflow.mflog import (
@@ -12,7 +13,6 @@ from metaflow.mflog import (
     tail_logs,
     get_log_tailer,
 )
-import requests
 from .exceptions import NvcfJobFailedException, NvcfPollingConnectionError
 
 # Redirect structured logs to $PWD/.logs/
@@ -23,6 +23,35 @@ STDOUT_PATH = os.path.join(LOGS_DIR, STDOUT_FILE)
 STDERR_PATH = os.path.join(LOGS_DIR, STDERR_FILE)
 
 
+RETRIABLE_STATUS_CODES = [500]
+
+
+def retry_on_status(status_codes=RETRIABLE_STATUS_CODES, max_retries=3, delay=1):
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            retries = 0
+            while retries <= max_retries:
+                try:
+                    return func(*args, **kwargs)
+                except HTTPError as e:
+                    if e.code in status_codes and retries < max_retries:
+                        retries += 1
+                        print(
+                            f"[@nvidia] Received {e.code} error, retrying ({retries}/{max_retries})..."
+                        )
+                        time.sleep(delay)
+                        continue
+                    raise
+                except Exception:
+                    raise
+            return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
 class Nvcf(object):
     def __init__(self, metadata, datastore, environment, function_id, ngc_api_key):
         self.metadata = metadata
@@ -220,7 +249,11 @@ class Job(object):
     @property
     def status(self):
         if self._status not in [JobStatus.SUCCESSFUL, JobStatus.FAILED]:
-            self._poll()
+            try:
+                self._poll()
+            except (HTTPError, URLError) as e:
+                self._status = JobStatus.FAILED
+                raise NvcfPollingConnectionError(e)
         return self._status
 
     @property
@@ -239,6 +272,7 @@ class Job(object):
     def result(self):
         return self._result
 
+    @retry_on_status(status_codes=RETRIABLE_STATUS_CODES, max_retries=3, delay=5)
     def _poll(self):
         try:
             headers = {
@@ -257,11 +291,16 @@ class Job(object):
                 else:
                     self._status = JobStatus.FAILED
                 self._result = data
-            elif response.getcode() in [400, 500]:
-                self._status = JobStatus.FAILED
             elif response.getcode() != 202:
                 print(
                     f"[@nvidia] Unexpected response code: {response.getcode()}. Please notify an Outerbounds support engineer if this error persists."
                 )
-        except (HTTPError, URLError) as e:
-            raise NvcfPollingConnectionError(e)
+                self._status = JobStatus.FAILED
+        # 4xx and 5xx responses go in 'except' block
+        except HTTPError as e:
+            if e.code not in RETRIABLE_STATUS_CODES:
+                self._status = JobStatus.FAILED
+            raise
+        except URLError:
+            self._status = JobStatus.FAILED
+            raise

{ob-metaflow-extensions-1.1.121 → ob-metaflow-extensions-1.1.123}/metaflow_extensions/outerbounds/plugins/snowpark/snowpark_job.py

@@ -207,7 +207,7 @@ class RunningJob(object):
 
     @property
     def is_waiting(self):
-        return self.status == "PENDING"
+        return self.status in ["PENDING", "UNKNOWN"]
 
     @property
     def is_running(self):

{ob-metaflow-extensions-1.1.121 → ob-metaflow-extensions-1.1.123}/ob_metaflow_extensions.egg-info/PKG-INFO

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

ob-metaflow-extensions-1.1.123/ob_metaflow_extensions.egg-info/requires.txt

@@ -0,0 +1,3 @@
+boto3
+kubernetes
+ob-metaflow==2.13.8.1

{ob-metaflow-extensions-1.1.121 → ob-metaflow-extensions-1.1.123}/setup.py

@@ -2,7 +2,7 @@ from setuptools import setup, find_namespace_packages
 from pathlib import Path
 
 
-version = "1.1.121"
+version = "1.1.123"
 this_directory = Path(__file__).parent
 long_description = (this_directory / "README.md").read_text()
 
@@ -18,6 +18,6 @@ setup(
     install_requires=[
         "boto3",
         "kubernetes",
-        "ob-metaflow == 2.13.4.1",
+        "ob-metaflow == 2.13.8.1",
     ],
 )

ob-metaflow-extensions-1.1.121/ob_metaflow_extensions.egg-info/requires.txt

@@ -1,3 +0,0 @@
-boto3
-kubernetes
-ob-metaflow==2.13.4.1