ob-metaflow-extensions 1.1.124__tar.gz → 1.1.125__tar.gz

{ob-metaflow-extensions-1.1.124 → ob-metaflow-extensions-1.1.125}/PKG-INFO

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

{ob-metaflow-extensions-1.1.124 → ob-metaflow-extensions-1.1.125}/metaflow_extensions/outerbounds/plugins/nvcf/exceptions.py

@@ -79,5 +79,16 @@ class NvcfTimeoutTooShortException(MetaflowException):
         super(NvcfTimeoutTooShortException, self).__init__(msg)
 
 
+class NvcfQueueTimeoutTooShortException(MetaflowException):
+    headline = "[@nvidia NvcfQueueTimeoutTooShortException] Queue Timeout too short"
+
+    def __init__(self, step):
+        msg = (
+            "The queue timeout for step *{step}* should be at least 60 seconds for "
+            "execution with @nvidia".format(step=step)
+        )
+        super(NvcfQueueTimeoutTooShortException, self).__init__(msg)
+
+
 class NvcfKilledException(MetaflowException):
     headline = "Nvidia job killed"

{ob-metaflow-extensions-1.1.124 → ob-metaflow-extensions-1.1.125}/metaflow_extensions/outerbounds/plugins/nvcf/nvcf.py

@@ -23,29 +23,53 @@ 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 retry_on_status(status_codes=[500], max_retries=3, delay=1):
     def decorator(func):
         @wraps(func)
-        def wrapper(*args, **kwargs):
+        def wrapper(instance, *args, **kwargs):
             retries = 0
-            while retries <= max_retries:
+
+            # Determine retry limit upfront
+            use_queue_timeout = 504 in status_codes
+            if use_queue_timeout:
+                poll_seconds = int(instance._poll_seconds)
+                retry_limit = (
+                    instance._queue_timeout + (poll_seconds - 1)
+                ) // poll_seconds
+                remainder = instance._queue_timeout % poll_seconds
+                last_timeout = remainder if remainder != 0 else poll_seconds
+            else:
+                retry_limit = max_retries
+
+            while retries < retry_limit:
                 try:
-                    return func(*args, **kwargs)
+                    return func(instance, *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})..."
-                        )
+                    if e.code not in status_codes or retries >= retry_limit:
+                        instance._status = JobStatus.FAILED
+                        if e.code == 504 and retries >= retry_limit:
+                            raise NvcfPollingConnectionError(
+                                "Request timed out after all retries"
+                            )
+                        raise
+
+                    if e.code == 504 and retries == retry_limit - 1:
+                        instance._poll_seconds = str(last_timeout)
+
+                    print(
+                        f"[@nvidia] {'Queue timeout' if e.code == 504 else f'Received {e.code}'}, "
+                        f"retrying ({retries + 1}/{retry_limit})... with poll seconds as {instance._poll_seconds}"
+                    )
+
+                    if e.code != 504:
                         time.sleep(delay)
-                        continue
-                    raise
-                except Exception:
+
+                    retries += 1
+                except URLError as e:
+                    instance._status = JobStatus.FAILED
                     raise
-            return func(*args, **kwargs)
+            # final attempt
+            return func(instance, *args, **kwargs)
 
         return wrapper
 
@@ -53,12 +77,15 @@ def retry_on_status(status_codes=RETRIABLE_STATUS_CODES, max_retries=3, delay=1)
 
 
 class Nvcf(object):
-    def __init__(self, metadata, datastore, environment, function_id, ngc_api_key):
+    def __init__(
+        self, metadata, datastore, environment, function_id, ngc_api_key, queue_timeout
+    ):
         self.metadata = metadata
         self.datastore = datastore
         self.environment = environment
         self._function_id = function_id
         self._ngc_api_key = ngc_api_key
+        self._queue_timeout = queue_timeout
 
     def launch_job(
         self,
@@ -120,6 +147,7 @@ class Nvcf(object):
             self.datastore._storage_impl,
             self._function_id,
             self._ngc_api_key,
+            self._queue_timeout,
         )
         self.job.submit()
 
@@ -172,7 +200,9 @@ result_endpoint = f"{nvcf_url}/v2/nvcf/pexec/status"
 
 
 class Job(object):
-    def __init__(self, command, env, task_spec, backend, function_id, ngc_api_key):
+    def __init__(
+        self, command, env, task_spec, backend, function_id, ngc_api_key, queue_timeout
+    ):
         self._payload = {
             "command": command,
             "env": {k: v for k, v in env.items() if v is not None},
@@ -180,6 +210,8 @@ class Job(object):
         self._result = {}
         self._function_id = function_id
         self._ngc_api_key = ngc_api_key
+        self._queue_timeout = queue_timeout
+        self._poll_seconds = "3600"
 
         flow_name = task_spec.get("flow_name")
         run_id = task_spec.get("run_id")
@@ -214,11 +246,14 @@ class Job(object):
         )
         self.heartbeat_thread.start()
 
+    @retry_on_status(status_codes=[504])
     def submit(self):
         try:
             headers = {
                 "Authorization": f"Bearer {self._ngc_api_key}",
                 "Content-Type": "application/json",
+                "nvcf-feature-enable-gateway-timeout": "true",
+                "NVCF-POLL-SECONDS": self._poll_seconds,
             }
             request_data = json.dumps(self._payload).encode()
             request = Request(
@@ -239,12 +274,9 @@ class Job(object):
                 self._status = JobStatus.SUBMITTED
             else:
                 self._status = JobStatus.FAILED
-            # TODO: Handle 404s nicely
-        except (HTTPError, URLError) as e:
-            # TODO: If queue is full, wait in line and retry?
-            # without that, branching over concurrent requests causes error.
-            self._state = JobStatus.FAILED
-            raise e
+        except URLError:
+            self._status = JobStatus.FAILED
+            raise
 
     @property
     def status(self):
@@ -272,12 +304,15 @@ class Job(object):
     def result(self):
         return self._result
 
-    @retry_on_status(status_codes=RETRIABLE_STATUS_CODES, max_retries=3, delay=5)
+    @retry_on_status(status_codes=[500], max_retries=3, delay=5)
+    @retry_on_status(status_codes=[504])
     def _poll(self):
         try:
             headers = {
                 "Authorization": f"Bearer {self._ngc_api_key}",
                 "Content-Type": "application/json",
+                "nvcf-feature-enable-gateway-timeout": "true",
+                "NVCF-POLL-SECONDS": self._poll_seconds,
             }
             request = Request(
                 f"{result_endpoint}/{self._invocation_id}", headers=headers
@@ -296,11 +331,6 @@ class Job(object):
                     f"[@nvidia] Unexpected response code: {response.getcode()}. Please notify an Outerbounds support engineer if this error persists."
                 )
                 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

{ob-metaflow-extensions-1.1.124 → ob-metaflow-extensions-1.1.125}/metaflow_extensions/outerbounds/plugins/nvcf/nvcf_cli.py

@@ -111,6 +111,9 @@ def kill(ctx, run_id):
 @click.argument("code-package-url")
 @click.option("--function-id", help="NVCF function id.")
 @click.option("--ngc-api-key", help="NGC API key.")
+@click.option(
+    "--queue-timeout", default=5 * 24 * 3600, help="Queue timeout in seconds."
+)
 @click.option("--run-id", help="Passed to the top-level 'step'.")
 @click.option("--task-id", help="Passed to the top-level 'step'.")
 @click.option("--input-paths", help="Passed to the top-level 'step'.")
@@ -133,6 +136,7 @@ def step(
     code_package_url,
     function_id,
     ngc_api_key,
+    queue_timeout,
     **kwargs,
 ):
     def echo(msg, stream="stderr", _id=None, **kwargs):
@@ -249,6 +253,7 @@ def step(
         ctx.obj.environment,
         function_id,
         ngc_api_key,
+        queue_timeout,
     )
     try:
         with ctx.obj.monitor.measure("metaflow.nvcf.launch_job"):

{ob-metaflow-extensions-1.1.124 → ob-metaflow-extensions-1.1.125}/metaflow_extensions/outerbounds/plugins/nvcf/nvcf_decorator.py

@@ -20,6 +20,7 @@ from .exceptions import (
     UnsupportedNvcfConfigurationException,
     UnsupportedNvcfDatastoreException,
     NvcfTimeoutTooShortException,
+    NvcfQueueTimeoutTooShortException,
 )
 
 from metaflow.metaflow_config import SERVICE_URL
@@ -36,10 +37,16 @@ class NvcfDecorator(StepDecorator):
         Number of GPUs to use.
     gpu_type : str
         Type of Nvidia GPU to use.
+    queue_timeout : int
+        Time to keep the job in NVCF's queue.
     """
 
     name = "nvidia"
-    defaults = {"gpu": 1, "gpu_type": None}
+    defaults = {
+        "gpu": 1,
+        "gpu_type": None,
+        "queue_timeout": 5 * 24 * 3600,  # Default 5 days in seconds
+    }
 
     package_url = None
     package_sha = None
@@ -118,6 +125,10 @@ class NvcfDecorator(StepDecorator):
             )
         self.attributes["function_id"] = available_configurations[desired_configuration]
 
+        queue_timeout = self.attributes["queue_timeout"]
+        if not isinstance(queue_timeout, int) or queue_timeout < 60:
+            raise NvcfQueueTimeoutTooShortException(step)
+
     def runtime_init(self, flow, graph, package, run_id):
         # Set some more internal state.
         self.flow = flow
@@ -144,6 +155,7 @@ class NvcfDecorator(StepDecorator):
             cli_options = {
                 "function_id": self.attributes["function_id"],
                 "ngc_api_key": self.attributes["ngc_api_key"],
+                "queue_timeout": self.attributes["queue_timeout"],
             }
             cli_args.command_options.update(cli_options)
             cli_args.entrypoint[0] = sys.executable

{ob-metaflow-extensions-1.1.124 → ob-metaflow-extensions-1.1.125}/ob_metaflow_extensions.egg-info/PKG-INFO

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

{ob-metaflow-extensions-1.1.124 → ob-metaflow-extensions-1.1.125}/setup.py

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