ob-metaflow-extensions 1.1.114__py2.py3-none-any.whl → 1.1.116__py2.py3-none-any.whl

metaflow_extensions/outerbounds/plugins/nvcf/constants.py

@@ -0,0 +1,3 @@
+SUPPORTABLE_GPU_TYPES = ["L40", "L40S", "L40G", "H100"]
+DEFAULT_GPU_TYPE = "H100"
+MAX_N_GPU_BY_TYPE = {"L40": 1, "L40S": 1, "L40G": 1, "H100": 4}

metaflow_extensions/outerbounds/plugins/nvcf/exceptions.py

@@ -0,0 +1,83 @@
+from metaflow.exception import MetaflowException
+from .constants import SUPPORTABLE_GPU_TYPES
+
+
+class NvcfJobFailedException(MetaflowException):
+    headline = "[@nvidia] error"
+
+    def __init__(self, msg):
+        super(NvcfJobFailedException, self).__init__(msg)
+
+
+class NvcfPollingConnectionError(MetaflowException):
+    headline = "[@nvidia] polling error."
+
+    def __init__(self, og_error_msg):
+        msg = (
+            "An error occurred while polling the job status. "
+            "\n\nOriginal error message: %s" % (og_error_msg)
+        )
+
+        super(NvcfPollingConnectionError, self).__init__(msg)
+
+
+class RequestedGPUTypeUnavailableException(MetaflowException):
+    headline = "[@nvidia RequestedGPUTypeUnavailableException] GPU type unavailable."
+
+    def __init__(self, requested_gpu_type):
+        msg = (
+            f"The requested GPU type @nvidia(..., gpu_type='{requested_gpu_type}') is not available. "
+            f"Please choose from the following supported GPU types when using @nvidia: {SUPPORTABLE_GPU_TYPES}"
+        )
+        super(RequestedGPUTypeUnavailableException, self).__init__(msg)
+
+
+class UnsupportedNvcfConfigurationException(MetaflowException):
+    headline = (
+        "[@nvidia UnsupportedNvcfConfigurationException] Unsupported GPU configuration"
+    )
+
+    def __init__(self, n_gpu, gpu_type, available_configurations, step):
+        msg = f"The requested configuration of @nvidia(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 @nvidia."
+            )
+        else:
+            msg += f"\n\nAvailable configurations for your deployment include: \n\t- {self._display(available_configurations)}"
+            msg += "\n\nPlease contact Outerbounds support if you wish to use a configuration not listed above."
+        super(UnsupportedNvcfConfigurationException, self).__init__(msg)
+
+    def _display(self, configs):
+        _available_decos = []
+        for cfg in configs:
+            n_gpu, gpu_type = cfg[0], cfg[1]
+            _available_decos.append(f"@nvidia(gpu={n_gpu}, gpu_type='{gpu_type}')")
+        return "\n\t- ".join(_available_decos)
+
+
+class UnsupportedNvcfDatastoreException(MetaflowException):
+    headline = "[@nvidia UnsupportedNvcfDatastoreException] Unsupported datastore"
+
+    def __init__(self, ds_type):
+        msg = (
+            "The *@nvidia* decorator requires --datastore=s3 or --datastore=azure or --datastore=gs at the moment."
+            f"Current datastore type: {ds_type}."
+        )
+        super(UnsupportedNvcfDatastoreException, self).__init__(msg)
+
+
+class NvcfTimeoutTooShortException(MetaflowException):
+    headline = "[@nvidia NvcfTimeoutTooShortException] Timeout too short"
+
+    def __init__(self, step):
+        msg = (
+            "The timeout for step *{step}* should be at least 60 seconds for "
+            "execution with @nvidia".format(step=step)
+        )
+        super(NvcfTimeoutTooShortException, self).__init__(msg)
+
+
+class NvcfKilledException(MetaflowException):
+    headline = "Nvidia job killed"

metaflow_extensions/outerbounds/plugins/nvcf/nvcf.py

@@ -5,7 +5,6 @@ from urllib.parse import urlparse
 from urllib.request import HTTPError, Request, URLError, urlopen
 
 from metaflow import util
-from metaflow.exception import MetaflowException
 from metaflow.mflog import (
     BASH_SAVE_LOGS,
     bash_capture_logs,
@@ -14,16 +13,7 @@ from metaflow.mflog import (
     get_log_tailer,
 )
 import requests
-from metaflow.metaflow_config_funcs import init_config
-
-
-class NvcfException(MetaflowException):
-    headline = "Nvidia error"
-
-
-class NvcfKilledException(MetaflowException):
-    headline = "Nvidia job killed"
-
+from .exceptions import NvcfJobFailedException, NvcfPollingConnectionError
 
 # Redirect structured logs to $PWD/.logs/
 LOGS_DIR = "$PWD/.logs"
@@ -34,10 +24,12 @@ STDERR_PATH = os.path.join(LOGS_DIR, STDERR_FILE)
 
 
 class Nvcf(object):
-    def __init__(self, metadata, datastore, environment):
+    def __init__(self, metadata, datastore, environment, function_id, ngc_api_key):
         self.metadata = metadata
         self.datastore = datastore
         self.environment = environment
+        self._function_id = function_id
+        self._ngc_api_key = ngc_api_key
 
     def launch_job(
         self,
@@ -93,7 +85,12 @@ class Nvcf(object):
             % cmd_str
         )
         self.job = Job(
-            'bash -c "%s"' % cmd_str, env, task_spec, self.datastore._storage_impl
+            'bash -c "%s"' % cmd_str,
+            env,
+            task_spec,
+            self.datastore._storage_impl,
+            self._function_id,
+            self._ngc_api_key,
         )
         self.job.submit()
 
@@ -128,7 +125,9 @@ class Nvcf(object):
             _id=self.job.id,
         )
         if self.job.has_failed:
-            raise NvcfException("This could be a transient error. Use @retry to retry.")
+            raise NvcfJobFailedException(
+                "This could be a transient error. Use @retry to retry."
+            )
 
 
 class JobStatus(object):
@@ -144,37 +143,14 @@ result_endpoint = f"{nvcf_url}/v2/nvcf/pexec/status"
 
 
 class Job(object):
-    def __init__(self, command, env, task_spec, backend):
+    def __init__(self, command, env, task_spec, backend, function_id, ngc_api_key):
         self._payload = {
             "command": command,
             "env": {k: v for k, v in env.items() if v is not None},
         }
         self._result = {}
-
-        conf = init_config()
-        if "OBP_AUTH_SERVER" in conf:
-            auth_host = conf["OBP_AUTH_SERVER"]
-        else:
-            auth_host = "auth." + urlparse(SERVICE_URL).hostname.split(".", 1)[1]
-
-        # NOTE: reusing the same auth_host as the one used in NimMetadata,
-        # however, user should not need to use nim container to use @nvidia.
-        # May want to refactor this to a common endpoint.
-        nim_info_url = "https://" + auth_host + "/generate/nim"
-
-        if "METAFLOW_SERVICE_AUTH_KEY" in conf:
-            headers = {"x-api-key": conf["METAFLOW_SERVICE_AUTH_KEY"]}
-            res = requests.get(nim_info_url, headers=headers)
-        else:
-            headers = json.loads(os.environ.get("METAFLOW_SERVICE_HEADERS"))
-            res = requests.get(nim_info_url, headers=headers)
-
-        res.raise_for_status()
-        self._ngc_api_key = res.json()["nvcf"]["api_key"]
-
-        for f in res.json()["nvcf"]["functions"]:
-            if f["model_key"] == "metaflow_task_executor":
-                self._function_id = f["id"]
+        self._function_id = function_id
+        self._ngc_api_key = ngc_api_key
 
         flow_name = task_spec.get("flow_name")
         run_id = task_spec.get("run_id")
@@ -236,6 +212,8 @@ class Job(object):
                 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
 
@@ -281,5 +259,9 @@ class Job(object):
                 self._result = data
             elif response.getcode() in [400, 500]:
                 self._status = JobStatus.FAILED
+            elif response.getcode() != 202:
+                print(
+                    f"[@nvidia] Unexpected response code: {response.getcode()}. Please notify an Outerbounds support engineer if this error persists."
+                )
         except (HTTPError, URLError) as e:
-            print(f"Error occurred while polling for result: {e}")
+            raise NvcfPollingConnectionError(e)

metaflow_extensions/outerbounds/plugins/nvcf/nvcf_cli.py

@@ -31,7 +31,8 @@ from metaflow.metaflow_config import (
 )
 from metaflow.mflog import TASK_LOG_SOURCE
 
-from .nvcf import Nvcf, NvcfKilledException
+from .nvcf import Nvcf
+from .exceptions import NvcfKilledException
 
 
 @click.group()
@@ -108,6 +109,8 @@ def kill(ctx, run_id):
 @click.argument("step-name")
 @click.argument("code-package-sha")
 @click.argument("code-package-url")
+@click.option("--function-id", help="NVCF function id.")
+@click.option("--ngc-api-key", help="NGC API key.")
 @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'.")
@@ -123,7 +126,15 @@ def kill(ctx, run_id):
     "--max-user-code-retries", default=0, help="Passed to the top-level 'step'."
 )
 @click.pass_context
-def step(ctx, step_name, code_package_sha, code_package_url, **kwargs):
+def step(
+    ctx,
+    step_name,
+    code_package_sha,
+    code_package_url,
+    function_id,
+    ngc_api_key,
+    **kwargs,
+):
     def echo(msg, stream="stderr", _id=None, **kwargs):
         msg = util.to_unicode(msg)
         if _id:
@@ -232,7 +243,13 @@ def step(ctx, step_name, code_package_sha, code_package_url, **kwargs):
                 ),
             )
 
-    nvcf = Nvcf(ctx.obj.metadata, ctx.obj.flow_datastore, ctx.obj.environment)
+    nvcf = Nvcf(
+        ctx.obj.metadata,
+        ctx.obj.flow_datastore,
+        ctx.obj.environment,
+        function_id,
+        ngc_api_key,
+    )
     try:
         with ctx.obj.monitor.measure("metaflow.nvcf.launch_job"):
             nvcf.launch_job(
@@ -242,7 +259,6 @@ def step(ctx, step_name, code_package_sha, code_package_url, **kwargs):
                 code_package_sha,
                 code_package_url,
                 ctx.obj.flow_datastore.TYPE,
-                # function_id=function_id,
                 env=env,
             )
     except Exception as e:

metaflow_extensions/outerbounds/plugins/nvcf/nvcf_decorator.py

@@ -1,23 +1,45 @@
 import os
 import sys
+import json
+import requests
+from urllib.parse import urlparse
 
+from metaflow import current
 from metaflow.exception import MetaflowException
 from metaflow.decorators import StepDecorator
+from metaflow.plugins.parallel_decorator import ParallelDecorator
 from metaflow.metadata_provider.util import sync_local_metadata_to_datastore
 from metaflow.metaflow_config import DATASTORE_LOCAL_DIR
 from metaflow.sidecar import Sidecar
 from metaflow.plugins.timeout_decorator import get_run_time_limit_for_task
-from .nvcf import NvcfException
-
 from metaflow.metadata_provider import MetaDatum
+from metaflow.metaflow_config_funcs import init_config
+from .constants import SUPPORTABLE_GPU_TYPES, DEFAULT_GPU_TYPE
+from .exceptions import (
+    RequestedGPUTypeUnavailableException,
+    UnsupportedNvcfConfigurationException,
+    UnsupportedNvcfDatastoreException,
+    NvcfTimeoutTooShortException,
+)
+
+from metaflow.metaflow_config import SERVICE_URL
 
 
 class NvcfDecorator(StepDecorator):
+
+    """
+    Specifies that this step should execute on DGX cloud.
+
+    Parameters
+    ----------
+    gpu : int
+        Number of GPUs to use.
+    gpu_type : str
+        Type of Nvidia GPU to use.
+    """
+
     name = "nvidia"
-    # defaults = {"function_id": "9e5647f2-740f-4101-a129-1c961a075575"}
-    defaults = {}
-    # "0817006f-018b-4590-b2a5-6cf9d64d9d9a"}
-    #
+    defaults = {"gpu": 1, "gpu_type": None}
 
     package_url = None
     package_sha = None
@@ -28,13 +50,8 @@ class NvcfDecorator(StepDecorator):
     def step_init(self, flow, graph, step, decos, environment, flow_datastore, logger):
         # Executing NVCF functions requires a non-local datastore.
         if flow_datastore.TYPE not in ("s3", "azure", "gs"):
-            raise NvcfException(
-                "The *@nvidia* decorator requires --datastore=s3 or --datastore=azure or --datastore=gs at the moment."
-            )
-        # if self.attributes["function_id"] is None:
-        #     raise NvcfException(
-        #         "The *@nvidia* decorator requires a function_id. Please reach out to Outerbounds if you are unsure how to get access to one."
-        #     )
+            raise UnsupportedNvcfDatastoreException(flow_datastore.TYPE)
+
         # Set internal state.
         self.logger = logger
         self.environment = environment
@@ -46,10 +63,10 @@ class NvcfDecorator(StepDecorator):
                 "Step *{step}* is marked for execution both on Kubernetes and "
                 "Nvidia. Please use one or the other.".format(step=step)
             )
-        if any([deco.name == "parallel" for deco in decos]):
+        if any([isinstance(deco, ParallelDecorator) for deco in decos]):
             raise MetaflowException(
                 "Step *{step}* contains a @parallel decorator "
-                "with the @nvidia decorator. @parallel is not supported with @nvidia.".format(
+                "with the @nvidia decorator. @parallel decorators are not currently supported with @nvidia.".format(
                     step=step
                 )
             )
@@ -57,10 +74,49 @@ class NvcfDecorator(StepDecorator):
         # Set run time limit for the NVCF function.
         self.run_time_limit = get_run_time_limit_for_task(decos)
         if self.run_time_limit < 60:
-            raise NvcfException(
-                "The timeout for step *{step}* should be at least 60 seconds for "
-                "execution with @nvidia.".format(step=step)
+            raise NvcfTimeoutTooShortException(step)
+
+        conf = init_config()
+        if "OBP_AUTH_SERVER" in conf:
+            auth_host = conf["OBP_AUTH_SERVER"]
+        else:
+            auth_host = "auth." + urlparse(SERVICE_URL).hostname.split(".", 1)[1]
+
+        # NOTE: reusing the same auth_host as the one used in NimMetadata,
+        # however, user should not need to use nim container to use @nvidia.
+        # May want to refactor this to a common endpoint.
+        nim_info_url = "https://" + auth_host + "/generate/nim"
+
+        if "METAFLOW_SERVICE_AUTH_KEY" in conf:
+            headers = {"x-api-key": conf["METAFLOW_SERVICE_AUTH_KEY"]}
+            res = requests.get(nim_info_url, headers=headers)
+        else:
+            headers = json.loads(os.environ.get("METAFLOW_SERVICE_HEADERS"))
+            res = requests.get(nim_info_url, headers=headers)
+
+        res.raise_for_status()
+        self.attributes["ngc_api_key"] = res.json()["nvcf"]["api_key"]
+
+        available_functions_info = res.json()["nvcf"]["functions"]
+        requested_gpu_type = self.attributes["gpu_type"]
+        n_gpu = self.attributes["gpu"]
+
+        if requested_gpu_type is None:
+            requested_gpu_type = DEFAULT_GPU_TYPE
+        if requested_gpu_type not in SUPPORTABLE_GPU_TYPES:
+            raise RequestedGPUTypeUnavailableException(requested_gpu_type)
+
+        desired_configuration = (n_gpu, requested_gpu_type)
+        available_configurations = {}
+        for f in available_functions_info:
+            if f["model_key"] == "metaflow_task_executor":
+                available_configurations[(f["gpu"], f["gpu_type"])] = f["id"]
+
+        if desired_configuration not in available_configurations:
+            raise UnsupportedNvcfConfigurationException(
+                n_gpu, requested_gpu_type, available_configurations, step
             )
+        self.attributes["function_id"] = available_configurations[desired_configuration]
 
     def runtime_init(self, flow, graph, package, run_id):
         # Set some more internal state.
@@ -85,8 +141,11 @@ class NvcfDecorator(StepDecorator):
             cli_args.commands = ["nvidia", "step"]
             cli_args.command_args.append(self.package_sha)
             cli_args.command_args.append(self.package_url)
-            cli_args.command_options.update(self.attributes)
-            # cli_args.command_options["run-time-limit"] = self.run_time_limit
+            cli_options = {
+                "function_id": self.attributes["function_id"],
+                "ngc_api_key": self.attributes["ngc_api_key"],
+            }
+            cli_args.command_options.update(cli_options)
             cli_args.entrypoint[0] = sys.executable
 
     def task_pre_step(

metaflow_extensions/outerbounds/plugins/nvcf/utils.py

@@ -0,0 +1,6 @@
+import sys
+
+
+def warning_message(message, prefix="[@nvidia]"):
+    msg = "%s %s" % (prefix, message)
+    print(msg, file=sys.stderr)

{ob_metaflow_extensions-1.1.114.dist-info → ob_metaflow_extensions-1.1.116.dist-info}/METADATA

@@ -1,13 +1,13 @@
 Metadata-Version: 2.1
 Name: ob-metaflow-extensions
-Version: 1.1.114
+Version: 1.1.116
 Summary: Outerbounds Platform Extensions for Metaflow
 Author: Outerbounds, Inc.
 License: Commercial
 Description-Content-Type: text/markdown
 Requires-Dist: boto3
 Requires-Dist: kubernetes
-Requires-Dist: ob-metaflow (==2.12.39.1)
+Requires-Dist: ob-metaflow (==2.13.0.1)
 
 # Outerbounds platform package
 

{ob_metaflow_extensions-1.1.114.dist-info → ob_metaflow_extensions-1.1.116.dist-info}/RECORD

@@ -14,10 +14,13 @@ metaflow_extensions/outerbounds/plugins/kubernetes/kubernetes_client.py,sha256=f
 metaflow_extensions/outerbounds/plugins/nim/__init__.py,sha256=GVnvSTjqYVj5oG2yh8KJFt7iZ33cEadDD5HbdmC9hJ0,1457
 metaflow_extensions/outerbounds/plugins/nim/nim_manager.py,sha256=SWieODDxtIaeZwdMYtObDi57Kjyfw2DUuE6pJtU750w,9206
 metaflow_extensions/outerbounds/plugins/nvcf/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+metaflow_extensions/outerbounds/plugins/nvcf/constants.py,sha256=aGHdNw_hqBu8i0zWXcatQM6e769wUXox0l8g0f6fNZ8,146
+metaflow_extensions/outerbounds/plugins/nvcf/exceptions.py,sha256=Wn5WvE_sY-L2jEz-iObMLii5Ds_HQJuE437ufadPFLk,3258
 metaflow_extensions/outerbounds/plugins/nvcf/heartbeat_store.py,sha256=pOWwm8LFQBbtku0zNBBwCyXxLK8U-hhC4naQcmU69nE,6217
-metaflow_extensions/outerbounds/plugins/nvcf/nvcf.py,sha256=1ayR5YLiEy3XRGVEnRUpsDCH_UyK7z-7A8L5GxI4qOE,9535
-metaflow_extensions/outerbounds/plugins/nvcf/nvcf_cli.py,sha256=dX14geUdLcj2T7epv0jDcdLfxiZ_AgK2w29JDcCbx5E,9143
-metaflow_extensions/outerbounds/plugins/nvcf/nvcf_decorator.py,sha256=IsUTT_ijpmSw6-15ldK8ZVcFIJDp9aOqrR3NSF0IoXU,6543
+metaflow_extensions/outerbounds/plugins/nvcf/nvcf.py,sha256=zOIDbN4PhRM2VMHczfoAHUeo1df2UrqWMgTwcppsTwc,8990
+metaflow_extensions/outerbounds/plugins/nvcf/nvcf_cli.py,sha256=5pLEekiw3krlwpcjfjjfUL-URep6soZgmfTqtzLz4Vo,9362
+metaflow_extensions/outerbounds/plugins/nvcf/nvcf_decorator.py,sha256=yGv_6EmrBZNiQQP0rEWWE3akAL-KfI3Wd4ZFrcgl3VQ,8663
+metaflow_extensions/outerbounds/plugins/nvcf/utils.py,sha256=DxWSCayfa95e0HJkWacey1s1nxoTpaunGhrb_0Ayv28,133
 metaflow_extensions/outerbounds/plugins/profilers/deco_injector.py,sha256=oI_C3c64XBm7n88FILqHwn-Nnc5DeT_68I67lM9rXaI,2434
 metaflow_extensions/outerbounds/plugins/profilers/gpu_profile_decorator.py,sha256=gDHQ2sMIp4NuZSzUspbSd8RGdFAoO5mgZAyFcZ2a51Y,2619
 metaflow_extensions/outerbounds/plugins/secrets/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -41,7 +44,7 @@ metaflow_extensions/outerbounds/toplevel/plugins/azure/__init__.py,sha256=WUuhz2
 metaflow_extensions/outerbounds/toplevel/plugins/gcp/__init__.py,sha256=BbZiaH3uILlEZ6ntBLKeNyqn3If8nIXZFq_Apd7Dhco,70
 metaflow_extensions/outerbounds/toplevel/plugins/kubernetes/__init__.py,sha256=5zG8gShSj8m7rgF4xgWBZFuY3GDP5n1T0ktjRpGJLHA,69
 metaflow_extensions/outerbounds/toplevel/plugins/snowflake/__init__.py,sha256=LptpH-ziXHrednMYUjIaosS1SXD3sOtF_9_eRqd8SJw,50
-ob_metaflow_extensions-1.1.114.dist-info/METADATA,sha256=Z_l0voS_e6k7Z9jqOSBg6lTx2cxbsc2d-zSfoPdXC8A,521
-ob_metaflow_extensions-1.1.114.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-ob_metaflow_extensions-1.1.114.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
-ob_metaflow_extensions-1.1.114.dist-info/RECORD,,
+ob_metaflow_extensions-1.1.116.dist-info/METADATA,sha256=J5dqyrzD2HvLSgr8g8jgFCVnfdQ6K1p1SwdKIA8t0F8,520
+ob_metaflow_extensions-1.1.116.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+ob_metaflow_extensions-1.1.116.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
+ob_metaflow_extensions-1.1.116.dist-info/RECORD,,