ob-metaflow-extensions 1.1.130__py2.py3-none-any.whl → 1.5.1__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",
@@ -188,12 +188,16 @@ 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)
 
 
+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,6 +217,14 @@ 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()
+
+        # 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")
         step_name = task_spec.get("step_name")
@@ -272,21 +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):
-        if self._status not in [JobStatus.SUCCESSFUL, JobStatus.FAILED]:
+    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()
-            except (HTTPError, URLError) as e:
-                self._status = JobStatus.FAILED
-                raise NvcfPollingConnectionError(e)
-        return self._status
+                # 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 id(self):
@@ -294,11 +316,12 @@ 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):
-        return self.status == JobStatus.FAILED
+        return self._status == JobStatus.FAILED
 
     @property
     def result(self):
@@ -308,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",
@@ -317,20 +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 status code: {response.getcode()}")
+
             if response.getcode() == 200:
-                data = json.loads(response.read())
-                # TODO: Propagate the internal error forward
+                data = json.loads(body)
                 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)

metaflow_extensions/outerbounds/plugins/nvcf/nvcf_decorator.py

@@ -198,7 +198,12 @@ class NvcfDecorator(StepDecorator):
             meta["nvcf-nspectid"] = os.environ.get("NVCF_NSPECTID")
 
             entries = [
-                MetaDatum(field=k, value=v, type=k, tags=[])
+                MetaDatum(
+                    field=k,
+                    value=v,
+                    type=k,
+                    tags=["attempt_id:{0}".format(retry_count)],
+                )
                 for k, v in meta.items()
                 if v is not None
             ]

metaflow_extensions/outerbounds/plugins/nvct/exceptions.py

@@ -0,0 +1,71 @@
+from metaflow.exception import MetaflowException
+
+
+class NvctExecutionException(MetaflowException):
+    headline = "Nvct task couldn't be executed"
+
+
+class NvctTaskFailedException(MetaflowException):
+    headline = "Nvct task failed"
+
+
+class NvctKilledException(MetaflowException):
+    headline = "Nvct job killed"
+
+
+class RequestedGPUTypeUnavailableException(MetaflowException):
+    headline = "[@nvct RequestedGPUTypeUnavailableException] GPU type unavailable."
+
+    def __init__(self, requested_gpu_type, available_gpus):
+        msg = (
+            f"The requested GPU type @nvct(..., gpu_type='{requested_gpu_type}') is not available. "
+            f"Please choose from the following supported GPU types when using @nvct: {available_gpus}"
+        )
+        super(RequestedGPUTypeUnavailableException, self).__init__(msg)
+
+
+class UnsupportedNvctConfigurationException(MetaflowException):
+    headline = (
+        "[@nvct UnsupportedNvctConfigurationException] Unsupported GPU configuration"
+    )
+
+    def __init__(self, n_gpu, gpu_type, available_configurations, step):
+        msg = f"The requested configuration of @nvct(gpu={n_gpu}, gpu_type='{gpu_type}') for @step {step} is not available."
+        if len(available_configurations) == 0:
+            msg += (
+                "\n\nNo configurations are available in your Outerbounds deployment."
+                " Please contact Outerbounds support if you wish to use @nvct."
+            )
+        else:
+            msg += f"\n\nAvailable configurations for your deployment with {gpu_type} include: \n\t- {self._display(gpu_type, available_configurations)}"
+            msg += "\n\nPlease contact Outerbounds support if you wish to use a configuration not listed above."
+        super(UnsupportedNvctConfigurationException, self).__init__(msg)
+
+    def _display(self, gpu_type, configs):
+        _available_decos = []
+        for cfg in configs:
+            n_gpu = cfg["n_gpus"]
+            _available_decos.append(f"@nvct(gpu={n_gpu}, gpu_type='{gpu_type}')")
+        return "\n\t- ".join(_available_decos)
+
+
+class UnsupportedNvctDatastoreException(MetaflowException):
+    headline = "[@nvct UnsupportedNvctDatastoreException] Unsupported datastore"
+
+    def __init__(self, ds_type):
+        msg = (
+            "The *@nvct* decorator requires --datastore=s3 or --datastore=azure or --datastore=gs at the moment."
+            f"Current datastore type: {ds_type}."
+        )
+        super(UnsupportedNvctDatastoreException, self).__init__(msg)
+
+
+class NvctTimeoutTooShortException(MetaflowException):
+    headline = "[@nvct NvctTimeoutTooShortException] Timeout too short"
+
+    def __init__(self, step):
+        msg = (
+            "The timeout for step *{step}* should be at least 60 seconds for "
+            "execution with @nvct".format(step=step)
+        )
+        super(NvctTimeoutTooShortException, self).__init__(msg)

metaflow_extensions/outerbounds/plugins/nvct/nvct.py

@@ -0,0 +1,131 @@
+import requests
+from requests.adapters import HTTPAdapter, Retry
+
+BASE_URL = "https://api.ngc.nvidia.com/v2/orgs/zhxkmsaasxhw/"
+POLL_SEC = 1
+
+
+def _session(api_key):
+    s = requests.Session()
+    s.headers.update(
+        {
+            "Authorization": f"Bearer {api_key}",
+            "Content-Type": "application/json",
+        }
+    )
+    retry = Retry(total=5, backoff_factor=1.5, status_forcelist=[502, 503, 504])
+    s.mount("https://", HTTPAdapter(max_retries=retry))
+    return s
+
+
+def _url(path):
+    return BASE_URL.rstrip("/") + path
+
+
+class NVCTClient:
+    def __init__(self, api_key):
+        self.sess = _session(api_key)
+
+    # TODO: Handle https://outerboundsco.slack.com/archives/C05QGNR4E06/p1745970955540289
+    def create(self, spec):
+        r = self.sess.post(_url("/nvct/tasks"), json=spec, timeout=30)
+        r.raise_for_status()
+        return r.json().get("task", {}).get("id")
+
+    def get(self, task_id):
+        r = self.sess.get(_url(f"/nvct/tasks/{task_id}"), timeout=30)
+        r.raise_for_status()
+        return r.json().get("task", {})
+
+    def cancel(self, task_id):
+        r = self.sess.post(_url(f"/nvct/tasks/{task_id}/cancel"), timeout=30)
+        r.raise_for_status()
+
+
+class NVCTRequest(object):
+    def __init__(self, name):
+        self._spec = {}
+        self._spec["name"] = name
+        self._spec["gpuSpecification"] = {}
+        self._spec["resultHandlingStrategy"] = "NONE"
+        self._spec["terminationGracePeriodDuration"] = "PT10M"
+
+    def container_image(self, image):
+        self._spec["containerImage"] = image
+        return self
+
+    def container_args(self, args):
+        self._spec["containerArgs"] = args
+        return self
+
+    def env(self, key, value):
+        env_list = self._spec.setdefault("containerEnvironment", [])
+        env_list.append({"key": key, "value": value})
+        return self
+
+    def gpu(self, gpu, instance_type, backend):
+        gpu_spec = self._spec["gpuSpecification"]
+        gpu_spec["gpu"] = gpu
+        gpu_spec["instanceType"] = instance_type
+        gpu_spec["backend"] = backend
+        return self
+
+    def max_runtime(self, iso_duration):
+        self._spec["maxRuntimeDuration"] = iso_duration
+        return self
+
+    def max_queued(self, iso_duration="PT72H"):
+        self._spec["maxQueuedDuration"] = iso_duration
+        return self
+
+    def termination_grace(self, iso_duration="PT10M"):
+        self._spec["terminationGracePeriodDuration"] = iso_duration
+        return self
+
+    def extra(self, key, value):
+        self._spec[key] = value
+        return self
+
+    def to_dict(self):
+        return self._spec
+
+
+class NVCTTask:
+    def __init__(self, client: NVCTClient, spec):
+        self.client = client
+        self.spec = spec
+        self.id = None
+        self.record = None
+
+    def submit(self):
+        self.id = self.client.create(self.spec)
+        return self.id
+
+    def cancel(self):
+        if not self.has_finished:
+            self.client.cancel(self.id)
+
+    @property
+    def status(self):
+        self.record = self.client.get(self.id)
+        return self.record["status"]
+
+    @property
+    def is_waiting(self):
+        return self.status == "QUEUED"
+
+    @property
+    def is_running(self):
+        return self.status in {"RUNNING", "LAUNCHED"}
+
+    @property
+    def has_failed(self):
+        return self.status in {"ERRORED", "CANCELED"}
+
+    @property
+    def has_succeeded(self):
+        return self.status == "COMPLETED"
+
+    @property
+    def has_finished(self):
+        return self.has_succeeded or self.has_failed