ob-metaflow-extensions 1.1.149__py2.py3-none-any.whl → 1.1.151__py2.py3-none-any.whl

metaflow_extensions/outerbounds/plugins/__init__.py

@@ -312,11 +312,13 @@ class ObpGcpAuthProvider(object):
 GCP_CLIENT_PROVIDERS_DESC = [("obp", ".ObpGcpAuthProvider")]
 CLIS_DESC = [
     ("nvidia", ".nvcf.nvcf_cli.cli"),
+    ("nvct", ".nvct.nvct_cli.cli"),
     ("fast-bakery", ".fast_bakery.fast_bakery_cli.cli"),
     ("snowpark", ".snowpark.snowpark_cli.cli"),
 ]
 STEP_DECORATORS_DESC = [
     ("nvidia", ".nvcf.nvcf_decorator.NvcfDecorator"),
+    ("nvct", ".nvct.nvct_decorator.NvctDecorator"),
     (
         "fast_bakery_internal",
         ".fast_bakery.fast_bakery_decorator.InternalFastBakeryDecorator",

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

metaflow_extensions/outerbounds/plugins/nvct/nvct_cli.py

@@ -0,0 +1,289 @@
+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,
+)
+from metaflow.mflog import TASK_LOG_SOURCE
+from .nvct_runner import NvctRunner
+from .nvct import NVCTClient
+from .utils import get_ngc_api_key
+from .exceptions import NvctKilledException
+
+
+@click.group()
+def cli():
+    pass
+
+
+@cli.group(help="Commands related to nvct.")
+def nvct():
+    pass
+
+
+@nvct.command(help="List steps / tasks running as an nvct job.")
+@click.option(
+    "--run-id",
+    default=None,
+    required=True,
+    help="List unfinished and running 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 "nvct-task-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)
+
+
+@nvct.command(help="Cancel steps / tasks running as an nvct 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):
+    ngc_api_key = get_ngc_api_key()
+    nvct_client = NVCTClient(api_key=ngc_api_key)
+
+    flow_name = ctx.obj.flow.name
+    run_obj = Run(pathspec=f"{flow_name}/{run_id}", _namespace_check=False)
+    tasks_cancelled = []
+
+    for each_step in run_obj:
+        for each_task in each_step:
+            if not each_task.finished and "nvct-task-id" in each_task.metadata_dict:
+                task_pathspec = each_task.pathspec
+                attempt = each_task.metadata_dict.get("attempt")
+                _, _, step_name, task_id = task_pathspec.split("/")
+
+                nvct_task_id = each_task.metadata_dict.get("nvct-task-id")
+                nvct_client.cancel(nvct_task_id)
+
+                tasks_cancelled.append(
+                    f"[{nvct_task_id}] -- Flow Name: {flow_name}, Run ID: {run_id}, Step Name: {step_name}, Task ID: {task_id}, Retry Count: {attempt} is cancelled."
+                )
+
+    if tasks_cancelled:
+        for each_cancelled_task in tasks_cancelled:
+            ctx.obj.echo(each_cancelled_task)
+
+
+@nvct.command(
+    help="Execute a single task using @nvct. This command calls the "
+    "top-level step command inside an nvct 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("--gpu-type", help="Type of Nvidia GPU to use.")
+@click.option("--instance-type", help="Instance type to use.")
+@click.option("--backend", help="Backend to use.")
+@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'.")
+@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,
+    gpu_type,
+    instance_type,
+    backend,
+    ngc_api_key,
+    **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": "nvct",
+        "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,
+        "NVCT_CONTEXT": "1",
+    }
+
+    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"]
+                ),
+            )
+
+    nvct = NvctRunner(
+        ctx.obj.metadata,
+        ctx.obj.flow_datastore,
+        ctx.obj.environment,
+        gpu_type,
+        instance_type,
+        backend,
+        ngc_api_key,
+    )
+    try:
+        with ctx.obj.monitor.measure("metaflow.nvct.launch_task"):
+            nvct.launch_task(
+                step_name,
+                step_cli,
+                task_spec,
+                code_package_sha,
+                code_package_url,
+                ctx.obj.flow_datastore.TYPE,
+                env=env,
+            )
+    except Exception:
+        traceback.print_exc()
+        _sync_metadata()
+        sys.exit(METAFLOW_EXIT_DISALLOW_RETRY)
+    try:
+        nvct.wait_for_completion(stdout_location, stderr_location, echo=echo)
+    except NvctKilledException:
+        # don't retry killed tasks
+        traceback.print_exc()
+        sys.exit(METAFLOW_EXIT_DISALLOW_RETRY)
+    finally:
+        _sync_metadata()

metaflow_extensions/outerbounds/plugins/nvct/nvct_decorator.py

@@ -0,0 +1,262 @@
+import os
+import sys
+
+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 .utils import get_ngc_api_key
+from .exceptions import (
+    UnsupportedNvctDatastoreException,
+    NvctTimeoutTooShortException,
+    RequestedGPUTypeUnavailableException,
+    UnsupportedNvctConfigurationException,
+)
+
+
+DEFAULT_GPU_TYPE = "H100"
+
+SUPPORTABLE_GPU_TYPES = {
+    "L40": [
+        {
+            "n_gpus": 1,
+            "instance_type": "gl40_1.br20_2xlarge",
+            "backend": "GFN",
+        },
+    ],
+    "L40S": [
+        {
+            "n_gpus": 1,
+            "instance_type": "gl40s_1.br25_2xlarge",
+            "backend": "GFN",
+        },
+    ],
+    "L40G": [
+        {
+            "n_gpus": 1,
+            "instance_type": "gl40g_1.br25_2xlarge",
+            "backend": "GFN",
+        },
+    ],
+    "H100": [
+        {
+            "n_gpus": 1,
+            "instance_type": "GCP.GPU.H100_1x",
+            "backend": "gcp-asia-se-1a",
+        },
+        {
+            "n_gpus": 2,
+            "instance_type": "GCP.GPU.H100_2x",
+            "backend": "gcp-asia-se-1a",
+        },
+        {
+            "n_gpus": 4,
+            "instance_type": "GCP.GPU.H100_4x",
+            "backend": "gcp-asia-se-1a",
+        },
+        {
+            "n_gpus": 8,
+            "instance_type": "GCP.GPU.H100_8x",
+            "backend": "gcp-asia-se-1a",
+        },
+    ],
+}
+
+
+class NvctDecorator(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 = "nvct"
+    defaults = {
+        "gpu": 1,
+        "gpu_type": None,
+        "ngc_api_key": None,
+        "instance_type": None,
+        "backend": None,
+    }
+
+    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 NVCT functions requires a non-local datastore.
+        if flow_datastore.TYPE not in ("s3", "azure", "gs"):
+            raise UnsupportedNvctDatastoreException(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 "
+                "Nvct. 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 @nvct decorator. @parallel decorators are not currently supported with @nvct.".format(
+                    step=step
+                )
+            )
+
+        # Set run time limit for NVCT.
+        self.run_time_limit = get_run_time_limit_for_task(decos)
+        if self.run_time_limit < 60:
+            raise NvctTimeoutTooShortException(step)
+
+        self.attributes["ngc_api_key"] = get_ngc_api_key()
+
+        requested_gpu_type = self.attributes["gpu_type"]
+        requested_n_gpus = 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, list(SUPPORTABLE_GPU_TYPES.keys())
+            )
+
+        valid_config = None
+        available_configurations = SUPPORTABLE_GPU_TYPES[requested_gpu_type]
+        for each_config in available_configurations:
+            if each_config["n_gpus"] == requested_n_gpus:
+                valid_config = each_config
+                break
+
+        if valid_config is None:
+            raise UnsupportedNvctConfigurationException(
+                requested_n_gpus,
+                requested_gpu_type,
+                available_configurations,
+                step,
+            )
+
+        self.attributes["instance_type"] = valid_config["instance_type"]
+        self.attributes["gpu_type"] = requested_gpu_type
+        self.attributes["backend"] = valid_config["backend"]
+
+    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 = ["nvct", "step"]
+            cli_args.command_args.append(self.package_sha)
+            cli_args.command_args.append(self.package_url)
+            cli_options = {
+                "gpu_type": self.attributes["gpu_type"],
+                "instance_type": self.attributes["instance_type"],
+                "backend": self.attributes["backend"],
+                "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(
+        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 "NVCT_CONTEXT" in os.environ:
+            meta = {}
+
+            meta["nvct-task-id"] = os.environ.get("NVCT_TASK_ID")
+            meta["nvct-task-name"] = os.environ.get("NVCT_TASK_NAME")
+            meta["nvct-ncaid"] = os.environ.get("NVCT_NCA_ID")
+            meta["nvct-progress-file-path"] = os.environ.get("NVCT_PROGRESS_FILE_PATH")
+            meta["nvct-results-dir"] = os.environ.get("NVCT_RESULTS_DIR")
+
+            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 "NVCT_CONTEXT" in os.environ:
+            # If `local` metadata is configured, we would need to copy task
+            # execution metadata from the NVCT 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/nvct/nvct_runner.py

@@ -0,0 +1,218 @@
+import os
+import re
+import time
+import math
+import shlex
+import atexit
+
+from metaflow import util
+from metaflow.mflog import (
+    BASH_SAVE_LOGS,
+    bash_capture_logs,
+    export_mflog_env_vars,
+    tail_logs,
+    get_log_tailer,
+)
+
+from .nvct import NVCTClient, NVCTTask, NVCTRequest
+from .exceptions import (
+    NvctKilledException,
+    NvctExecutionException,
+    NvctTaskFailedException,
+)
+
+# Constants for Metaflow logs
+LOGS_DIR = "$PWD/.logs"
+STDOUT_FILE = "mflog_stdout"
+STDERR_FILE = "mflog_stderr"
+STDOUT_PATH = os.path.join(LOGS_DIR, STDOUT_FILE)
+STDERR_PATH = os.path.join(LOGS_DIR, STDERR_FILE)
+NVCT_WRAPPER = "/usr/local/bin/nvct-wrapper.sh"
+
+
+class NvctRunner:
+    def __init__(
+        self,
+        metadata,
+        datastore,
+        environment,
+        gpu_type,
+        instance_type,
+        backend,
+        ngc_api_key,
+    ):
+        self.metadata = metadata
+        self.datastore = datastore
+        self.environment = environment
+        self.gpu_type = gpu_type
+        self.instance_type = instance_type
+        self.backend = backend
+        self._ngc_api_key = ngc_api_key
+        self.client = None
+        self.task = None
+        atexit.register(lambda: self.task.cancel() if hasattr(self, "task") else None)
+
+    def launch_task(
+        self,
+        step_name,
+        step_cli,
+        task_spec,
+        code_package_sha,
+        code_package_url,
+        code_package_ds,
+        env={},
+        max_runtime="PT7H",  # <8H allowed for GFN backend
+        max_queued="PT120H",  # 5 days
+    ):
+        mflog_expr = export_mflog_env_vars(
+            datastore_type=code_package_ds,
+            stdout_path=STDOUT_PATH,
+            stderr_path=STDERR_PATH,
+            **task_spec,
+        )
+        init_cmds = self.environment.get_package_commands(
+            code_package_url, code_package_ds
+        )
+        init_expr = " && ".join(init_cmds)
+        step_expr = bash_capture_logs(
+            " && ".join(
+                self.environment.bootstrap_commands(step_name, code_package_ds)
+                + [step_cli]
+            )
+        )
+        cmd_str = "mkdir -p %s && %s && %s && %s; c=$?; %s; exit $c" % (
+            LOGS_DIR,
+            mflog_expr,
+            init_expr,
+            step_expr,
+            BASH_SAVE_LOGS,
+        )
+
+        # Add optional initialization script execution
+        cmd_str = (
+            '${METAFLOW_INIT_SCRIPT:+eval \\"${METAFLOW_INIT_SCRIPT}\\"} && %s'
+            % cmd_str
+        )
+
+        cmd_str = shlex.split('bash -c "%s"' % cmd_str)[-1]
+
+        def modify_python_c(match):
+            content = match.group(1)
+            # Escape double quotes within the python -c command
+            content = content.replace('"', r"\"")
+            # Replace outermost double quotes with single quotes
+            return 'python -c "%s"' % content
+
+        # Convert python -c single quotes to double quotes
+        cmd_str = re.sub(r"python -c '(.*?)'", modify_python_c, cmd_str)
+        cmd_str = cmd_str.replace("'", '"')
+        # Create the final command with outer single quotes to pass to NVCT wrapper
+        nvct_cmd = f"{NVCT_WRAPPER} bash -c '{cmd_str}'"
+
+        flow_name = task_spec.get("flow_name")
+        run_id = task_spec.get("run_id")
+        task_id = task_spec.get("task_id")
+        retry_count = task_spec.get("retry_count")
+        task_name = f"{flow_name}-{run_id}-{step_name}-{task_id}-{retry_count}"
+
+        if self.backend != "GFN":
+            # if maxRuntimeDuration exceeds 8 hours for a Task on the GFN backend,
+            # the request will be rejected.
+            # (https://docs.nvidia.com/cloud-functions/user-guide/latest/cloud-function/tasks.html#create-task)
+            ## thus, if it is non GFN backend, we increase it to 3 days
+            max_runtime = "PT72H"
+
+        request = (
+            NVCTRequest(task_name)
+            .container_image("nvcr.io/zhxkmsaasxhw/nvct-base:1.0-jovyan")
+            .container_args(nvct_cmd)
+            .gpu(
+                gpu=self.gpu_type,
+                instance_type=self.instance_type,
+                backend=self.backend,
+            )
+            .max_runtime(max_runtime)
+            .max_queued(max_queued)
+        )
+
+        for k, v in env.items():
+            if v is not None:
+                request.env(k, str(v))
+
+        self.client = NVCTClient(self._ngc_api_key)
+        self.task = NVCTTask(self.client, request.to_dict())
+
+        self.task.submit()
+        return self.task.id
+
+    def wait_for_completion(self, stdout_location, stderr_location, echo=None):
+        if not self.task:
+            raise NvctExecutionException("No task has been launched")
+
+        def update_delay(secs_since_start):
+            # this sigmoid function reaches
+            # - 0.1 after 11 minutes
+            # - 0.5 after 15 minutes
+            # - 1.0 after 23 minutes
+            # in other words, the user will see very frequent updates
+            # during the first 10 minutes
+            sigmoid = 1.0 / (1.0 + math.exp(-0.01 * secs_since_start + 9.0))
+            return 0.5 + sigmoid * 30.0
+
+        def wait_for_launch(task):
+            status = task.status
+            echo(
+                "Task is starting (%s)..." % status,
+                "stderr",
+                _id=task.id,
+            )
+
+            t = time.time()
+            start_time = time.time()
+            while task.is_waiting:
+                new_status = task.status
+                if status != new_status or (time.time() - t) > 30:
+                    status = new_status
+                    echo(
+                        "Task is starting (%s)..." % status,
+                        "stderr",
+                        _id=task.id,
+                    )
+                    t = time.time()
+                time.sleep(update_delay(time.time() - start_time))
+
+        _make_prefix = lambda: b"[%s] " % util.to_bytes(self.task.id)
+        stdout_tail = get_log_tailer(stdout_location, self.datastore.TYPE)
+        stderr_tail = get_log_tailer(stderr_location, self.datastore.TYPE)
+
+        # 1) Loop until the job has started
+        wait_for_launch(self.task)
+
+        echo(
+            "Task is starting (%s)..." % self.task.status,
+            "stderr",
+            _id=self.task.id,
+        )
+
+        # 2) Tail logs until the job has finished
+        tail_logs(
+            prefix=_make_prefix(),
+            stdout_tail=stdout_tail,
+            stderr_tail=stderr_tail,
+            echo=echo,
+            has_log_updates=lambda: self.task.is_running,
+        )
+
+        if self.task.has_failed:
+            raise NvctTaskFailedException(
+                f"Task failed with status: {self.task.status}. This could be a transient error. Use @retry to retry."
+            )
+        else:
+            if self.task.is_running:
+                # Kill the job if it is still running by throwing an exception.
+                raise NvctKilledException("Task failed!")
+            echo(
+                f"Task finished with status: {self.task.status}",
+                "stderr",
+                _id=self.task.id,
+            )

metaflow_extensions/outerbounds/plugins/nvct/utils.py

@@ -0,0 +1,29 @@
+import os
+import json
+import requests
+from urllib.parse import urlparse
+from metaflow.metaflow_config import SERVICE_URL
+from metaflow.metaflow_config_funcs import init_config
+
+
+def get_ngc_api_key():
+    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 @nvct.
+    # 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()
+    return res.json()["nvcf"]["api_key"]

{ob_metaflow_extensions-1.1.149.dist-info → ob_metaflow_extensions-1.1.151.dist-info}/METADATA

@@ -1,13 +1,13 @@
 Metadata-Version: 2.1
 Name: ob-metaflow-extensions
-Version: 1.1.149
+Version: 1.1.151
 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.15.10.1)
+Requires-Dist: ob-metaflow (==2.15.11.1)
 
 # Outerbounds platform package
 

{ob_metaflow_extensions-1.1.149.dist-info → ob_metaflow_extensions-1.1.151.dist-info}/RECORD

@@ -1,7 +1,7 @@
 metaflow_extensions/outerbounds/__init__.py,sha256=TRGvIUMjkfneWtYUFSWoubu_Kf2ekAL4WLbV3IxOj9k,499
 metaflow_extensions/outerbounds/remote_config.py,sha256=Zpfpjgz68_ZgxlXezjzlsDLo4840rkWuZgwDB_5H57U,4059
 metaflow_extensions/outerbounds/config/__init__.py,sha256=JsQGRuGFz28fQWjUvxUgR8EKBLGRdLUIk_buPLJplJY,1225
-metaflow_extensions/outerbounds/plugins/__init__.py,sha256=LOQLHsbvHWhhjZUQ99kMVl57wWCZFOxFQGblQ5kYarQ,13140
+metaflow_extensions/outerbounds/plugins/__init__.py,sha256=eHcM_t2Mzlge7B9Dv3VGVM5x8qNZYdLyqBOAC6uRxec,13228
 metaflow_extensions/outerbounds/plugins/auth_server.py,sha256=_Q9_2EL0Xy77bCRphkwT1aSu8gQXRDOH-Z-RxTUO8N4,2202
 metaflow_extensions/outerbounds/plugins/perimeters.py,sha256=QXh3SFP7GQbS-RAIxUOPbhPzQ7KDFVxZkTdKqFKgXjI,2697
 metaflow_extensions/outerbounds/plugins/apps/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -35,6 +35,13 @@ metaflow_extensions/outerbounds/plugins/nvcf/nvcf.py,sha256=3ZFdYItVpFWnHMOeyV1n
 metaflow_extensions/outerbounds/plugins/nvcf/nvcf_cli.py,sha256=3D-r5XO88Yh2k1EAZFJTe_PwdbhWp5qXflG8AgE4ZIU,9500
 metaflow_extensions/outerbounds/plugins/nvcf/nvcf_decorator.py,sha256=pSWKaPyGXBEfxG35QA1FQljio8ADjwf-DnPgEsqXoUM,9251
 metaflow_extensions/outerbounds/plugins/nvcf/utils.py,sha256=DxWSCayfa95e0HJkWacey1s1nxoTpaunGhrb_0Ayv28,133
+metaflow_extensions/outerbounds/plugins/nvct/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+metaflow_extensions/outerbounds/plugins/nvct/exceptions.py,sha256=1PiV6FdH36CvkmHh5jtsfrsoe3Q_Fo1NomHw5wvgoDM,2886
+metaflow_extensions/outerbounds/plugins/nvct/nvct.py,sha256=Z2ZPWGuHe58au_d6GfHiw6Nl5d8INdLDI5exlsPEOSA,3564
+metaflow_extensions/outerbounds/plugins/nvct/nvct_cli.py,sha256=bB9AURhRep9PV_-b-qLHpgw_GPG_xFoq1PeHEgFP1mQ,10104
+metaflow_extensions/outerbounds/plugins/nvct/nvct_decorator.py,sha256=LaJ_Tk-vNjvrglzSTR-U6pk8f9MtQRKObU9m7vBYtkI,8695
+metaflow_extensions/outerbounds/plugins/nvct/nvct_runner.py,sha256=D2sEtVFVWXBqaQEsminZYI_WesR2kADLmwv4lsxVBhk,7091
+metaflow_extensions/outerbounds/plugins/nvct/utils.py,sha256=U4_Fu8H94j_Bbox7mmMhNnlRhlYHqnK28R5w_TMWEFM,1029
 metaflow_extensions/outerbounds/plugins/ollama/__init__.py,sha256=HEsI5U4ckQby7K2NsGBOdizhPY3WWqXSnXx_IHL7_No,2307
 metaflow_extensions/outerbounds/plugins/ollama/ollama.py,sha256=KlP8_EmnUoi8-PidyU0IDuENYxKjQaHFC33yGsvaeic,13320
 metaflow_extensions/outerbounds/plugins/profilers/deco_injector.py,sha256=oI_C3c64XBm7n88FILqHwn-Nnc5DeT_68I67lM9rXaI,2434
@@ -61,7 +68,7 @@ metaflow_extensions/outerbounds/toplevel/plugins/gcp/__init__.py,sha256=BbZiaH3u
 metaflow_extensions/outerbounds/toplevel/plugins/kubernetes/__init__.py,sha256=5zG8gShSj8m7rgF4xgWBZFuY3GDP5n1T0ktjRpGJLHA,69
 metaflow_extensions/outerbounds/toplevel/plugins/ollama/__init__.py,sha256=GRSz2zwqkvlmFS6bcfYD_CX6CMko9DHQokMaH1iBshA,47
 metaflow_extensions/outerbounds/toplevel/plugins/snowflake/__init__.py,sha256=LptpH-ziXHrednMYUjIaosS1SXD3sOtF_9_eRqd8SJw,50
-ob_metaflow_extensions-1.1.149.dist-info/METADATA,sha256=e-FFRtty5ShVlDzjNq_jpVLZbMA7Ul5LtSDa9_oGcaQ,521
-ob_metaflow_extensions-1.1.149.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-ob_metaflow_extensions-1.1.149.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
-ob_metaflow_extensions-1.1.149.dist-info/RECORD,,
+ob_metaflow_extensions-1.1.151.dist-info/METADATA,sha256=x5PLR9aAaWhuhLDERP5mSl44te5I0ZNmNxmjBsRgrzg,521
+ob_metaflow_extensions-1.1.151.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+ob_metaflow_extensions-1.1.151.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
+ob_metaflow_extensions-1.1.151.dist-info/RECORD,,