metaflow 2.14.3__py2.py3-none-any.whl → 2.15.1__py2.py3-none-any.whl

metaflow/__init__.py

@@ -119,6 +119,13 @@ from .includefile import IncludeFile
 # Decorators
 from .decorators import step, _import_plugin_decorators
 
+# Config parsers
+from .plugins.pypi.parsers import (
+    requirements_txt_parser,
+    pyproject_toml_parser,
+    conda_environment_yml_parser,
+)
+
 # this auto-generates decorator functions from Decorator objects
 # in the top-level metaflow namespace
 _import_plugin_decorators(globals())

metaflow/cmd/code/__init__.py

@@ -0,0 +1,230 @@
+import os
+import shutil
+import sys
+from subprocess import PIPE, CompletedProcess, run
+from tempfile import TemporaryDirectory
+from typing import Any, Callable, List, Mapping, Optional, cast
+
+from metaflow import Run
+from metaflow._vendor import click
+from metaflow.cli import echo_always
+
+
+@click.group()
+def cli():
+    pass
+
+
+@cli.group(help="Access, compare, and manage code associated with Metaflow runs.")
+def code():
+    pass
+
+
+def echo(line: str) -> None:
+    echo_always(line, err=True, fg="magenta")
+
+
+def extract_code_package(runspec: str) -> TemporaryDirectory:
+    try:
+        mf_run = Run(runspec, _namespace_check=False)
+        echo(f"✅  Run *{runspec}* found, downloading code..")
+    except Exception as e:
+        echo(f"❌  Run **{runspec}** not found")
+        raise e
+
+    if mf_run.code is None:
+        echo(
+            f"❌  Run **{runspec}** doesn't have a code package. Maybe it's a local run?"
+        )
+        raise RuntimeError("no code package found")
+
+    return mf_run.code.extract()
+
+
+def perform_diff(
+    source_dir: str,
+    target_dir: Optional[str] = None,
+    output: bool = False,
+    **kwargs: Mapping[str, Any],
+) -> Optional[List[str]]:
+    if target_dir is None:
+        target_dir = os.getcwd()
+
+    diffs = []
+    for dirpath, dirnames, filenames in os.walk(source_dir, followlinks=True):
+        for fname in filenames:
+            # NOTE: the paths below need to be set up carefully
+            # for the `patch` command to work. Better not to touch
+            # the directories below. If you must, test that patches
+            # work after your changes.
+            #
+            # target_file is the git repo in the current working directory
+            rel = os.path.relpath(dirpath, source_dir)
+            target_file = os.path.join(rel, fname)
+            # source_file is the run file loaded in a tmp directory
+            source_file = os.path.join(dirpath, fname)
+
+            if sys.stdout.isatty() and not output:
+                color = ["--color"]
+            else:
+                color = ["--no-color"]
+
+            if os.path.exists(os.path.join(target_dir, target_file)):
+                cmd = (
+                    ["git", "diff", "--no-index", "--exit-code"]
+                    + color
+                    + [
+                        target_file,
+                        source_file,
+                    ]
+                )
+                result: CompletedProcess = run(
+                    cmd, text=True, stdout=PIPE, cwd=target_dir
+                )
+                if result.returncode == 0:
+                    if not output:
+                        echo(f"✅  {target_file} is identical, skipping")
+                    continue
+
+                if output:
+                    diffs.append(result.stdout)
+                else:
+                    run(["less", "-R"], input=result.stdout, text=True)
+            else:
+                if not output:
+                    echo(f"❗  {target_file} not in the target directory, skipping")
+    return diffs if output else None
+
+
+def run_op(
+    runspec: str, op: Callable[..., Optional[List[str]]], **op_args: Mapping[str, Any]
+) -> Optional[List[str]]:
+    tmp = None
+    try:
+        tmp = extract_code_package(runspec)
+        return op(tmp.name, **op_args)
+    finally:
+        if tmp and os.path.exists(tmp.name):
+            shutil.rmtree(tmp.name)
+
+
+def run_op_diff_runs(
+    source_run_pathspec: str, target_run_pathspec: str, **op_args: Mapping[str, Any]
+) -> Optional[List[str]]:
+    source_tmp = None
+    target_tmp = None
+    try:
+        source_tmp = extract_code_package(source_run_pathspec)
+        target_tmp = extract_code_package(target_run_pathspec)
+        return perform_diff(source_tmp.name, target_tmp.name, **op_args)
+    finally:
+        for d in [source_tmp, target_tmp]:
+            if d and os.path.exists(d.name):
+                shutil.rmtree(d.name)
+
+
+def op_diff(tmpdir: str, **kwargs: Mapping[str, Any]) -> Optional[List[str]]:
+    kwargs_dict = dict(kwargs)
+    target_dir = cast(Optional[str], kwargs_dict.pop("target_dir", None))
+    output: bool = bool(kwargs_dict.pop("output", False))
+    op_args: Mapping[str, Any] = {**kwargs_dict}
+    return perform_diff(tmpdir, target_dir=target_dir, output=output, **op_args)
+
+
+def op_pull(tmpdir: str, dst: str, **op_args: Mapping[str, Any]) -> None:
+    if os.path.exists(dst):
+        echo(f"❌  Directory *{dst}* already exists")
+    else:
+        shutil.move(tmpdir, dst)
+        echo(f"Code downloaded to *{dst}*")
+
+
+def op_patch(tmpdir: str, dst: str, **kwargs: Mapping[str, Any]) -> None:
+    diffs = perform_diff(tmpdir, output=True) or []
+    with open(dst, "w", encoding="utf-8") as f:
+        for out in diffs:
+            out = out.replace(tmpdir, "/.")
+            out = out.replace("+++ b/./", "+++ b/")
+            out = out.replace("--- b/./", "--- b/")
+            out = out.replace("--- a/./", "--- a/")
+            out = out.replace("+++ a/./", "+++ a/")
+            f.write(out)
+    echo(f"Patch saved in *{dst}*")
+    path = run(
+        ["git", "rev-parse", "--show-prefix"], text=True, stdout=PIPE
+    ).stdout.strip()
+    if path:
+        diropt = f" --directory={path.rstrip('/')}"
+    else:
+        diropt = ""
+    echo("Apply the patch by running:")
+    echo_always(
+        f"git apply --verbose{diropt} {dst}", highlight=True, bold=True, err=True
+    )
+
+
+@code.command()
+@click.argument("run_pathspec")
+def diff(run_pathspec: str, **kwargs: Mapping[str, Any]) -> None:
+    """
+    Do a 'git diff' of the current directory and a Metaflow run.
+    """
+    _ = run_op(run_pathspec, op_diff, **kwargs)
+
+
+@code.command()
+@click.argument("source_run_pathspec")
+@click.argument("target_run_pathspec")
+def diff_runs(
+    source_run_pathspec: str, target_run_pathspec: str, **kwargs: Mapping[str, Any]
+) -> None:
+    """
+    Do a 'git diff' between two Metaflow runs.
+    """
+    _ = run_op_diff_runs(source_run_pathspec, target_run_pathspec, **kwargs)
+
+
+@code.command()
+@click.argument("run_pathspec")
+@click.option(
+    "--dir", help="Destination directory (default: {run_pathspec}_code)", default=None
+)
+def pull(
+    run_pathspec: str, dir: Optional[str] = None, **kwargs: Mapping[str, Any]
+) -> None:
+    """
+    Pull the code of a Metaflow run.
+    """
+    if dir is None:
+        dir = run_pathspec.lower().replace("/", "_") + "_code"
+    op_args: Mapping[str, Any] = {**kwargs, "dst": dir}
+    run_op(run_pathspec, op_pull, **op_args)
+
+
+@code.command()
+@click.argument("run_pathspec")
+@click.option(
+    "--file_path",
+    help="Patch file name. If not provided, defaults to a sanitized version of RUN_PATHSPEC "
+    "with slashes replaced by underscores, plus '.patch'.",
+    show_default=False,
+)
+@click.option(
+    "--overwrite", is_flag=True, help="Overwrite the patch file if it exists."
+)
+def patch(
+    run_pathspec: str,
+    file_path: Optional[str] = None,
+    overwrite: bool = False,
+    **kwargs: Mapping[str, Any],
+) -> None:
+    """
+    Create a patch by comparing current dir with a Metaflow run.
+    """
+    if file_path is None:
+        file_path = run_pathspec.lower().replace("/", "_") + ".patch"
+    if os.path.exists(file_path) and not overwrite:
+        echo(f"File *{file_path}* already exists. To overwrite, specify --overwrite.")
+        return
+    op_args: Mapping[str, Any] = {**kwargs, "dst": file_path}
+    run_op(run_pathspec, op_patch, **op_args)

metaflow/cmd/develop/stub_generator.py

@@ -1133,13 +1133,16 @@ class StubGenerator:
             result = result[1:]
         # Add doc to first and last overloads. Jedi uses the last one and pycharm
         # the first one. Go figure.
+        result_docstring = docs["func_doc"]
+        if docs["param_doc"]:
+            result_docstring += "\nParameters\n----------\n" + docs["param_doc"]
         result[0] = (
             result[0][0],
-            docs["func_doc"] + "\nParameters\n----------\n" + docs["param_doc"],
+            result_docstring,
         )
         result[-1] = (
             result[-1][0],
-            docs["func_doc"] + "\nParameters\n----------\n" + docs["param_doc"],
+            result_docstring,
         )
         return result
 

metaflow/cmd/main_cli.py

@@ -67,6 +67,7 @@ CMDS_DESC = [
     ("configure", ".configure_cmd.cli"),
     ("tutorials", ".tutorials_cmd.cli"),
     ("develop", ".develop.cli"),
+    ("code", ".code.cli"),
 ]
 
 process_cmds(globals())

metaflow/cmd/make_wrapper.py

@@ -0,0 +1,16 @@
+import sys
+import subprocess
+from pathlib import Path
+import sysconfig
+
+
+def main():
+    share_dir = Path(sysconfig.get_paths()["data"]) / "share" / "metaflow" / "devtools"
+    makefile_path = share_dir / "Makefile"
+    cmd = ["make", "-f", str(makefile_path)] + sys.argv[1:]
+    # subprocess.run(cmd, check=True)
+    try:
+        completed = subprocess.run(cmd, check=True)
+        sys.exit(completed.returncode)
+    except subprocess.CalledProcessError as ex:
+        sys.exit(ex.returncode)

metaflow/metaflow_config.py

@@ -381,6 +381,8 @@ KUBERNETES_DISK = from_conf("KUBERNETES_DISK", None)
 # Default kubernetes QoS class
 KUBERNETES_QOS = from_conf("KUBERNETES_QOS", "burstable")
 
+# Architecture of kubernetes nodes - used for @conda/@pypi in metaflow-dev
+KUBERNETES_CONDA_ARCH = from_conf("KUBERNETES_CONDA_ARCH")
 ARGO_WORKFLOWS_KUBERNETES_SECRETS = from_conf("ARGO_WORKFLOWS_KUBERNETES_SECRETS", "")
 ARGO_WORKFLOWS_ENV_VARS_TO_SKIP = from_conf("ARGO_WORKFLOWS_ENV_VARS_TO_SKIP", "")
 

metaflow/metaflow_environment.py

@@ -6,7 +6,7 @@ from .util import get_username
 from . import metaflow_version
 from metaflow.exception import MetaflowException
 from metaflow.extension_support import dump_module_info
-from metaflow.mflog import BASH_MFLOG
+from metaflow.mflog import BASH_MFLOG, BASH_FLUSH_LOGS
 from . import R
 
 
@@ -159,6 +159,7 @@ class MetaflowEnvironment(object):
     def get_package_commands(self, code_package_url, datastore_type):
         cmds = [
             BASH_MFLOG,
+            BASH_FLUSH_LOGS,
             "mflog 'Setting up task environment.'",
             self._get_install_dependencies_cmd(datastore_type),
             "mkdir metaflow",
@@ -176,6 +177,7 @@ class MetaflowEnvironment(object):
             "fi" % code_package_url,
             "TAR_OPTIONS='--warning=no-timestamp' tar xf job.tar",
             "mflog 'Task is starting.'",
+            "flush_mflogs",
         ]
         return cmds
 

metaflow/mflog/__init__.py

@@ -44,6 +44,8 @@ BASH_MFLOG = (
 BASH_SAVE_LOGS_ARGS = ["python", "-m", "metaflow.mflog.save_logs"]
 BASH_SAVE_LOGS = " ".join(BASH_SAVE_LOGS_ARGS)
 
+BASH_FLUSH_LOGS = "flush_mflogs(){ " f"{BASH_SAVE_LOGS}; " "}"
+
 
 # this function returns a bash expression that redirects stdout
 # and stderr of the given bash expression to mflog.tee
@@ -63,7 +65,7 @@ def bash_capture_logs(bash_expr, var_transform=None):
 # update_delay determines how often logs should be uploaded to S3
 # as a function of the task execution time
 
-MIN_UPDATE_DELAY = 1.0  # the most frequent update interval
+MIN_UPDATE_DELAY = 0.25  # the most frequent update interval
 MAX_UPDATE_DELAY = 30.0  # the least frequent update interval
 
 
@@ -110,7 +112,6 @@ def export_mflog_env_vars(
 
 def tail_logs(prefix, stdout_tail, stderr_tail, echo, has_log_updates):
     def _available_logs(tail, stream, echo, should_persist=False):
-        # print the latest batch of lines
         try:
             for line in tail:
                 if should_persist:
@@ -128,7 +129,7 @@ def tail_logs(prefix, stdout_tail, stderr_tail, echo, has_log_updates):
 
     start_time = time.time()
     next_log_update = start_time
-    log_update_delay = 1
+    log_update_delay = update_delay(0)
     while has_log_updates():
         if time.time() > next_log_update:
             _available_logs(stdout_tail, "stdout", echo)

metaflow/plugins/aws/batch/batch_cli.py

@@ -211,7 +211,7 @@ def step(
     log_driver=None,
     log_options=None,
     num_parallel=None,
-    **kwargs
+    **kwargs,
 ):
     def echo(msg, stream="stderr", batch_id=None, **kwargs):
         msg = util.to_unicode(msg)
@@ -273,11 +273,11 @@ def step(
         "metaflow_version"
     ]
 
+    env = {"METAFLOW_FLOW_FILENAME": os.path.basename(sys.argv[0])}
+
     env_deco = [deco for deco in node.decorators if deco.name == "environment"]
     if env_deco:
-        env = env_deco[0].attributes["vars"]
-    else:
-        env = {}
+        env.update(env_deco[0].attributes["vars"])
 
     # Add the environment variables related to the input-paths argument
     if split_vars:

metaflow/plugins/aws/batch/batch_decorator.py

@@ -298,6 +298,13 @@ class BatchDecorator(StepDecorator):
             self._save_logs_sidecar = Sidecar("save_logs_periodically")
             self._save_logs_sidecar.start()
 
+            # Start spot termination monitor sidecar.
+            current._update_env(
+                {"spot_termination_notice": "/tmp/spot_termination_notice"}
+            )
+            self._spot_monitor_sidecar = Sidecar("spot_termination_monitor")
+            self._spot_monitor_sidecar.start()
+
         num_parallel = int(os.environ.get("AWS_BATCH_JOB_NUM_NODES", 0))
         if num_parallel >= 1 and ubf_context == UBF_CONTROL:
             # UBF handling for multinode case
@@ -350,6 +357,7 @@ class BatchDecorator(StepDecorator):
 
         try:
             self._save_logs_sidecar.terminate()
+            self._spot_monitor_sidecar.terminate()
         except:
             # Best effort kill
             pass

metaflow/plugins/kubernetes/kubernetes_decorator.py

@@ -29,6 +29,7 @@ from metaflow.metaflow_config import (
     KUBERNETES_SHARED_MEMORY,
     KUBERNETES_TOLERATIONS,
     KUBERNETES_QOS,
+    KUBERNETES_CONDA_ARCH,
 )
 from metaflow.plugins.resources_decorator import ResourcesDecorator
 from metaflow.plugins.timeout_decorator import get_run_time_limit_for_task
@@ -158,7 +159,7 @@ class KubernetesDecorator(StepDecorator):
 
     # Conda environment support
     supports_conda_environment = True
-    target_platform = "linux-64"
+    target_platform = KUBERNETES_CONDA_ARCH or "linux-64"
 
     def init(self):
         super(KubernetesDecorator, self).init()

metaflow/plugins/kubernetes/kubernetes_jobsets.py

@@ -319,6 +319,8 @@ class RunningJobSet(object):
     def kill(self):
         plural = "jobsets"
         client = self._client.get()
+        if not (self.is_running or self.is_waiting):
+            return
         try:
             # Killing the control pod will trigger the jobset to mark everything as failed.
             # Since jobsets have a successPolicy set to `All` which ensures that everything has

metaflow/plugins/pypi/bootstrap.py

@@ -8,6 +8,7 @@ import subprocess
 import sys
 import tarfile
 import time
+import platform
 from urllib.error import URLError
 from urllib.request import urlopen
 from metaflow.metaflow_config import DATASTORE_LOCAL_DIR, CONDA_USE_FAST_INIT
@@ -19,7 +20,7 @@ import warnings
 
 from . import MAGIC_FILE, _datastore_packageroot
 
-FAST_INIT_BIN_URL = "https://fast-flow-init.outerbounds.sh/{platform}/fast-env-0.1.1.gz"
+FAST_INIT_BIN_URL = "https://fast-flow-init.outerbounds.sh/{platform}/latest"
 
 # Bootstraps a valid conda virtual environment composed of conda and pypi packages
 
@@ -36,29 +37,6 @@ def timer(func):
 
 
 if __name__ == "__main__":
-    # TODO: Detect architecture on the fly when dealing with arm architectures.
-    # ARCH=$(uname -m)
-    # OS=$(uname)
-
-    # if [[ "$OS" == "Linux" ]]; then
-    #     PLATFORM="linux"
-    #     if [[ "$ARCH" == "aarch64" ]]; then
-    #         ARCH="aarch64";
-    #     elif [[ $ARCH == "ppc64le" ]]; then
-    #         ARCH="ppc64le";
-    #     else
-    #         ARCH="64";
-    #     fi
-    # fi
-
-    # if [[ "$OS" == "Darwin" ]]; then
-    #     PLATFORM="osx";
-    #     if [[ "$ARCH" == "arm64" ]]; then
-    #         ARCH="arm64";
-    #     else
-    #         ARCH="64"
-    #     fi
-    # fi
 
     def run_cmd(cmd, stdin_str=None):
         result = subprocess.run(
@@ -350,12 +328,25 @@ if __name__ == "__main__":
         cmd = f"fast-initializer --prefix {prefix} --packages-dir {pkgs_dir}"
         run_cmd(cmd, all_package_urls)
 
-    if len(sys.argv) != 5:
-        print("Usage: bootstrap.py <flow_name> <id> <datastore_type> <architecture>")
+    if len(sys.argv) != 4:
+        print("Usage: bootstrap.py <flow_name> <id> <datastore_type>")
         sys.exit(1)
 
     try:
-        _, flow_name, id_, datastore_type, architecture = sys.argv
+        _, flow_name, id_, datastore_type = sys.argv
+
+        system = platform.system().lower()
+        arch_machine = platform.machine().lower()
+
+        if system == "darwin" and arch_machine == "arm64":
+            architecture = "osx-arm64"
+        elif system == "darwin":
+            architecture = "osx-64"
+        elif system == "linux" and arch_machine == "aarch64":
+            architecture = "linux-aarch64"
+        else:
+            # default fallback
+            architecture = "linux-64"
 
         prefix = os.path.join(os.getcwd(), architecture, id_)
         pkgs_dir = os.path.join(os.getcwd(), ".pkgs")

metaflow/plugins/pypi/conda_environment.py

@@ -190,7 +190,6 @@ class CondaEnvironment(MetaflowEnvironment):
         #  4. Start PyPI solves in parallel after each conda environment is created
         #  5. Download PyPI packages sequentially
         #  6. Create and cache PyPI environments in parallel
-
         with ThreadPoolExecutor() as executor:
             # Start all conda solves in parallel
             conda_futures = [
@@ -213,14 +212,14 @@ class CondaEnvironment(MetaflowEnvironment):
 
                 # Queue PyPI solve to start after conda create
                 if result[0] in pypi_envs:
+                    # solve pypi envs uniquely
+                    pypi_env = pypi_envs.pop(result[0])
 
                     def pypi_solve(env):
                         create_future.result()  # Wait for conda create
                         return solve(*env, "pypi")
 
-                    pypi_futures.append(
-                        executor.submit(pypi_solve, pypi_envs[result[0]])
-                    )
+                    pypi_futures.append(executor.submit(pypi_solve, pypi_env))
 
             # Process PyPI results sequentially for downloads
             for solve_future in pypi_futures:
@@ -242,7 +241,7 @@ class CondaEnvironment(MetaflowEnvironment):
         if id_:
             # bootstrap.py is responsible for ensuring the validity of this executable.
             # -s is important! Can otherwise leak packages to other environments.
-            return os.path.join("linux-64", id_, "bin/python -s")
+            return os.path.join("$MF_ARCH", id_, "bin/python -s")
         else:
             # for @conda/@pypi(disabled=True).
             return super().executable(step_name, default)
@@ -315,7 +314,6 @@ class CondaEnvironment(MetaflowEnvironment):
         # 5. All resolved packages (Conda or PyPI) are cached
         # 6. PyPI packages are only installed for local platform
 
-        # Resolve `linux-64` Conda environments if @batch or @kubernetes are in play
         target_platform = conda_platform()
         for decorator in step.decorators:
             # NOTE: Keep the list of supported decorator names for backward compatibility purposes.
@@ -329,7 +327,6 @@ class CondaEnvironment(MetaflowEnvironment):
                 "snowpark",
                 "slurm",
             ]:
-                # TODO: Support arm architectures
                 target_platform = getattr(decorator, "target_platform", "linux-64")
                 break
 
@@ -424,15 +421,18 @@ class CondaEnvironment(MetaflowEnvironment):
         if id_:
             return [
                 "echo 'Bootstrapping virtual environment...'",
+                "flush_mflogs",
                 # We have to prevent the tracing module from loading,
                 # as the bootstrapping process uses the internal S3 client which would fail to import tracing
                 # due to the required dependencies being bundled into the conda environment,
                 # which is yet to be initialized at this point.
-                'DISABLE_TRACING=True python -m metaflow.plugins.pypi.bootstrap "%s" %s "%s" linux-64'
+                'DISABLE_TRACING=True python -m metaflow.plugins.pypi.bootstrap "%s" %s "%s"'
                 % (self.flow.name, id_, self.datastore_type),
                 "echo 'Environment bootstrapped.'",
+                "flush_mflogs",
                 # To avoid having to install micromamba in the PATH in micromamba.py, we add it to the PATH here.
                 "export PATH=$PATH:$(pwd)/micromamba/bin",
+                "export MF_ARCH=$(case $(uname)/$(uname -m) in Darwin/arm64)echo osx-arm64;;Darwin/*)echo osx-64;;Linux/aarch64)echo linux-aarch64;;*)echo linux-64;;esac)",
             ]
         else:
             # for @conda/@pypi(disabled=True).