ob-metaflow-extensions 1.1.144__py2.py3-none-any.whl → 1.1.146__py2.py3-none-any.whl

metaflow_extensions/outerbounds/plugins/nvcf/nvcf.py

@@ -153,7 +153,7 @@ class Nvcf(object):
 
     def wait(self, stdout_location, stderr_location, echo=None):
         def wait_for_launch(job):
-            status = job.status
+            status = job._status
             echo(
                 "Task status: %s..." % status,
                 "stderr",
@@ -196,6 +196,8 @@ class JobStatus(object):
     DISAPPEARED = "DISAPPEARED"  # Job disappeared from NVCF but was previously polled (likely successful)
 
 
+terminal_states = [JobStatus.SUCCESSFUL, JobStatus.FAILED, JobStatus.DISAPPEARED]
+
 nvcf_url = "https://api.nvcf.nvidia.com"
 submit_endpoint = f"{nvcf_url}/v2/nvcf/pexec/functions"
 result_endpoint = f"{nvcf_url}/v2/nvcf/pexec/status"
@@ -218,7 +220,10 @@ class Job(object):
         # Initialize status and tracking variables
         self._status = JobStatus.CREATED
         self._last_poll_time = time.time()
-        self._force_poll_interval = 30
+
+        # State tracking for long polling
+        self._long_polling_active = False
+        self._poll_response = None
 
         flow_name = task_spec.get("flow_name")
         run_id = task_spec.get("run_id")
@@ -279,62 +284,31 @@ class Job(object):
                 self._result = data
             elif response.getcode() == 202:
                 self._status = JobStatus.SUBMITTED
+                # Start long polling immediately after receiving 202
+                self._start_long_polling()
             else:
                 self._status = JobStatus.FAILED
         except URLError:
             self._status = JobStatus.FAILED
             raise
 
-    @property
-    def status(self):
-        terminal_states = [
-            JobStatus.SUCCESSFUL,
-            JobStatus.FAILED,
-            JobStatus.DISAPPEARED,
-        ]
-
-        # If status is already terminal, don't poll again
-        if self._status in terminal_states:
-            return self._status
-
-        current_time = time.time()
-        if (
-            current_time - self._last_poll_time > self._force_poll_interval
-            or self._status not in terminal_states
-        ):
+    def _start_long_polling(self):
+        if not self._long_polling_active:
+            self._long_polling_active = True
+            polling_thread = threading.Thread(target=self._long_poll_loop, daemon=True)
+            polling_thread.start()
+
+    def _long_poll_loop(self):
+        while self._long_polling_active and self._status not in terminal_states:
             try:
                 self._poll()
-                self._last_poll_time = current_time
+                # No sleep needed - the request itself will block for up to self._poll_seconds
+            except Exception as e:
+                print(f"[@nvidia] Long polling error: {e}")
+                # Brief pause before retry on error
+                time.sleep(1)
 
-                # Update job status to POLLED if this is our first successful poll
-                if self._status == JobStatus.SUBMITTED:
-                    self._status = JobStatus.POLLED
-
-                if self._status == JobStatus.SUCCESSFUL:
-                    return self._status
-
-            except HTTPError as e:
-                if e.code == 404:
-                    # 404 interpretation depends on job lifecycle
-                    if self._status in [JobStatus.POLLED, JobStatus.SUBMITTED]:
-                        # We've submitted or successfully polled this job before,
-                        # so a 404 likely means it completed and was removed
-                        self._status = JobStatus.DISAPPEARED
-                        self._result = {"exit_code": 0}
-                    else:
-                        # Job was never successfully tracked
-                        print(
-                            f"[@nvidia] 404 received for job that was never successfully tracked - treating as failure"
-                        )
-                        self._status = JobStatus.FAILED
-                        raise NvcfPollingConnectionError(e)
-                else:
-                    self._status = JobStatus.FAILED
-                    raise NvcfPollingConnectionError(e)
-            except URLError as e:
-                self._status = JobStatus.FAILED
-                raise NvcfPollingConnectionError(e)
-        return self._status
+        self._long_polling_active = False
 
     @property
     def id(self):
@@ -343,11 +317,11 @@ class Job(object):
     @property
     def is_running(self):
         # Job is running if it's in SUBMITTED or POLLED state
-        return self.status in [JobStatus.SUBMITTED, JobStatus.POLLED]
+        return self._status in [JobStatus.SUBMITTED, JobStatus.POLLED]
 
     @property
     def has_failed(self):
-        return self.status == JobStatus.FAILED
+        return self._status == JobStatus.FAILED
 
     @property
     def result(self):
@@ -357,6 +331,14 @@ class Job(object):
     @retry_on_status(status_codes=[504])
     def _poll(self):
         try:
+            # Implement rate limiting to prevent more than 1 request per second
+            current_time = time.time()
+            if (
+                hasattr(self, "_last_poll_time")
+                and current_time - self._last_poll_time < 1
+            ):
+                time.sleep(1 - (current_time - self._last_poll_time))
+
             headers = {
                 "Authorization": f"Bearer {self._ngc_api_key}",
                 "Content-Type": "application/json",
@@ -366,22 +348,70 @@ class Job(object):
             request = Request(
                 f"{result_endpoint}/{self._invocation_id}", headers=headers
             )
+
+            # Record time before making the request
+            self._last_poll_time = time.time()
+
             response = urlopen(request)
             body = response.read()
-            print(f"[@nvidia] polling response: {body}")
+            print(f"[@nvidia] polling status code: {response.getcode()}")
+
             if response.getcode() == 200:
                 data = json.loads(body)
-                # TODO: Propagate the internal error forward
                 if data.get("exit_code") == 0:
                     self._status = JobStatus.SUCCESSFUL
                 else:
                     self._status = JobStatus.FAILED
                 self._result = data
-            elif response.getcode() != 202:
+                self._long_polling_active = False  # Stop polling once job completes
+            elif response.getcode() == 202:
+                # Job is still running - status remains SUBMITTED or POLLED
+                if self._status == JobStatus.SUBMITTED:
+                    self._status = JobStatus.POLLED
+            elif response.getcode() == 302:
+                # Handle redirects for large responses or requests in different regions
+                redirect_location = response.headers.get("Location")
+                if redirect_location:
+                    redirect_request = Request(redirect_location, headers=headers)
+                    redirect_response = urlopen(redirect_request)
+                    if redirect_response.getcode() == 200:
+                        data = json.loads(redirect_response.read())
+                        if data.get("exit_code") == 0:
+                            self._status = JobStatus.SUCCESSFUL
+                        else:
+                            self._status = JobStatus.FAILED
+                        self._result = data
+                        self._long_polling_active = False
+            else:
                 print(
                     f"[@nvidia] Unexpected response code: {response.getcode()}. Please notify an Outerbounds support engineer if this error persists."
                 )
                 self._status = JobStatus.FAILED
-        except URLError:
+
+        except HTTPError as e:
+            if e.code == 404:
+                # 404 interpretation depends on job lifecycle
+                if self._status in [JobStatus.POLLED, JobStatus.SUBMITTED]:
+                    # We've submitted or successfully polled this job before,
+                    # so a 404 likely means it completed and was removed
+                    self._status = JobStatus.DISAPPEARED
+                    self._result = {"exit_code": 0}
+                    print(
+                        f"[@nvidia] 404 received for job that was previously tracked - assuming job completed"
+                    )
+                else:
+                    # Job was never successfully tracked
+                    print(
+                        f"[@nvidia] 404 received for job that was never successfully tracked - treating as failure"
+                    )
+                    self._status = JobStatus.FAILED
+                    raise NvcfPollingConnectionError(e)
+            elif e.code in [500, 504]:
+                # Don't set status to FAILED, just re-raise for retry decorator
+                raise
+            else:
+                self._status = JobStatus.FAILED
+                raise NvcfPollingConnectionError(e)
+        except URLError as e:
             self._status = JobStatus.FAILED
-            raise
+            raise NvcfPollingConnectionError(e)

{ob_metaflow_extensions-1.1.144.dist-info → ob_metaflow_extensions-1.1.146.dist-info}/METADATA

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

{ob_metaflow_extensions-1.1.144.dist-info → ob_metaflow_extensions-1.1.146.dist-info}/RECORD

@@ -28,7 +28,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=-Pm9cOWUzpv94TvVUeq-FenAWdfLBJd5N7WPqIGZVqU,3671
 metaflow_extensions/outerbounds/plugins/nvcf/heartbeat_store.py,sha256=pOWwm8LFQBbtku0zNBBwCyXxLK8U-hhC4naQcmU69nE,6217
-metaflow_extensions/outerbounds/plugins/nvcf/nvcf.py,sha256=pDnd30uGyBLgIMROI0ClN0wAaxvt2XrbBbc0QIZ9LmE,13872
+metaflow_extensions/outerbounds/plugins/nvcf/nvcf.py,sha256=3ZFdYItVpFWnHMOeyV1nslUyelfvX5rknh2d2IWxVws,15591
 metaflow_extensions/outerbounds/plugins/nvcf/nvcf_cli.py,sha256=3D-r5XO88Yh2k1EAZFJTe_PwdbhWp5qXflG8AgE4ZIU,9500
 metaflow_extensions/outerbounds/plugins/nvcf/nvcf_decorator.py,sha256=pxxNxW4bW3kbB6ybRam657GyKIhvIkMuidB94iFfCD8,9116
 metaflow_extensions/outerbounds/plugins/nvcf/utils.py,sha256=DxWSCayfa95e0HJkWacey1s1nxoTpaunGhrb_0Ayv28,133
@@ -58,7 +58,7 @@ metaflow_extensions/outerbounds/toplevel/plugins/gcp/__init__.py,sha256=BbZiaH3u
 metaflow_extensions/outerbounds/toplevel/plugins/kubernetes/__init__.py,sha256=5zG8gShSj8m7rgF4xgWBZFuY3GDP5n1T0ktjRpGJLHA,69
 metaflow_extensions/outerbounds/toplevel/plugins/ollama/__init__.py,sha256=GRSz2zwqkvlmFS6bcfYD_CX6CMko9DHQokMaH1iBshA,47
 metaflow_extensions/outerbounds/toplevel/plugins/snowflake/__init__.py,sha256=LptpH-ziXHrednMYUjIaosS1SXD3sOtF_9_eRqd8SJw,50
-ob_metaflow_extensions-1.1.144.dist-info/METADATA,sha256=zphrWhng37KADGpWv9zJHp4jZBfivvktH_a07HwU-i0,520
-ob_metaflow_extensions-1.1.144.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-ob_metaflow_extensions-1.1.144.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
-ob_metaflow_extensions-1.1.144.dist-info/RECORD,,
+ob_metaflow_extensions-1.1.146.dist-info/METADATA,sha256=Vtp-8yJaPeg8Hh33WXQYAerLoLlx5Ldya5M2U1lp5z0,520
+ob_metaflow_extensions-1.1.146.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+ob_metaflow_extensions-1.1.146.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
+ob_metaflow_extensions-1.1.146.dist-info/RECORD,,