ob-metaflow-extensions 1.1.150__tar.gz → 1.1.151__tar.gz

{ob-metaflow-extensions-1.1.150 → ob-metaflow-extensions-1.1.151}/PKG-INFO

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

{ob-metaflow-extensions-1.1.150 → ob-metaflow-extensions-1.1.151}/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",

ob-metaflow-extensions-1.1.151/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)

ob-metaflow-extensions-1.1.151/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

ob-metaflow-extensions-1.1.151/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()