ob-metaflow-extensions 1.1.45rc3__py2.py3-none-any.whl → 1.5.1__py2.py3-none-any.whl

metaflow_extensions/outerbounds/plugins/nvcf/nvcf_cli.py

@@ -0,0 +1,280 @@
+import json
+import os
+import sys
+import time
+import traceback
+
+from metaflow import util, Run
+from metaflow._vendor import click
+from metaflow.exception import METAFLOW_EXIT_DISALLOW_RETRY
+from metaflow.metadata_provider.util import sync_local_metadata_from_datastore
+from metaflow.metaflow_config import (
+    CARD_S3ROOT,
+    DATASTORE_LOCAL_DIR,
+    DATASTORE_SYSROOT_S3,
+    DATATOOLS_S3ROOT,
+    DEFAULT_METADATA,
+    SERVICE_HEADERS,
+    SERVICE_URL,
+    DEFAULT_SECRETS_BACKEND_TYPE,
+    DEFAULT_AWS_CLIENT_PROVIDER,
+    AWS_SECRETS_MANAGER_DEFAULT_REGION,
+    S3_ENDPOINT_URL,
+    AZURE_STORAGE_BLOB_SERVICE_ENDPOINT,
+    DATASTORE_SYSROOT_AZURE,
+    CARD_AZUREROOT,
+    DATASTORE_SYSROOT_GS,
+    CARD_GSROOT,
+    KUBERNETES_SANDBOX_INIT_SCRIPT,
+    OTEL_ENDPOINT,
+    NVIDIA_HEARTBEAT_THRESHOLD,
+)
+from metaflow.mflog import TASK_LOG_SOURCE
+
+from .nvcf import Nvcf
+from .exceptions import NvcfKilledException
+
+
+@click.group()
+def cli():
+    pass
+
+
+@cli.group(help="Commands related to nvidia.")
+def nvidia():
+    pass
+
+
+@nvidia.command(help="List steps / tasks running as an nvidia job.")
+@click.option(
+    "--run-id",
+    default=None,
+    required=True,
+    help="List unfinished tasks corresponding to the run id.",
+)
+@click.pass_context
+def list(ctx, run_id):
+    flow_name = ctx.obj.flow.name
+    run_obj = Run(pathspec=f"{flow_name}/{run_id}", _namespace_check=False)
+    running_invocations = []
+
+    for each_step in run_obj:
+        for each_task in each_step:
+            if not each_task.finished and "nvcf-function-id" in each_task.metadata_dict:
+
+                task_pathspec = each_task.pathspec
+                attempt = each_task.metadata_dict.get("attempt")
+                flow_name, run_id, step_name, task_id = task_pathspec.split("/")
+                running_invocations.append(
+                    f"Flow Name: {flow_name}, Run ID: {run_id}, Step Name: {step_name}, Task ID: {task_id}, Retry Count: {attempt}"
+                )
+
+    if running_invocations:
+        for each_invocation in running_invocations:
+            ctx.obj.echo(each_invocation)
+
+
+@nvidia.command(help="Kill steps / tasks running as an nvidia job.")
+@click.option(
+    "--run-id",
+    default=None,
+    required=True,
+    help="Terminate unfinished tasks corresponding to the run id.",
+)
+@click.pass_context
+def kill(ctx, run_id):
+    from metaflow_extensions.outerbounds.plugins.nvcf.heartbeat_store import (
+        HeartbeatStore,
+    )
+
+    datastore_root = ctx.obj.datastore_impl.datastore_root
+    store = HeartbeatStore(
+        main_pid=None,
+        storage_backend=ctx.obj.datastore_impl(datastore_root),
+    )
+
+    flow_name = ctx.obj.flow.name
+    tombstone_prefix = f"{flow_name}/{run_id}"
+    store.emit_tombstone(
+        tombstone_prefix=tombstone_prefix, folder_name="nvcf_heartbeats"
+    )
+
+
+@nvidia.command(
+    help="Execute a single task using @nvidia. This command calls the "
+    "top-level step command inside an nvidia job with the given options. "
+    "Typically you do not call this command directly; it is used internally by "
+    "Metaflow."
+)
+@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(
+    "--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'.")
+@click.option("--split-index", help="Passed to the top-level 'step'.")
+@click.option("--clone-path", help="Passed to the top-level 'step'.")
+@click.option("--clone-run-id", help="Passed to the top-level 'step'.")
+@click.option(
+    "--tag", multiple=True, default=None, help="Passed to the top-level 'step'."
+)
+@click.option("--namespace", default=None, help="Passed to the top-level 'step'.")
+@click.option("--retry-count", default=0, help="Passed to the top-level 'step'.")
+@click.option(
+    "--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,
+    function_id,
+    ngc_api_key,
+    queue_timeout,
+    **kwargs,
+):
+    def echo(msg, stream="stderr", _id=None, **kwargs):
+        msg = util.to_unicode(msg)
+        if _id:
+            msg = "[%s] %s" % (_id, msg)
+        ctx.obj.echo_always(msg, err=(stream == sys.stderr), **kwargs)
+
+    executable = ctx.obj.environment.executable(step_name)
+    entrypoint = "%s -u %s" % (executable, os.path.basename(sys.argv[0]))
+
+    top_args = " ".join(util.dict_to_cli_options(ctx.parent.parent.params))
+
+    input_paths = kwargs.get("input_paths")
+    split_vars = None
+    if input_paths:
+        max_size = 30 * 1024
+        split_vars = {
+            "METAFLOW_INPUT_PATHS_%d" % (i // max_size): input_paths[i : i + max_size]
+            for i in range(0, len(input_paths), max_size)
+        }
+        kwargs["input_paths"] = "".join("${%s}" % s for s in split_vars.keys())
+
+    step_args = " ".join(util.dict_to_cli_options(kwargs))
+    step_cli = "{entrypoint} {top_args} step {step} {step_args}".format(
+        entrypoint=entrypoint,
+        top_args=top_args,
+        step=step_name,
+        step_args=step_args,
+    )
+    node = ctx.obj.graph[step_name]
+
+    # Get retry information
+    retry_count = kwargs.get("retry_count", 0)
+    retry_deco = [deco for deco in node.decorators if deco.name == "retry"]
+    minutes_between_retries = None
+    if retry_deco:
+        minutes_between_retries = int(
+            retry_deco[0].attributes.get("minutes_between_retries", 1)
+        )
+
+    task_spec = {
+        "flow_name": ctx.obj.flow.name,
+        "step_name": step_name,
+        "run_id": kwargs["run_id"],
+        "task_id": kwargs["task_id"],
+        "retry_count": str(retry_count),
+    }
+
+    env = {
+        "METAFLOW_CODE_SHA": code_package_sha,
+        "METAFLOW_CODE_URL": code_package_url,
+        "METAFLOW_CODE_DS": ctx.obj.flow_datastore.TYPE,
+        "METAFLOW_SERVICE_URL": SERVICE_URL,
+        "METAFLOW_SERVICE_HEADERS": json.dumps(SERVICE_HEADERS),
+        "METAFLOW_DATASTORE_SYSROOT_S3": DATASTORE_SYSROOT_S3,
+        "METAFLOW_DATATOOLS_S3ROOT": DATATOOLS_S3ROOT,
+        "METAFLOW_DEFAULT_DATASTORE": ctx.obj.flow_datastore.TYPE,
+        "METAFLOW_USER": util.get_username(),
+        "METAFLOW_DEFAULT_METADATA": DEFAULT_METADATA,
+        "METAFLOW_CARD_S3ROOT": CARD_S3ROOT,
+        "METAFLOW_RUNTIME_ENVIRONMENT": "nvcf",
+        "METAFLOW_DEFAULT_SECRETS_BACKEND_TYPE": DEFAULT_SECRETS_BACKEND_TYPE,
+        "METAFLOW_DEFAULT_AWS_CLIENT_PROVIDER": DEFAULT_AWS_CLIENT_PROVIDER,
+        "METAFLOW_AWS_SECRETS_MANAGER_DEFAULT_REGION": AWS_SECRETS_MANAGER_DEFAULT_REGION,
+        "METAFLOW_S3_ENDPOINT_URL": S3_ENDPOINT_URL,
+        "METAFLOW_AZURE_STORAGE_BLOB_SERVICE_ENDPOINT": AZURE_STORAGE_BLOB_SERVICE_ENDPOINT,
+        "METAFLOW_DATASTORE_SYSROOT_AZURE": DATASTORE_SYSROOT_AZURE,
+        "METAFLOW_CARD_AZUREROOT": CARD_AZUREROOT,
+        "METAFLOW_DATASTORE_SYSROOT_GS": DATASTORE_SYSROOT_GS,
+        "METAFLOW_CARD_GSROOT": CARD_GSROOT,
+        "METAFLOW_INIT_SCRIPT": KUBERNETES_SANDBOX_INIT_SCRIPT,
+        "METAFLOW_OTEL_ENDPOINT": OTEL_ENDPOINT,
+        "METAFLOW_NVIDIA_HEARTBEAT_THRESHOLD": str(NVIDIA_HEARTBEAT_THRESHOLD),
+    }
+
+    env_deco = [deco for deco in node.decorators if deco.name == "environment"]
+    if env_deco:
+        env.update(env_deco[0].attributes["vars"])
+
+    # Add the environment variables related to the input-paths argument
+    if split_vars:
+        env.update(split_vars)
+
+    if retry_count:
+        ctx.obj.echo_always(
+            "Sleeping %d minutes before the next retry" % minutes_between_retries
+        )
+        time.sleep(minutes_between_retries * 60)
+
+    # this information is needed for log tailing
+    ds = ctx.obj.flow_datastore.get_task_datastore(
+        mode="w",
+        run_id=kwargs["run_id"],
+        step_name=step_name,
+        task_id=kwargs["task_id"],
+        attempt=int(retry_count),
+    )
+    stdout_location = ds.get_log_location(TASK_LOG_SOURCE, "stdout")
+    stderr_location = ds.get_log_location(TASK_LOG_SOURCE, "stderr")
+
+    def _sync_metadata():
+        if ctx.obj.metadata.TYPE == "local":
+            sync_local_metadata_from_datastore(
+                DATASTORE_LOCAL_DIR,
+                ctx.obj.flow_datastore.get_task_datastore(
+                    kwargs["run_id"], step_name, kwargs["task_id"]
+                ),
+            )
+
+    nvcf = Nvcf(
+        ctx.obj.metadata,
+        ctx.obj.flow_datastore,
+        ctx.obj.environment,
+        function_id,
+        ngc_api_key,
+        queue_timeout,
+    )
+    try:
+        with ctx.obj.monitor.measure("metaflow.nvcf.launch_job"):
+            nvcf.launch_job(
+                step_name,
+                step_cli,
+                task_spec,
+                code_package_sha,
+                code_package_url,
+                ctx.obj.flow_datastore.TYPE,
+                env=env,
+            )
+    except Exception as e:
+        traceback.print_exc()
+        _sync_metadata()
+        sys.exit(METAFLOW_EXIT_DISALLOW_RETRY)
+    try:
+        nvcf.wait(stdout_location, stderr_location, echo=echo)
+    except NvcfKilledException:
+        # don't retry killed tasks
+        traceback.print_exc()
+        sys.exit(METAFLOW_EXIT_DISALLOW_RETRY)
+    finally:
+        _sync_metadata()

metaflow_extensions/outerbounds/plugins/nvcf/nvcf_decorator.py

@@ -0,0 +1,242 @@
+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 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,
+    NvcfQueueTimeoutTooShortException,
+)
+
+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.
+    queue_timeout : int
+        Time to keep the job in NVCF's queue.
+    """
+
+    name = "nvidia"
+    defaults = {
+        "gpu": 1,
+        "gpu_type": None,
+        "queue_timeout": 5 * 24 * 3600,  # Default 5 days in seconds
+    }
+
+    package_url = None
+    package_sha = None
+
+    # Refer https://github.com/Netflix/metaflow/blob/master/docs/lifecycle.png
+    # to understand where these functions are invoked in the lifecycle of a
+    # Metaflow flow.
+    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 UnsupportedNvcfDatastoreException(flow_datastore.TYPE)
+
+        # Set internal state.
+        self.logger = logger
+        self.environment = environment
+        self.step = step
+        self.flow_datastore = flow_datastore
+
+        if any([deco.name == "kubernetes" for deco in decos]):
+            raise MetaflowException(
+                "Step *{step}* is marked for execution both on Kubernetes and "
+                "Nvidia. Please use one or the other.".format(step=step)
+            )
+        if any([isinstance(deco, ParallelDecorator) for deco in decos]):
+            raise MetaflowException(
+                "Step *{step}* contains a @parallel decorator "
+                "with the @nvidia decorator. @parallel decorators are not currently supported with @nvidia.".format(
+                    step=step
+                )
+            )
+
+        # 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 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]
+
+        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
+        self.graph = graph
+        self.package = package
+        self.run_id = run_id
+
+    def runtime_task_created(
+        self, task_datastore, task_id, split_index, input_paths, is_cloned, ubf_context
+    ):
+        if not is_cloned:
+            self._save_package_once(self.flow_datastore, self.package)
+
+    def runtime_step_cli(
+        self, cli_args, retry_count, max_user_code_retries, ubf_context
+    ):
+        if retry_count <= max_user_code_retries:
+            # after all attempts to run the user code have failed, we don't need
+            # to execute on NVCF anymore. We can execute possible fallback
+            # code locally.
+            cli_args.commands = ["nvidia", "step"]
+            cli_args.command_args.append(self.package_sha)
+            cli_args.command_args.append(self.package_url)
+            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
+
+    def task_pre_step(
+        self,
+        step_name,
+        task_datastore,
+        metadata,
+        run_id,
+        task_id,
+        flow,
+        graph,
+        retry_count,
+        max_retries,
+        ubf_context,
+        inputs,
+    ):
+        self.metadata = metadata
+        self.task_datastore = task_datastore
+
+        # task_pre_step may run locally if fallback is activated for @catch
+        # decorator.
+
+        if "NVCF_CONTEXT" in os.environ:
+            meta = {}
+
+            meta["nvcf-function-id"] = os.environ.get("NVCF_FUNCTION_ID")
+            meta["nvcf-function-version-id"] = os.environ.get(
+                "NVCF_FUNCTION_VERSION_ID"
+            )
+            meta["nvcf-region"] = os.environ.get("NVCF_REGION")
+            meta["nvcf-ncaid"] = os.environ.get("NVCF_NCAID")
+            meta["nvcf-sub"] = os.environ.get("NVCF_SUB")
+            meta["nvcf-instancetype"] = os.environ.get("NVCF_INSTANCETYPE")
+            meta["nvcf-reqid"] = os.environ.get("NVCF_REQID")
+            meta["nvcf-env"] = os.environ.get("NVCF_ENV")
+            meta["nvcf-backend"] = os.environ.get("NVCF_BACKEND")
+            meta["nvcf-function-name"] = os.environ.get("NVCF_FUNCTION_NAME")
+            meta["nvcf-nspectid"] = os.environ.get("NVCF_NSPECTID")
+
+            entries = [
+                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
+            ]
+            # Register book-keeping metadata for debugging.
+            metadata.register_metadata(run_id, step_name, task_id, entries)
+
+            self._save_logs_sidecar = Sidecar("save_logs_periodically")
+            self._save_logs_sidecar.start()
+
+    def task_finished(
+        self, step_name, flow, graph, is_task_ok, retry_count, max_retries
+    ):
+        # task_finished may run locally if fallback is activated for @catch
+        # decorator.
+        if "NVCF_CONTEXT" in os.environ:
+            # If `local` metadata is configured, we would need to copy task
+            # execution metadata from the NVCF container to user's
+            # local file system after the user code has finished execution.
+            # This happens via datastore as a communication bridge.
+            if hasattr(self, "metadata") and self.metadata.TYPE == "local":
+                sync_local_metadata_to_datastore(
+                    DATASTORE_LOCAL_DIR, self.task_datastore
+                )
+
+        try:
+            self._save_logs_sidecar.terminate()
+        except:
+            # Best effort kill
+            pass
+
+    @classmethod
+    def _save_package_once(cls, flow_datastore, package):
+        if cls.package_url is None:
+            cls.package_url, cls.package_sha = flow_datastore.save_data(
+                [package.blob], len_hint=1
+            )[0]

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)

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)