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

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

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

{ob-metaflow-extensions-1.1.142 → ob-metaflow-extensions-1.1.144}/metaflow_extensions/outerbounds/plugins/__init__.py

@@ -162,13 +162,18 @@ class ObpAuthProvider(object):
             client_params = {}
 
         from botocore.exceptions import ClientError
+        from botocore.config import Config
 
         with hide_access_keys():
             session = get_boto3_session(role_arn, session_vars)
+            _client_params = client_params.copy()
+            if _client_params.get("config") and type(_client_params["config"]) == dict:
+                _client_params["config"] = Config(**_client_params["config"])
+
             if with_error:
-                return session.client(module, **client_params), ClientError
+                return session.client(module, **_client_params), ClientError
             else:
-                return session.client(module, **client_params)
+                return session.client(module, **_client_params)
 
 
 AWS_CLIENT_PROVIDERS_DESC = [("obp", ".ObpAuthProvider")]

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

@@ -188,10 +188,12 @@ class Nvcf(object):
 
 
 class JobStatus(object):
-    SUBMITTED = "SUBMITTED"
-    RUNNING = "RUNNING"
-    SUCCESSFUL = "SUCCESSFUL"
-    FAILED = "FAILED"
+    CREATED = "CREATED"  # Job object created but not submitted
+    SUBMITTED = "SUBMITTED"  # Job submitted to NVCF
+    POLLED = "POLLED"  # Job has been successfully polled at least once
+    SUCCESSFUL = "SUCCESSFUL"  # Job completed successfully
+    FAILED = "FAILED"  # Job failed
+    DISAPPEARED = "DISAPPEARED"  # Job disappeared from NVCF but was previously polled (likely successful)
 
 
 nvcf_url = "https://api.nvcf.nvidia.com"
@@ -213,6 +215,11 @@ class Job(object):
         self._queue_timeout = queue_timeout
         self._poll_seconds = "3600"
 
+        # Initialize status and tracking variables
+        self._status = JobStatus.CREATED
+        self._last_poll_time = time.time()
+        self._force_poll_interval = 30
+
         flow_name = task_spec.get("flow_name")
         run_id = task_spec.get("run_id")
         step_name = task_spec.get("step_name")
@@ -280,10 +287,51 @@ class Job(object):
 
     @property
     def status(self):
-        if self._status not in [JobStatus.SUCCESSFUL, JobStatus.FAILED]:
+        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()
-            except (HTTPError, URLError) as e:
+                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 self._status
@@ -294,7 +342,8 @@ class Job(object):
 
     @property
     def is_running(self):
-        return self.status == JobStatus.SUBMITTED
+        # Job is running if it's in SUBMITTED or POLLED state
+        return self.status in [JobStatus.SUBMITTED, JobStatus.POLLED]
 
     @property
     def has_failed(self):
@@ -318,8 +367,10 @@ class Job(object):
                 f"{result_endpoint}/{self._invocation_id}", headers=headers
             )
             response = urlopen(request)
+            body = response.read()
+            print(f"[@nvidia] polling response: {body}")
             if response.getcode() == 200:
-                data = json.loads(response.read())
+                data = json.loads(body)
                 # TODO: Propagate the internal error forward
                 if data.get("exit_code") == 0:
                     self._status = JobStatus.SUCCESSFUL

{ob-metaflow-extensions-1.1.142 → ob-metaflow-extensions-1.1.144}/ob_metaflow_extensions.egg-info/PKG-INFO

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

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

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