ob-metaflow-extensions 1.1.122__py2.py3-none-any.whl → 1.1.123__py2.py3-none-any.whl

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

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.122.dist-info → ob_metaflow_extensions-1.1.123.dist-info}/METADATA

@@ -1,13 +1,13 @@
 Metadata-Version: 2.1
 Name: ob-metaflow-extensions
-Version: 1.1.122
+Version: 1.1.123
 Summary: Outerbounds Platform Extensions for Metaflow
 Author: Outerbounds, Inc.
 License: Commercial
 Description-Content-Type: text/markdown
 Requires-Dist: boto3
 Requires-Dist: kubernetes
-Requires-Dist: ob-metaflow (==2.13.7.1)
+Requires-Dist: ob-metaflow (==2.13.8.1)
 
 # Outerbounds platform package
 

{ob_metaflow_extensions-1.1.122.dist-info → ob_metaflow_extensions-1.1.123.dist-info}/RECORD

@@ -17,7 +17,7 @@ metaflow_extensions/outerbounds/plugins/nvcf/__init__.py,sha256=47DEQpj8HBSa-_TI
 metaflow_extensions/outerbounds/plugins/nvcf/constants.py,sha256=aGHdNw_hqBu8i0zWXcatQM6e769wUXox0l8g0f6fNZ8,146
 metaflow_extensions/outerbounds/plugins/nvcf/exceptions.py,sha256=Wn5WvE_sY-L2jEz-iObMLii5Ds_HQJuE437ufadPFLk,3258
 metaflow_extensions/outerbounds/plugins/nvcf/heartbeat_store.py,sha256=pOWwm8LFQBbtku0zNBBwCyXxLK8U-hhC4naQcmU69nE,6217
-metaflow_extensions/outerbounds/plugins/nvcf/nvcf.py,sha256=zOIDbN4PhRM2VMHczfoAHUeo1df2UrqWMgTwcppsTwc,8990
+metaflow_extensions/outerbounds/plugins/nvcf/nvcf.py,sha256=sfWaKZyKuM02v5DujPdfLbm-WoecxHfGn8g432Roct4,10273
 metaflow_extensions/outerbounds/plugins/nvcf/nvcf_cli.py,sha256=5pLEekiw3krlwpcjfjjfUL-URep6soZgmfTqtzLz4Vo,9362
 metaflow_extensions/outerbounds/plugins/nvcf/nvcf_decorator.py,sha256=yGv_6EmrBZNiQQP0rEWWE3akAL-KfI3Wd4ZFrcgl3VQ,8663
 metaflow_extensions/outerbounds/plugins/nvcf/utils.py,sha256=DxWSCayfa95e0HJkWacey1s1nxoTpaunGhrb_0Ayv28,133
@@ -33,7 +33,7 @@ metaflow_extensions/outerbounds/plugins/snowpark/snowpark_cli.py,sha256=ezJ2Jr8J
 metaflow_extensions/outerbounds/plugins/snowpark/snowpark_client.py,sha256=JEW0EUxj_mNZXo9OFkJFmWfg-P7_CEgvNbgsMTCBTAE,4273
 metaflow_extensions/outerbounds/plugins/snowpark/snowpark_decorator.py,sha256=a7LqSKULVh8IrR1StrVPbemHOLojR0nEqh-mMX-M1i4,9904
 metaflow_extensions/outerbounds/plugins/snowpark/snowpark_exceptions.py,sha256=FTfYlJu-sn9DkPOs2R1V1ChWb1vZthOgeq0BZdT1ucY,296
-metaflow_extensions/outerbounds/plugins/snowpark/snowpark_job.py,sha256=d_5UhXqZ_12rCvatH1capPQZYGLx1FVqq_rtW65OXyk,6874
+metaflow_extensions/outerbounds/plugins/snowpark/snowpark_job.py,sha256=aQphxX6jqYgfa83w387pEWl0keuLm38V53I8P8UL2ck,6887
 metaflow_extensions/outerbounds/plugins/snowpark/snowpark_service_spec.py,sha256=AI_kcm1hZV3JRxJkookcH6twiGnAYjk9Dx-MeoYz60Y,8511
 metaflow_extensions/outerbounds/plugins/tensorboard/__init__.py,sha256=9lUM4Cqi5RjrHBRfG6AQMRz8-R96eZC8Ih0KD2lv22Y,1858
 metaflow_extensions/outerbounds/profilers/__init__.py,sha256=wa_jhnCBr82TBxoS0e8b6_6sLyZX0fdHicuGJZNTqKw,29
@@ -44,7 +44,7 @@ metaflow_extensions/outerbounds/toplevel/plugins/azure/__init__.py,sha256=WUuhz2
 metaflow_extensions/outerbounds/toplevel/plugins/gcp/__init__.py,sha256=BbZiaH3uILlEZ6ntBLKeNyqn3If8nIXZFq_Apd7Dhco,70
 metaflow_extensions/outerbounds/toplevel/plugins/kubernetes/__init__.py,sha256=5zG8gShSj8m7rgF4xgWBZFuY3GDP5n1T0ktjRpGJLHA,69
 metaflow_extensions/outerbounds/toplevel/plugins/snowflake/__init__.py,sha256=LptpH-ziXHrednMYUjIaosS1SXD3sOtF_9_eRqd8SJw,50
-ob_metaflow_extensions-1.1.122.dist-info/METADATA,sha256=cSdliiwbdi-0Cr7GI4nuQfRQyGfmdrXuU3BPR3UJQf0,520
-ob_metaflow_extensions-1.1.122.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-ob_metaflow_extensions-1.1.122.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
-ob_metaflow_extensions-1.1.122.dist-info/RECORD,,
+ob_metaflow_extensions-1.1.123.dist-info/METADATA,sha256=JVisSFX7kwYfSyPmziBWAVYI87W1R7nmS5AJ2_AFtfQ,520
+ob_metaflow_extensions-1.1.123.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+ob_metaflow_extensions-1.1.123.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
+ob_metaflow_extensions-1.1.123.dist-info/RECORD,,