ob-metaflow-extensions 1.1.144__tar.gz → 1.1.145__tar.gz

{ob-metaflow-extensions-1.1.144 → ob-metaflow-extensions-1.1.145}/PKG-INFO

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

{ob-metaflow-extensions-1.1.144 → ob-metaflow-extensions-1.1.145}/metaflow_extensions/outerbounds/plugins/nvcf/nvcf.py

@@ -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"
@@ -213,12 +215,15 @@ class Job(object):
         self._function_id = function_id
         self._ngc_api_key = ngc_api_key
         self._queue_timeout = queue_timeout
-        self._poll_seconds = "3600"
+        self._poll_seconds = "300"
 
         # 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,61 +284,40 @@ 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
 
+    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()
+                # 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)
+
+        self._long_polling_active = False
+
     @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
-        ):
-            try:
-                self._poll()
-                self._last_poll_time = current_time
-
-                # 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 cached status - no need to poll
+        # Long polling loop will update the status
         return self._status
 
     @property
@@ -357,6 +341,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 +358,67 @@ 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)
+            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 → ob-metaflow-extensions-1.1.145}/ob_metaflow_extensions.egg-info/PKG-INFO

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

{ob-metaflow-extensions-1.1.144 → ob-metaflow-extensions-1.1.145}/setup.py

@@ -2,7 +2,7 @@ from setuptools import setup, find_namespace_packages
 from pathlib import Path
 
 
-version = "1.1.144"
+version = "1.1.145"
 this_directory = Path(__file__).parent
 long_description = (this_directory / "README.md").read_text()