ob-metaflow-extensions 1.1.88__tar.gz → 1.1.89__tar.gz

{ob-metaflow-extensions-1.1.88 → ob-metaflow-extensions-1.1.89}/PKG-INFO

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

{ob-metaflow-extensions-1.1.88 → ob-metaflow-extensions-1.1.89}/metaflow_extensions/outerbounds/plugins/__init__.py

@@ -306,7 +306,7 @@ class ObpGcpAuthProvider(object):
 
 GCP_CLIENT_PROVIDERS_DESC = [("obp", ".ObpGcpAuthProvider")]
 CLIS_DESC = [
-    ("nvcf", ".nvcf.nvcf_cli.cli"),
+    ("nvidia", ".nvcf.nvcf_cli.cli"),
     ("fast-bakery", ".fast_bakery.fast_bakery_cli.cli"),
     ("snowpark", ".snowpark.snowpark_cli.cli"),
 ]

{ob-metaflow-extensions-1.1.88 → ob-metaflow-extensions-1.1.89}/metaflow_extensions/outerbounds/plugins/fast_bakery/docker_environment.py

@@ -15,6 +15,7 @@ from metaflow.plugins.kubernetes.kubernetes_decorator import KubernetesDecorator
 from metaflow.plugins.pypi.conda_decorator import CondaStepDecorator
 from metaflow.plugins.pypi.conda_environment import CondaEnvironment
 from metaflow.plugins.pypi.pypi_decorator import PyPIStepDecorator
+from metaflow import decorators
 
 from .fast_bakery import FastBakery, FastBakeryApiResponse, FastBakeryException
 
@@ -26,8 +27,6 @@ import os
 from concurrent.futures import ThreadPoolExecutor
 from functools import wraps
 
-# TODO - ensure that both @conda/@pypi are not assigned to the same step
-
 
 def cache_request(cache_file):
     def decorator(func):
@@ -103,7 +102,10 @@ class DockerEnvironment(MetaflowEnvironment):
         self.local_root = local_root
 
     def decospecs(self):
-        return ("conda", "fast_bakery_internal") + super().decospecs()
+        # Due to conflicts with the CondaEnvironment fallback and bakery,
+        # we can not simply attach 'conda' or 'pypi' to all steps here.
+        # Instead we do this on a per-step basis in init_environment
+        return ("fast_bakery_internal",) + super().decospecs()
 
     def validate_environment(self, logger, datastore_type):
         self.datastore_type = datastore_type
@@ -116,14 +118,7 @@ class DockerEnvironment(MetaflowEnvironment):
 
         # Mixing @pypi/@conda in a single step is not supported yet
         for step in self.flow:
-            if (
-                sum(
-                    1
-                    for deco in step.decorators
-                    if isinstance(deco, (PyPIStepDecorator, CondaStepDecorator))
-                )
-                > 1
-            ):
+            if sum(1 for deco in step.decorators if _is_env_deco(deco)) > 1:
                 raise MetaflowException(
                     "Mixing and matching PyPI packages and Conda packages within a\n"
                     "step is not yet supported. Use one of @pypi or @conda only for the *%s* step."
@@ -132,13 +127,31 @@ class DockerEnvironment(MetaflowEnvironment):
 
     def init_environment(self, echo):
         self.skipped_steps = {
-            step.name
-            for step in self.flow
-            if not any(
-                isinstance(deco, (BatchDecorator, KubernetesDecorator))
-                for deco in step.decorators
-            )
+            step.name for step in self.flow if not _step_executes_remotely(step)
         }
+        # Attach environment decorator as needed. This is done on a step-by-step basis
+        # as we require a conda decorator for fallback steps, but prefer pypi for the baked ones.
+        for step in self.flow:
+            if not _step_has_environment_deco(step):
+                if step.name in self.skipped_steps:
+                    # Conda fallback requires a conda decorator as the default for a step
+                    decorators._attach_decorators_to_step(step, ["conda"])
+                else:
+                    # We default to PyPI for steps that are going to be baked.
+                    decorators._attach_decorators_to_step(step, ["pypi"])
+                # Initialize the decorator we attached.
+                # This is crucial for the conda decorator to work properly in the fallback environment
+                for deco in step.decorators:
+                    if _is_env_deco(deco):
+                        deco.step_init(
+                            self.flow,
+                            None,  # not passing graph as it is not available, and not required by conda/pypi decorators
+                            step.name,
+                            step.decorators,
+                            self,
+                            self.datastore,
+                            echo,
+                        )
 
         steps_to_bake = [
             step for step in self.flow if step.name not in self.skipped_steps
@@ -149,7 +162,7 @@ class DockerEnvironment(MetaflowEnvironment):
             self.results = self._bake(steps_to_bake)
             for step in self.flow:
                 for d in step.decorators:
-                    if isinstance(d, (BatchDecorator, KubernetesDecorator)):
+                    if _is_remote_deco(d):
                         d.attributes["image"] = self.results[step.name].container_image
                         d.attributes["executable"] = self.results[step.name].python_path
             if self.images_baked > 0:
@@ -178,32 +191,32 @@ class DockerEnvironment(MetaflowEnvironment):
             conda_packages=None,
             base_image=None,
         ):
-            bakery = FastBakery(url=FAST_BAKERY_URL)
-            bakery._reset_payload()
-            bakery.python_version(python)
-            bakery.pypi_packages(pypi_packages)
-            bakery.conda_packages(conda_packages)
-            bakery.base_image(base_image)
-            # bakery.ignore_cache()
-
-            with logger_lock:
-                self.logger(f"🍳 Baking [{ref}] ...")
-                self.logger(f"     🐍 Python: {python}")
-
-                if pypi_packages:
-                    self.logger(f"     📦 PyPI packages:")
-                    for package, version in pypi_packages.items():
-                        self.logger(f"        🔧 {package}: {version}")
-
-                if conda_packages:
-                    self.logger(f"     📦 Conda packages:")
-                    for package, version in conda_packages.items():
-                        self.logger(f"        🔧 {package}: {version}")
-
-                self.logger(f"     🏗️  Base image: {base_image}")
-
-            start_time = time.time()
             try:
+                bakery = FastBakery(url=FAST_BAKERY_URL)
+                bakery._reset_payload()
+                bakery.python_version(python)
+                bakery.pypi_packages(pypi_packages)
+                bakery.conda_packages(conda_packages)
+                bakery.base_image(base_image)
+                # bakery.ignore_cache()
+
+                with logger_lock:
+                    self.logger(f"🍳 Baking [{ref}] ...")
+                    self.logger(f"     🐍 Python: {python}")
+
+                    if pypi_packages:
+                        self.logger(f"     📦 PyPI packages:")
+                        for package, version in pypi_packages.items():
+                            self.logger(f"        🔧 {package}: {version}")
+
+                    if conda_packages:
+                        self.logger(f"     📦 Conda packages:")
+                        for package, version in conda_packages.items():
+                            self.logger(f"        🔧 {package}: {version}")
+
+                    self.logger(f"     🏗️  Base image: {base_image}")
+
+                start_time = time.time()
                 res = bakery.bake()
                 # TODO: Get actual bake time from bakery
                 bake_time = time.time() - start_time
@@ -225,11 +238,7 @@ class DockerEnvironment(MetaflowEnvironment):
                 None,
             )
             dependencies = next(
-                (
-                    d
-                    for d in step.decorators
-                    if isinstance(d, (CondaStepDecorator, PyPIStepDecorator))
-                ),
+                (d for d in step.decorators if _is_env_deco(d)),
                 None,
             )
             python = next(
@@ -322,3 +331,22 @@ class DockerEnvironment(MetaflowEnvironment):
 
 def get_fastbakery_metafile_path(local_root, flow_name):
     return os.path.join(local_root, flow_name, BAKERY_METAFILE)
+
+
+def _is_remote_deco(deco):
+    return isinstance(deco, (BatchDecorator, KubernetesDecorator))
+
+
+def _step_executes_remotely(step):
+    "Check if a step is going to execute remotely or locally"
+    return any(_is_remote_deco(deco) for deco in step.decorators)
+
+
+def _is_env_deco(deco):
+    "Check if a decorator is a known environment decorator (conda/pypi)"
+    return isinstance(deco, (PyPIStepDecorator, CondaStepDecorator))
+
+
+def _step_has_environment_deco(step):
+    "Check if a step has a virtual environment decorator"
+    return any(_is_env_deco(deco) for deco in step.decorators)

{ob-metaflow-extensions-1.1.88 → ob-metaflow-extensions-1.1.89}/metaflow_extensions/outerbounds/plugins/fast_bakery/fast_bakery.py

@@ -73,6 +73,8 @@ class FastBakeryApiResponse:
 
 class FastBakery:
     def __init__(self, url: str):
+        if not url:
+            raise FastBakeryException("Specifying a url is required.")
         self.url = url
         self.headers = {"Content-Type": "application/json", "Connection": "keep-alive"}
         self._reset_payload()

{ob-metaflow-extensions-1.1.88 → ob-metaflow-extensions-1.1.89}/metaflow_extensions/outerbounds/plugins/nvcf/heartbeat_store.py

@@ -1,13 +1,29 @@
 import os
 import sys
 import time
-import signal
+import subprocess
 from io import BytesIO
 from datetime import datetime, timezone
 
 from metaflow.exception import MetaflowException
 
 
+def kill_process_and_descendants(pid, termination_timeout=5):
+    try:
+        subprocess.check_call(["pkill", "-TERM", "-P", str(pid)])
+        subprocess.check_call(["kill", "-TERM", str(pid)])
+    except subprocess.CalledProcessError:
+        pass
+
+    time.sleep(termination_timeout)
+
+    try:
+        subprocess.check_call(["pkill", "-KILL", "-P", str(pid)])
+        subprocess.check_call(["kill", "-KILL", str(pid)])
+    except subprocess.CalledProcessError:
+        pass
+
+
 class HeartbeatStore(object):
     def __init__(
         self,
@@ -67,7 +83,7 @@ class HeartbeatStore(object):
                 contents = f.read()
                 if "tombstone" in contents:
                     print("[Outerbounds] Tombstone detected. Terminating the task..")
-                    os.kill(self.main_pid, signal.SIGTERM)
+                    kill_process_and_descendants(self.main_pid)
                     sys.exit(1)
 
     def __handle_heartbeat(self, path):
@@ -86,7 +102,7 @@ class HeartbeatStore(object):
             print(
                 f"[Outerbounds] Missed {self.max_missed_heartbeats} consecutive heartbeats. Terminating the task.."
             )
-            os.kill(self.main_pid, signal.SIGTERM)
+            kill_process_and_descendants(self.main_pid)
             sys.exit(1)
 
     def is_main_process_running(self):

{ob-metaflow-extensions-1.1.88 → ob-metaflow-extensions-1.1.89}/metaflow_extensions/outerbounds/plugins/nvcf/nvcf_cli.py

@@ -39,13 +39,18 @@ def cli():
     pass
 
 
-@cli.group(help="Commands related to NVCF.")
-def nvcf():
+@cli.group(help="Commands related to nvidia.")
+def nvidia():
     pass
 
 
-@nvcf.command(help="List steps / tasks running as an NVCF job.")
-@click.argument("run-id")
+@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
@@ -65,13 +70,18 @@ def list(ctx, run_id):
 
     if running_invocations:
         for each_invocation in running_invocations:
-            print(each_invocation)
+            ctx.obj.echo(each_invocation)
     else:
-        print("No running NVCF invocations for Run ID: %s" % run_id)
+        ctx.obj.echo("No running @nvidia invocations for Run ID: %s" % run_id)
 
 
-@nvcf.command(help="Kill steps / tasks running as an NVCF job.")
-@click.argument("run-id")
+@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 (
@@ -100,12 +110,12 @@ def kill(ctx, run_id):
             )
             store.emit_tombstone(folder_name="nvcf_heartbeats")
         else:
-            print("No running NVCF invocations for Run ID: %s" % run_id)
+            ctx.obj.echo("No running @nvidia invocations for Run ID: %s" % run_id)
 
 
-@nvcf.command(
-    help="Execute a single task using NVCF. This command calls the "
-    "top-level step command inside a NVCF job with the given options. "
+@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."
 )

{ob-metaflow-extensions-1.1.88 → ob-metaflow-extensions-1.1.89}/metaflow_extensions/outerbounds/plugins/nvcf/nvcf_decorator.py

@@ -72,7 +72,7 @@ class NvcfDecorator(StepDecorator):
             # 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 = ["nvcf", "step"]
+            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)

{ob-metaflow-extensions-1.1.88 → ob-metaflow-extensions-1.1.89}/ob_metaflow_extensions.egg-info/PKG-INFO

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

ob-metaflow-extensions-1.1.89/ob_metaflow_extensions.egg-info/requires.txt

@@ -0,0 +1,3 @@
+boto3
+kubernetes
+ob-metaflow==2.12.19.1

{ob-metaflow-extensions-1.1.88 → ob-metaflow-extensions-1.1.89}/setup.py

@@ -2,7 +2,7 @@ from setuptools import setup, find_namespace_packages
 from pathlib import Path
 
 
-version = "1.1.88"
+version = "1.1.89"
 this_directory = Path(__file__).parent
 long_description = (this_directory / "README.md").read_text()
 
@@ -18,6 +18,6 @@ setup(
     install_requires=[
         "boto3",
         "kubernetes",
-        "ob-metaflow == 2.12.18.2",
+        "ob-metaflow == 2.12.19.1",
     ],
 )

ob-metaflow-extensions-1.1.88/ob_metaflow_extensions.egg-info/requires.txt

@@ -1,3 +0,0 @@
-boto3
-kubernetes
-ob-metaflow==2.12.18.2