metaflow 2.12.18__py2.py3-none-any.whl → 2.12.20__py2.py3-none-any.whl

metaflow/metaflow_current.py

@@ -4,7 +4,9 @@ from typing import Any, Optional, TYPE_CHECKING
 
 from metaflow.metaflow_config import TEMPDIR
 
-Parallel = namedtuple("Parallel", ["main_ip", "num_nodes", "node_index"])
+Parallel = namedtuple(
+    "Parallel", ["main_ip", "num_nodes", "node_index", "control_task_id"]
+)
 
 if TYPE_CHECKING:
     import metaflow

metaflow/plugins/argo/argo_workflows.py

@@ -1905,6 +1905,12 @@ class ArgoWorkflows(object):
                 jobset.environment_variable(
                     "MF_WORLD_SIZE", "{{inputs.parameters.num-parallel}}"
                 )
+                # We need this task-id set so that all the nodes are aware of the control
+                # task's task-id. These "MF_" variables populate the `current.parallel` namedtuple
+                jobset.environment_variable(
+                    "MF_PARALLEL_CONTROL_TASK_ID",
+                    "control-{{inputs.parameters.task-id-entropy}}-0",
+                )
                 # for k, v in .items():
                 jobset.environment_variables_from_selectors(
                     {
@@ -2552,8 +2558,8 @@ class ArgoWorkflows(object):
         cmd_str = " && ".join([init_cmds, heartbeat_cmds])
         cmds = shlex.split('bash -c "%s"' % cmd_str)
 
-        # TODO: Check that this is the minimal env.
         # Env required for sending heartbeats to the metadata service, nothing extra.
+        # prod token / runtime info is required to correctly register flow branches
         env = {
             # These values are needed by Metaflow to set it's internal
             # state appropriately.
@@ -2565,7 +2571,10 @@ class ArgoWorkflows(object):
             "METAFLOW_USER": "argo-workflows",
             "METAFLOW_DEFAULT_DATASTORE": self.flow_datastore.TYPE,
             "METAFLOW_DEFAULT_METADATA": DEFAULT_METADATA,
+            "METAFLOW_KUBERNETES_WORKLOAD": 1,
+            "METAFLOW_RUNTIME_ENVIRONMENT": "kubernetes",
             "METAFLOW_OWNER": self.username,
+            "METAFLOW_PRODUCTION_TOKEN": self.production_token,
         }
         # support Metaflow sandboxes
         env["METAFLOW_INIT_SCRIPT"] = KUBERNETES_SANDBOX_INIT_SCRIPT

metaflow/plugins/aws/batch/batch_client.py

@@ -89,6 +89,9 @@ class BatchJob(object):
         # Multinode
         if getattr(self, "num_parallel", 0) >= 1:
             num_nodes = self.num_parallel
+            # We need this task-id set so that all the nodes are aware of the control
+            # task's task-id. These "MF_" variables populate the `current.parallel` namedtuple
+            self.environment_variable("MF_PARALLEL_CONTROL_TASK_ID", self._task_id)
             main_task_override = copy.deepcopy(self.payload["containerOverrides"])
 
             # main

metaflow/plugins/kubernetes/kube_utils.py

@@ -0,0 +1,25 @@
+from metaflow.exception import CommandException
+from metaflow.util import get_username, get_latest_run_id
+
+
+def parse_cli_options(flow_name, run_id, user, my_runs, echo):
+    if user and my_runs:
+        raise CommandException("--user and --my-runs are mutually exclusive.")
+
+    if run_id and my_runs:
+        raise CommandException("--run_id and --my-runs are mutually exclusive.")
+
+    if my_runs:
+        user = get_username()
+
+    latest_run = True
+
+    if user and not run_id:
+        latest_run = False
+
+    if not run_id and latest_run:
+        run_id = get_latest_run_id(echo, flow_name)
+        if run_id is None:
+            raise CommandException("A previous run id was not found. Specify --run-id.")
+
+    return flow_name, run_id, user

metaflow/plugins/kubernetes/kubernetes.py

@@ -401,6 +401,9 @@ class Kubernetes(object):
             .label("app.kubernetes.io/name", "metaflow-task")
             .label("app.kubernetes.io/part-of", "metaflow")
         )
+        # We need this task-id set so that all the nodes are aware of the control
+        # task's task-id. These "MF_" variables populate the `current.parallel` namedtuple
+        jobset.environment_variable("MF_PARALLEL_CONTROL_TASK_ID", str(task_id))
 
         ## ----------- control/worker specific values START here -----------
         # We will now set the appropriate command for the control/worker job
@@ -698,7 +701,7 @@ class Kubernetes(object):
                     t = time.time()
                 time.sleep(update_delay(time.time() - start_time))
 
-        _make_prefix = lambda: b"[%s] " % util.to_bytes(self._job.id)
+        prefix = lambda: b"[%s] " % util.to_bytes(self._job.id)
 
         stdout_tail = get_log_tailer(stdout_location, self._datastore.TYPE)
         stderr_tail = get_log_tailer(stderr_location, self._datastore.TYPE)
@@ -708,7 +711,7 @@ class Kubernetes(object):
 
         # 2) Tail logs until the job has finished
         tail_logs(
-            prefix=_make_prefix(),
+            prefix=prefix(),
             stdout_tail=stdout_tail,
             stderr_tail=stderr_tail,
             echo=echo,

metaflow/plugins/kubernetes/kubernetes_cli.py

@@ -3,10 +3,12 @@ import sys
 import time
 import traceback
 
+from metaflow.plugins.kubernetes.kube_utils import parse_cli_options
+from metaflow.plugins.kubernetes.kubernetes_client import KubernetesClient
 import metaflow.tracing as tracing
 from metaflow import JSONTypeClass, util
 from metaflow._vendor import click
-from metaflow.exception import METAFLOW_EXIT_DISALLOW_RETRY, CommandException
+from metaflow.exception import METAFLOW_EXIT_DISALLOW_RETRY, MetaflowException
 from metaflow.metadata.util import sync_local_metadata_from_datastore
 from metaflow.metaflow_config import DATASTORE_LOCAL_DIR, KUBERNETES_LABELS
 from metaflow.mflog import TASK_LOG_SOURCE
@@ -305,3 +307,84 @@ def step(
         sys.exit(METAFLOW_EXIT_DISALLOW_RETRY)
     finally:
         _sync_metadata()
+
+
+@kubernetes.command(help="List unfinished Kubernetes tasks of this flow.")
+@click.option(
+    "--my-runs",
+    default=False,
+    is_flag=True,
+    help="List all my unfinished tasks.",
+)
+@click.option("--user", default=None, help="List unfinished tasks for the given user.")
+@click.option(
+    "--run-id",
+    default=None,
+    help="List unfinished tasks corresponding to the run id.",
+)
+@click.pass_obj
+def list(obj, run_id, user, my_runs):
+    flow_name, run_id, user = parse_cli_options(
+        obj.flow.name, run_id, user, my_runs, obj.echo
+    )
+    kube_client = KubernetesClient()
+    pods = kube_client.list(obj.flow.name, run_id, user)
+
+    def format_timestamp(timestamp=None):
+        if timestamp is None:
+            return "-"
+        return timestamp.strftime("%Y-%m-%d %H:%M:%S")
+
+    for pod in pods:
+        obj.echo(
+            "Run: *{run_id}* "
+            "Pod: *{pod_id}* "
+            "Started At: {startedAt} "
+            "Status: *{status}*".format(
+                run_id=pod.metadata.annotations.get(
+                    "metaflow/run_id",
+                    pod.metadata.labels.get("workflows.argoproj.io/workflow"),
+                ),
+                pod_id=pod.metadata.name,
+                startedAt=format_timestamp(pod.status.start_time),
+                status=pod.status.phase,
+            )
+        )
+
+    if not pods:
+        obj.echo("No active Kubernetes pods found.")
+
+
+@kubernetes.command(
+    help="Terminate unfinished Kubernetes tasks of this flow. Killed pods may result in newer attempts when using @retry."
+)
+@click.option(
+    "--my-runs",
+    default=False,
+    is_flag=True,
+    help="Kill all my unfinished tasks.",
+)
+@click.option(
+    "--user",
+    default=None,
+    help="Terminate unfinished tasks for the given user.",
+)
+@click.option(
+    "--run-id",
+    default=None,
+    help="Terminate unfinished tasks corresponding to the run id.",
+)
+@click.pass_obj
+def kill(obj, run_id, user, my_runs):
+    flow_name, run_id, user = parse_cli_options(
+        obj.flow.name, run_id, user, my_runs, obj.echo
+    )
+
+    if run_id is not None and run_id.startswith("argo-") or user == "argo-workflows":
+        raise MetaflowException(
+            "Killing pods launched by Argo Workflows is not supported. "
+            "Use *argo-workflows terminate* instead."
+        )
+
+    kube_client = KubernetesClient()
+    kube_client.kill_pods(flow_name, run_id, user, obj.echo)

metaflow/plugins/kubernetes/kubernetes_client.py

@@ -1,8 +1,10 @@
+from concurrent.futures import ThreadPoolExecutor
 import os
 import sys
 import time
 
 from metaflow.exception import MetaflowException
+from metaflow.metaflow_config import KUBERNETES_NAMESPACE
 
 from .kubernetes_job import KubernetesJob, KubernetesJobSet
 
@@ -28,6 +30,7 @@ class KubernetesClient(object):
                 % sys.executable
             )
         self._refresh_client()
+        self._namespace = KUBERNETES_NAMESPACE
 
     def _refresh_client(self):
         from kubernetes import client, config
@@ -60,6 +63,100 @@ class KubernetesClient(object):
 
         return self._client
 
+    def _find_active_pods(self, flow_name, run_id=None, user=None):
+        def _request(_continue=None):
+            # handle paginated responses
+            return self._client.CoreV1Api().list_namespaced_pod(
+                namespace=self._namespace,
+                # limited selector support for K8S api. We want to cover multiple statuses: Running / Pending / Unknown
+                field_selector="status.phase!=Succeeded,status.phase!=Failed",
+                limit=1000,
+                _continue=_continue,
+            )
+
+        results = _request()
+
+        if run_id is not None:
+            # handle argo prefixes in run_id
+            run_id = run_id[run_id.startswith("argo-") and len("argo-") :]
+
+        while results.metadata._continue or results.items:
+            for pod in results.items:
+                match = (
+                    # arbitrary pods might have no annotations at all.
+                    pod.metadata.annotations
+                    and pod.metadata.labels
+                    and (
+                        run_id is None
+                        or (pod.metadata.annotations.get("metaflow/run_id") == run_id)
+                        # we want to also match pods launched by argo-workflows
+                        or (
+                            pod.metadata.labels.get("workflows.argoproj.io/workflow")
+                            == run_id
+                        )
+                    )
+                    and (
+                        user is None
+                        or pod.metadata.annotations.get("metaflow/user") == user
+                    )
+                    and (
+                        pod.metadata.annotations.get("metaflow/flow_name") == flow_name
+                    )
+                )
+                if match:
+                    yield pod
+            if not results.metadata._continue:
+                break
+            results = _request(results.metadata._continue)
+
+    def list(self, flow_name, run_id, user):
+        results = self._find_active_pods(flow_name, run_id, user)
+
+        return list(results)
+
+    def kill_pods(self, flow_name, run_id, user, echo):
+        from kubernetes.stream import stream
+
+        api_instance = self._client.CoreV1Api()
+        job_api = self._client.BatchV1Api()
+        pods = self._find_active_pods(flow_name, run_id, user)
+
+        def _kill_pod(pod):
+            echo("Killing Kubernetes pod %s\n" % pod.metadata.name)
+            try:
+                stream(
+                    api_instance.connect_get_namespaced_pod_exec,
+                    name=pod.metadata.name,
+                    namespace=pod.metadata.namespace,
+                    command=[
+                        "/bin/sh",
+                        "-c",
+                        "/sbin/killall5",
+                    ],
+                    stderr=True,
+                    stdin=False,
+                    stdout=True,
+                    tty=False,
+                )
+            except Exception:
+                # best effort kill for pod can fail.
+                try:
+                    job_name = pod.metadata.labels.get("job-name", None)
+                    if job_name is None:
+                        raise Exception("Could not determine job name")
+
+                    job_api.patch_namespaced_job(
+                        name=job_name,
+                        namespace=pod.metadata.namespace,
+                        field_manager="metaflow",
+                        body={"spec": {"parallelism": 0}},
+                    )
+                except Exception as e:
+                    echo("failed to kill pod %s - %s" % (pod.metadata.name, str(e)))
+
+        with ThreadPoolExecutor() as executor:
+            executor.map(_kill_pod, list(pods))
+
     def jobset(self, **kwargs):
         return KubernetesJobSet(self, **kwargs)
 

metaflow/plugins/kubernetes/kubernetes_decorator.py

@@ -70,6 +70,10 @@ class KubernetesDecorator(StepDecorator):
         Kubernetes secrets to use when launching pod in Kubernetes. These
         secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
         in Metaflow configuration.
+    node_selector: Union[Dict[str,str], str], optional, default None
+        Kubernetes node selector(s) to apply to the pod running the task.
+        Can be passed in as a comma separated string of values e.g. "kubernetes.io/os=linux,kubernetes.io/arch=amd64"
+        or as a dictionary {"kubernetes.io/os": "linux", "kubernetes.io/arch": "amd64"}
     namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
         Kubernetes namespace to use when launching pod in Kubernetes.
     gpu : int, optional, default None

metaflow/plugins/kubernetes/kubernetes_jobsets.py

@@ -571,10 +571,8 @@ class JobSetSpec(object):
                                 namespace=self._kwargs["namespace"],
                             ),
                             spec=client.V1PodSpec(
-                                ##  --- jobset require podspec deets start----
                                 subdomain=self._kwargs["subdomain"],
                                 set_hostname_as_fqdn=True,
-                                ##  --- jobset require podspec deets end ----
                                 # Timeout is set on the pod and not the job (important!)
                                 active_deadline_seconds=self._kwargs[
                                     "timeout_in_seconds"

metaflow/plugins/parallel_decorator.py

@@ -24,6 +24,8 @@ class ParallelDecorator(StepDecorator):
                     The total number of tasks created by @parallel
                 - node_index : int
                     The index of the current task in all the @parallel tasks.
+                - control_task_id : Optional[str]
+                    The task ID of the control task. Available to all tasks.
 
     is_parallel -> bool
         True if the current step is a @parallel step.
@@ -67,6 +69,7 @@ class ParallelDecorator(StepDecorator):
                     main_ip=os.environ.get("MF_PARALLEL_MAIN_IP", "127.0.0.1"),
                     num_nodes=int(os.environ.get("MF_PARALLEL_NUM_NODES", "1")),
                     node_index=int(os.environ.get("MF_PARALLEL_NODE_INDEX", "0")),
+                    control_task_id=os.environ.get("MF_PARALLEL_CONTROL_TASK_ID", None),
                 )
             ),
         )
@@ -177,6 +180,7 @@ def _local_multinode_control_task_step_func(
     num_parallel = foreach_iter.num_parallel
     os.environ["MF_PARALLEL_NUM_NODES"] = str(num_parallel)
     os.environ["MF_PARALLEL_MAIN_IP"] = "127.0.0.1"
+    os.environ["MF_PARALLEL_CONTROL_TASK_ID"] = str(current.task_id)
 
     run_id = current.run_id
     step_name = current.step_name

metaflow/plugins/pypi/bootstrap.py

@@ -103,6 +103,7 @@ if __name__ == "__main__":
         echo "@EXPLICIT" > "$tmpfile";
         ls -d {conda_pkgs_dir}/*/* >> "$tmpfile";
         export PATH=$PATH:$(pwd)/micromamba;
+        export CONDA_PKGS_DIRS=$(pwd)/micromamba/pkgs;
         micromamba create --yes --offline --no-deps --safety-checks=disabled --no-extra-safety-checks --prefix {prefix} --file "$tmpfile";
         rm "$tmpfile"''',
     ]
@@ -123,6 +124,7 @@ if __name__ == "__main__":
             [
                 f"""set -e;
                 export PATH=$PATH:$(pwd)/micromamba;
+                export CONDA_PKGS_DIRS=$(pwd)/micromamba/pkgs;
                 micromamba run --prefix {prefix} python -m pip --disable-pip-version-check install --root-user-action=ignore --no-compile {pypi_pkgs_dir}/*.whl --no-user"""
             ]
         )

metaflow/plugins/pypi/conda_decorator.py

@@ -50,6 +50,9 @@ class CondaStepDecorator(StepDecorator):
     # conda channels, users can specify channel::package as the package name.
 
     def __init__(self, attributes=None, statically_defined=False):
+        self._user_defined_attributes = (
+            attributes.copy() if attributes is not None else {}
+        )
         super(CondaStepDecorator, self).__init__(attributes, statically_defined)
 
         # Support legacy 'libraries=' attribute for the decorator.
@@ -59,6 +62,9 @@ class CondaStepDecorator(StepDecorator):
         }
         del self.attributes["libraries"]
 
+    def is_attribute_user_defined(self, name):
+        return name in self._user_defined_attributes
+
     def step_init(self, flow, graph, step, decos, environment, flow_datastore, logger):
         # The init_environment hook for Environment creates the relevant virtual
         # environments. The step_init hook sets up the relevant state for that hook to
@@ -71,11 +77,16 @@ class CondaStepDecorator(StepDecorator):
 
         # Support flow-level decorator.
         if "conda_base" in self.flow._flow_decorators:
-            super_attributes = self.flow._flow_decorators["conda_base"][0].attributes
+            conda_base = self.flow._flow_decorators["conda_base"][0]
+            super_attributes = conda_base.attributes
             self.attributes["packages"] = {
                 **super_attributes["packages"],
                 **self.attributes["packages"],
             }
+            self._user_defined_attributes = {
+                **self._user_defined_attributes,
+                **conda_base._user_defined_attributes,
+            }
             self.attributes["python"] = (
                 self.attributes["python"] or super_attributes["python"]
             )
@@ -322,6 +333,9 @@ class CondaFlowDecorator(FlowDecorator):
     }
 
     def __init__(self, attributes=None, statically_defined=False):
+        self._user_defined_attributes = (
+            attributes.copy() if attributes is not None else {}
+        )
         super(CondaFlowDecorator, self).__init__(attributes, statically_defined)
 
         # Support legacy 'libraries=' attribute for the decorator.
@@ -333,9 +347,18 @@ class CondaFlowDecorator(FlowDecorator):
         if self.attributes["python"]:
             self.attributes["python"] = str(self.attributes["python"])
 
+    def is_attribute_user_defined(self, name):
+        return name in self._user_defined_attributes
+
     def flow_init(
         self, flow, graph, environment, flow_datastore, metadata, logger, echo, options
     ):
+        # NOTE: Important for extensions implementing custom virtual environments.
+        # Without this steps will not have an implicit conda step decorator on them unless the environment adds one in its decospecs.
+        from metaflow import decorators
+
+        decorators._attach_decorators(flow, ["conda"])
+
         # @conda uses a conda environment to create a virtual environment.
         # The conda environment can be created through micromamba.
         _supported_virtual_envs = ["conda"]

metaflow/plugins/pypi/pypi_decorator.py

@@ -40,7 +40,12 @@ class PyPIStepDecorator(StepDecorator):
 
         # Support flow-level decorator
         if "pypi_base" in self.flow._flow_decorators:
-            super_attributes = self.flow._flow_decorators["pypi_base"][0].attributes
+            pypi_base = self.flow._flow_decorators["pypi_base"][0]
+            super_attributes = pypi_base.attributes
+            self._user_defined_attributes = {
+                **self._user_defined_attributes,
+                **pypi_base._user_defined_attributes,
+            }
             self.attributes["packages"] = {
                 **super_attributes["packages"],
                 **self.attributes["packages"],
@@ -93,6 +98,12 @@ class PyPIStepDecorator(StepDecorator):
                     ),
                 )
             )
+        # TODO: This code snippet can be done away with by altering the constructor of
+        #       MetaflowEnvironment. A good first-task exercise.
+        # Avoid circular import
+        from metaflow.plugins.datastores.local_storage import LocalStorage
+
+        environment.set_local_root(LocalStorage.get_datastore_root_from_config(logger))
 
     def is_attribute_user_defined(self, name):
         return name in self._user_defined_attributes
@@ -117,6 +128,12 @@ class PyPIFlowDecorator(FlowDecorator):
     name = "pypi_base"
     defaults = {"packages": {}, "python": None, "disabled": None}
 
+    def __init__(self, attributes=None, statically_defined=False):
+        self._user_defined_attributes = (
+            attributes.copy() if attributes is not None else {}
+        )
+        super().__init__(attributes, statically_defined)
+
     def flow_init(
         self, flow, graph, environment, flow_datastore, metadata, logger, echo, options
     ):

metaflow/runner/click_api.py

@@ -1,3 +1,4 @@
+import os
 import sys
 
 if sys.version_info < (3, 7):
@@ -12,6 +13,7 @@ import importlib
 import inspect
 import itertools
 import uuid
+import json
 from collections import OrderedDict
 from typing import Any, Callable, Dict, List, Optional
 from typing import OrderedDict as TOrderedDict
@@ -37,6 +39,9 @@ from metaflow.exception import MetaflowException
 from metaflow.includefile import FilePathClass
 from metaflow.parameters import JSONTypeClass, flow_context
 
+# Define a recursive type alias for JSON
+JSON = Union[Dict[str, "JSON"], List["JSON"], str, int, float, bool, None]
+
 click_to_python_types = {
     StringParamType: str,
     IntParamType: int,
@@ -48,7 +53,7 @@ click_to_python_types = {
     Tuple: tuple,
     Choice: str,
     File: str,
-    JSONTypeClass: str,
+    JSONTypeClass: JSON,
     FilePathClass: str,
 }
 
@@ -82,6 +87,11 @@ def _method_sanity_check(
                 % (supplied_k, annotations[supplied_k], defaults[supplied_k])
             )
 
+        # because Click expects stringified JSON..
+        supplied_v = (
+            json.dumps(supplied_v) if annotations[supplied_k] == JSON else supplied_v
+        )
+
         if supplied_k in possible_arg_params:
             cli_name = possible_arg_params[supplied_k].opts[0].strip("-")
             method_params["args"][cli_name] = supplied_v
@@ -142,6 +152,8 @@ loaded_modules = {}
 
 
 def extract_flow_class_from_file(flow_file: str) -> FlowSpec:
+    if not os.path.exists(flow_file):
+        raise FileNotFoundError("Flow file not present at '%s'" % flow_file)
     # Check if the module has already been loaded
     if flow_file in loaded_modules:
         module = loaded_modules[flow_file]

metaflow/runner/deployer.py

@@ -5,6 +5,8 @@ import time
 import importlib
 import functools
 import tempfile
+
+from subprocess import CalledProcessError
 from typing import Optional, Dict, ClassVar
 
 from metaflow.exception import MetaflowNotFound
@@ -25,6 +27,8 @@ def handle_timeout(
         Temporary file that stores runner attribute data.
     command_obj : CommandManager
         Command manager object that encapsulates the running command details.
+    file_read_timeout : int
+        Timeout for reading the file.
 
     Returns
     -------
@@ -39,10 +43,10 @@ def handle_timeout(
     """
     try:
         content = read_from_file_when_ready(
-            tfp_runner_attribute.name, timeout=file_read_timeout
+            tfp_runner_attribute.name, command_obj, timeout=file_read_timeout
         )
         return content
-    except TimeoutError as e:
+    except (CalledProcessError, TimeoutError) as e:
         stdout_log = open(command_obj.log_files["stdout"]).read()
         stderr_log = open(command_obj.log_files["stderr"]).read()
         command = " ".join(command_obj.command)
@@ -397,6 +401,9 @@ class DeployerImpl(object):
             self.name = content.get("name")
             self.flow_name = content.get("flow_name")
             self.metadata = content.get("metadata")
+            # Additional info is used to pass additional deployer specific information.
+            # It is used in non-OSS deployers (extensions).
+            self.additional_info = content.get("additional_info", {})
 
             if command_obj.process.returncode == 0:
                 deployed_flow = DeployedFlow(deployer=self)

metaflow/runner/metaflow_runner.py

@@ -2,6 +2,8 @@ import importlib
 import os
 import sys
 import tempfile
+
+from subprocess import CalledProcessError
 from typing import Dict, Iterator, Optional, Tuple
 
 from metaflow import Run, metadata
@@ -275,13 +277,13 @@ class Runner(object):
 
             # Set the correct metadata from the runner_attribute file corresponding to this run.
             content = read_from_file_when_ready(
-                tfp_runner_attribute.name, timeout=self.file_read_timeout
+                tfp_runner_attribute.name, command_obj, timeout=self.file_read_timeout
             )
             metadata_for_flow, pathspec = content.rsplit(":", maxsplit=1)
             metadata(metadata_for_flow)
             run_object = Run(pathspec, _namespace_check=False)
             return ExecutingRun(self, command_obj, run_object)
-        except TimeoutError as e:
+        except (CalledProcessError, TimeoutError) as e:
             stdout_log = open(command_obj.log_files["stdout"]).read()
             stderr_log = open(command_obj.log_files["stderr"]).read()
             command = " ".join(command_obj.command)

metaflow/runner/subprocess_manager.py

@@ -1,18 +1,22 @@
 import asyncio
 import os
+import time
 import shutil
 import signal
 import subprocess
 import sys
 import tempfile
 import threading
-import time
 from typing import Callable, Dict, Iterator, List, Optional, Tuple
 
 
 def kill_process_and_descendants(pid, termination_timeout):
+    # TODO: there's a race condition that new descendants might
+    # spawn b/w the invocations of 'pkill' and 'kill'.
+    # Needs to be fixed in future.
     try:
         subprocess.check_call(["pkill", "-TERM", "-P", str(pid)])
+        subprocess.check_call(["kill", "-TERM", str(pid)])
     except subprocess.CalledProcessError:
         pass
 
@@ -20,6 +24,7 @@ def kill_process_and_descendants(pid, termination_timeout):
 
     try:
         subprocess.check_call(["pkill", "-KILL", "-P", str(pid)])
+        subprocess.check_call(["kill", "-KILL", str(pid)])
     except subprocess.CalledProcessError:
         pass
 
@@ -436,13 +441,13 @@ class CommandManager(object):
         if self.run_called:
             shutil.rmtree(self.temp_dir, ignore_errors=True)
 
-    async def kill(self, termination_timeout: float = 1):
+    async def kill(self, termination_timeout: float = 5):
         """
         Kill the subprocess and its descendants.
 
         Parameters
         ----------
-        termination_timeout : float, default 1
+        termination_timeout : float, default 5
             The time to wait after sending a SIGTERM to the process and its descendants
             before sending a SIGKILL.
         """

metaflow/runner/utils.py

@@ -1,7 +1,12 @@
 import os
 import ast
 import time
-from typing import Dict
+
+from subprocess import CalledProcessError
+from typing import Dict, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from .subprocess_manager import CommandManager
 
 
 def get_current_cell(ipython):
@@ -35,11 +40,23 @@ def clear_and_set_os_environ(env: Dict):
     os.environ.update(env)
 
 
-def read_from_file_when_ready(file_path: str, timeout: float = 5):
+def read_from_file_when_ready(
+    file_path: str, command_obj: "CommandManager", timeout: float = 5
+):
     start_time = time.time()
     with open(file_path, "r", encoding="utf-8") as file_pointer:
         content = file_pointer.read()
         while not content:
+            if command_obj.process.poll() is not None:
+                # Check to make sure the file hasn't been read yet to avoid a race
+                # where the file is written between the end of this while loop and the
+                # poll call above.
+                content = file_pointer.read()
+                if content:
+                    break
+                raise CalledProcessError(
+                    command_obj.process.returncode, command_obj.command
+                )
             if time.time() - start_time > timeout:
                 raise TimeoutError(
                     "Timeout while waiting for file content from '%s'" % file_path

metaflow/version.py

@@ -1 +1 @@
-metaflow_version = "2.12.18"
+metaflow_version = "2.12.20"

{metaflow-2.12.18.dist-info → metaflow-2.12.20.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: metaflow
-Version: 2.12.18
+Version: 2.12.20
 Summary: Metaflow: More Data Science, Less Engineering
 Author: Metaflow Developers
 Author-email: help@metaflow.org
@@ -26,7 +26,7 @@ License-File: LICENSE
 Requires-Dist: requests
 Requires-Dist: boto3
 Provides-Extra: stubs
-Requires-Dist: metaflow-stubs==2.12.18; extra == "stubs"
+Requires-Dist: metaflow-stubs==2.12.20; extra == "stubs"
 
 ![Metaflow_Logo_Horizontal_FullColor_Ribbon_Dark_RGB](https://user-images.githubusercontent.com/763451/89453116-96a57e00-d713-11ea-9fa6-82b29d4d6eff.png)
 

{metaflow-2.12.18.dist-info → metaflow-2.12.20.dist-info}/RECORD

@@ -17,7 +17,7 @@ metaflow/integrations.py,sha256=LlsaoePRg03DjENnmLxZDYto3NwWc9z_PtU6nJxLldg,1480
 metaflow/lint.py,sha256=5rj1MlpluxyPTSINjtMoJ7viotyNzfjtBJSAihlAwMU,10870
 metaflow/metaflow_config.py,sha256=4PUd2-JWJs35SaobDgMg4RdpZaKjhEqPUFh2f0pjqnU,22853
 metaflow/metaflow_config_funcs.py,sha256=5GlvoafV6SxykwfL8D12WXSfwjBN_NsyuKE_Q3gjGVE,6738
-metaflow/metaflow_current.py,sha256=5Kri7fzj-rtIJVr5xh5kPKwZ0T73_4egZybzlDR-fgc,7136
+metaflow/metaflow_current.py,sha256=pC-EMnAsnvBLvLd61W6MvfiCKcboryeui9f6r8z_sg8,7161
 metaflow/metaflow_environment.py,sha256=D7zHYqe8aLZVwJ20nx19vmsNW29Kf3PVE8hbBjVQin8,8115
 metaflow/metaflow_profile.py,sha256=jKPEW-hmAQO-htSxb9hXaeloLacAh41A35rMZH6G8pA,418
 metaflow/metaflow_version.py,sha256=mPQ6g_3XjNdi0NrxDzwlW8ZH0nMyYpwqmJ04P7TIdP0,4774
@@ -35,7 +35,7 @@ metaflow/tuple_util.py,sha256=_G5YIEhuugwJ_f6rrZoelMFak3DqAR2tt_5CapS1XTY,830
 metaflow/unbounded_foreach.py,sha256=p184WMbrMJ3xKYHwewj27ZhRUsSj_kw1jlye5gA9xJk,387
 metaflow/util.py,sha256=olAvJK3y1it_k99MhLulTaAJo7OFVt5rnrD-ulIFLCU,13616
 metaflow/vendor.py,sha256=FchtA9tH22JM-eEtJ2c9FpUdMn8sSb1VHuQS56EcdZk,5139
-metaflow/version.py,sha256=ENRBNUxFkqFdAswpvc-TIKbgracGkrn4SmFnZmnIlPI,29
+metaflow/version.py,sha256=kleNZLjxgidWFnFZyqIbSBKOHUU3wMG2-4RRKGiPktc,29
 metaflow/_vendor/__init__.py,sha256=y_CiwUD3l4eAKvTVDZeqgVujMy31cAM1qjAB-HfI-9s,353
 metaflow/_vendor/typing_extensions.py,sha256=0nUs5p1A_UrZigrAVBoOEM6TxU37zzPDUtiij1ZwpNc,110417
 metaflow/_vendor/zipp.py,sha256=ajztOH-9I7KA_4wqDYygtHa6xUBVZgFpmZ8FE74HHHI,8425
@@ -150,7 +150,7 @@ metaflow/plugins/environment_decorator.py,sha256=6m9j2B77d-Ja_l_9CTJ__0O6aB2a8Qt
 metaflow/plugins/events_decorator.py,sha256=c2GcH6Mspbey3wBkjM5lqxaNByFOzYDQdllLpXzRNv8,18283
 metaflow/plugins/logs_cli.py,sha256=77W5UNagU2mOKSMMvrQxQmBLRzvmjK-c8dWxd-Ygbqs,11410
 metaflow/plugins/package_cli.py,sha256=-J6D4cupHfWSZ4GEFo2yy9Je9oL3owRWm5pEJwaiqd4,1649
-metaflow/plugins/parallel_decorator.py,sha256=wXBmEnnOmWI46-Lcry_CArkujj6YOrfzAwVVw1kgpPI,8591
+metaflow/plugins/parallel_decorator.py,sha256=hRYTdw3iJjxBEws2rtrZaCHojOaDoO0Pu08ckl-Z5bU,8876
 metaflow/plugins/project_decorator.py,sha256=eJOe0Ea7CbUCReEhR_XQvRkhV6jyRqDxM72oZI7EMCk,5336
 metaflow/plugins/resources_decorator.py,sha256=3MMZ7uptDf99795_RcSOq4h0N3OFlKpd3ahIEsozBBs,1333
 metaflow/plugins/retry_decorator.py,sha256=tz_2Tq6GLg3vjDBZp0KKVTk3ADlCvqaWTSf7blmFdUw,1548
@@ -174,7 +174,7 @@ metaflow/plugins/airflow/sensors/s3_sensor.py,sha256=iDReG-7FKnumrtQg-HY6cCUAAqN
 metaflow/plugins/argo/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 metaflow/plugins/argo/argo_client.py,sha256=MKKhMCbWOPzf6z5zQQiyDRHHkAXcO7ipboDZDqAAvOk,15849
 metaflow/plugins/argo/argo_events.py,sha256=_C1KWztVqgi3zuH57pInaE9OzABc2NnncC-zdwOMZ-w,5909
-metaflow/plugins/argo/argo_workflows.py,sha256=6xUkz1LKdCLbl-O-D83Y2G5mCKYcIciKts3x1PNAzCk,170173
+metaflow/plugins/argo/argo_workflows.py,sha256=8YajGVnvBjGC22L308p8zMXzvfJP0fI9RsxGl46nO1k,170748
 metaflow/plugins/argo/argo_workflows_cli.py,sha256=X2j_F0xF8-K30ebM4dSLOTteDKXbr-jMN18oMpl5S6Y,36313
 metaflow/plugins/argo/argo_workflows_decorator.py,sha256=yprszMdbE3rBTcEA9VR0IEnPjTprUauZBc4SBb-Q7sA,7878
 metaflow/plugins/argo/argo_workflows_deployer.py,sha256=wSSZtThn_VPvE_Wu6NB1L0Q86LmBJh9g009v_lpvBPM,8125
@@ -188,7 +188,7 @@ metaflow/plugins/aws/aws_utils.py,sha256=dk92IRZ2QTF3PicBOtZMMOmS_FIncFqZPeL9EbC
 metaflow/plugins/aws/batch/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 metaflow/plugins/aws/batch/batch.py,sha256=e9ssahWM18GnipPK2sqYB-ztx9w7Eoo7YtWyEtufYxs,17787
 metaflow/plugins/aws/batch/batch_cli.py,sha256=6PTbyajRgdy0XmjyJLBTdKdiOB84dcovQQ8sFXlJqko,11749
-metaflow/plugins/aws/batch/batch_client.py,sha256=s9ZHhxQPPoBQijLUgn6_16QOaD4-22U_44uJbp-yLkI,28565
+metaflow/plugins/aws/batch/batch_client.py,sha256=ddlGG0Vk1mkO7tcvJjDvNAVsVLOlqddF7MA1kKfHSqM,28830
 metaflow/plugins/aws/batch/batch_decorator.py,sha256=kwgxEPCEoI6eZIpU5PuL442Ohg4_BfvwowoYgAnCzKE,17520
 metaflow/plugins/aws/secrets_manager/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 metaflow/plugins/aws/secrets_manager/aws_secrets_manager_secrets_provider.py,sha256=JtFUVu00Cg0FzAizgrPLXmrMqsT7YeQMkQlgeivUxcE,7986
@@ -279,35 +279,36 @@ metaflow/plugins/gcp/gs_tail.py,sha256=Jl_wvnzU7dub07A-DOAuP5FeccNIrPM-CeL1xKFs1
 metaflow/plugins/gcp/gs_utils.py,sha256=ZmIGFse1qYyvAVrwga23PQUzF6dXEDLLsZ2F-YRmvow,2030
 metaflow/plugins/gcp/includefile_support.py,sha256=vIDeR-MiJuUh-2S2pV7Z7FBkhIWwtHXaRrj76MWGRiY,3869
 metaflow/plugins/kubernetes/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-metaflow/plugins/kubernetes/kubernetes.py,sha256=cr3TheUasxIBEwFZ3GEVbctaf8gW57BM5BDk80ikjPI,31063
-metaflow/plugins/kubernetes/kubernetes_cli.py,sha256=qBDdr1Lvtt-RO9pB-9_HTOPdzAmDvvJ0aiQ1OoCcrMU,10892
-metaflow/plugins/kubernetes/kubernetes_client.py,sha256=GKg-gT3qhXMRQV-sG1YyoOf3Z32NXr_wwEN2ytMVSEg,2471
-metaflow/plugins/kubernetes/kubernetes_decorator.py,sha256=CXStYHomuJJK_Yocpdo6OJadEQv5hDfSpO7GPL61ltw,25322
+metaflow/plugins/kubernetes/kube_utils.py,sha256=fYDlvqi8jYPsWijDwT6Z2qhQswyFqv7tiwtic_I80Vg,749
+metaflow/plugins/kubernetes/kubernetes.py,sha256=bKBqgZXnIDkoa4xKtKoV6InPtYQy4CujfvcbQ3Pvsbc,31305
+metaflow/plugins/kubernetes/kubernetes_cli.py,sha256=sFZ9Zrjef85vCO0MGpUF-em8Pw3dePFb3hbX3PtAH4I,13463
+metaflow/plugins/kubernetes/kubernetes_client.py,sha256=LuKX9rejRySIajnfmDoNkQ6HrOMjwIMbHAOT7wgF9i8,6362
+metaflow/plugins/kubernetes/kubernetes_decorator.py,sha256=MDW7AJozpFavrFEZAzFlxbF0vIxU0KWRpbvAtjcRIi4,25671
 metaflow/plugins/kubernetes/kubernetes_job.py,sha256=Cfkee8LbXC17jSXWoeNdomQRvF_8YSeXNg1gvxm6E_M,31806
-metaflow/plugins/kubernetes/kubernetes_jobsets.py,sha256=OBmLtX-ZUDQdCCfftUmRMernfmTNMwdTxPoCAp_NmwE,40957
+metaflow/plugins/kubernetes/kubernetes_jobsets.py,sha256=sCVuy-PRN60s5HWpnRvFmnljS41k70sBoPO3nOtl2ZE,40800
 metaflow/plugins/metadata/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
 metaflow/plugins/metadata/local.py,sha256=YhLJC5zjVJrvQFIyQ92ZBByiUmhCC762RUX7ITX12O8,22428
 metaflow/plugins/metadata/service.py,sha256=ihq5F7KQZlxvYwzH_-jyP2aWN_I96i2vp92j_d697s8,20204
 metaflow/plugins/pypi/__init__.py,sha256=0YFZpXvX7HCkyBFglatual7XGifdA1RwC3U4kcizyak,1037
-metaflow/plugins/pypi/bootstrap.py,sha256=Tvc4_QKIx-A8j5Aq8ccWZrrxNM8csN40rK8HmxDx-Z8,5106
-metaflow/plugins/pypi/conda_decorator.py,sha256=fTJVbEfgOUtsDXIfnfsNk46sKeA9uTuTqGey9OFs9Ig,14738
+metaflow/plugins/pypi/bootstrap.py,sha256=LSRjpDhtpvYqCC0lDPIWvZsrHW7PXx626NDzdYjFeVM,5224
+metaflow/plugins/pypi/conda_decorator.py,sha256=sIAfvkjhUfTMK5yLrPOkqxeLUhlC5Buqmc9aITaf22E,15686
 metaflow/plugins/pypi/conda_environment.py,sha256=--q-8lypKupCdGsASpqABNpNqRxtQi6UCDgq8iHDFe4,19476
 metaflow/plugins/pypi/micromamba.py,sha256=67FiIZZz0Kig9EcN7bZLObsE6Z1MFyo4Dp93fd3Grcc,12178
 metaflow/plugins/pypi/pip.py,sha256=7B06mPOs5MvY33xbzPVYZlBr1iKMYaN-n8uulL9zSVg,13649
-metaflow/plugins/pypi/pypi_decorator.py,sha256=h5cAnxkWjmj4Ad4q0AkABKwhHQHYfeexy12yMaaLgXQ,6443
+metaflow/plugins/pypi/pypi_decorator.py,sha256=rDMbHl7r81Ye7-TuIlKAVJ_CDnfjl9jV44ZPws-UsTY,7229
 metaflow/plugins/pypi/pypi_environment.py,sha256=FYMg8kF3lXqcLfRYWD83a9zpVjcoo_TARqMGZ763rRk,230
 metaflow/plugins/pypi/utils.py,sha256=ds1Mnv_DaxGnLAYp7ozg_K6oyguGyNhvHfE-75Ia1YA,2836
 metaflow/plugins/secrets/__init__.py,sha256=mhJaN2eMS_ZZVewAMR2E-JdP5i0t3v9e6Dcwd-WpruE,310
 metaflow/plugins/secrets/inline_secrets_provider.py,sha256=EChmoBGA1i7qM3jtYwPpLZDBybXLergiDlN63E0u3x8,294
 metaflow/plugins/secrets/secrets_decorator.py,sha256=s-sFzPWOjahhpr5fMj-ZEaHkDYAPTO0isYXGvaUwlG8,11273
 metaflow/runner/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-metaflow/runner/click_api.py,sha256=vrrIb5DtVYhCW_hB7wPgj6Q0o-h1bBSWiYnKTOyC454,13452
-metaflow/runner/deployer.py,sha256=nArjnErc0rOaZW612VRKDOT5594jwzeu86w5zW1LX6U,13558
-metaflow/runner/metaflow_runner.py,sha256=AO9nwr5qUbZWmsbFdjkUJrvFlaylz7WvxslvHsIqDYc,15371
+metaflow/runner/click_api.py,sha256=Qfg4BOz5K2LaXTYBsi1y4zTfNIsGGHBVF3UkorX_-o8,13878
+metaflow/runner/deployer.py,sha256=ddYv2UFE7lOFD3EPg4REzMdeLqg5WQhAXkdgMD5bUY8,13920
+metaflow/runner/metaflow_runner.py,sha256=zRlnM2j5QLBh9OboPUuaQ4_WpwFaqrCedKBOpjjtxDw,15449
 metaflow/runner/nbdeploy.py,sha256=fP1s_5MeiDyT_igP82pB5EUqX9rOy2s06Hyc-OUbOvQ,4115
 metaflow/runner/nbrun.py,sha256=lmvhzMCz7iC9LSPGRijifW1wMXxa4RW_jVmpdjQi22E,7261
-metaflow/runner/subprocess_manager.py,sha256=0knxWZYJx8srMv6wTPYKOC6tn4-airnyI7Vbqfb3iXY,19567
-metaflow/runner/utils.py,sha256=FibdEj8CDnx1a-Je5KUQTwHuNbtkFm1unXGarj0D8ok,1394
+metaflow/runner/subprocess_manager.py,sha256=FJh2C6wJlNnvwE00wyK6zxjf8R49PRRHeVcnPvlO4-0,19839
+metaflow/runner/utils.py,sha256=jFlQf9fVvNRd4c6VKv0KwQ8kez_re7kC2VCRgHVMSsM,2051
 metaflow/sidecar/__init__.py,sha256=1mmNpmQ5puZCpRmmYlCOeieZ4108Su9XQ4_EqF1FGOU,131
 metaflow/sidecar/sidecar.py,sha256=EspKXvPPNiyRToaUZ51PS5TT_PzrBNAurn_wbFnmGr0,1334
 metaflow/sidecar/sidecar_messages.py,sha256=zPsCoYgDIcDkkvdC9MEpJTJ3y6TSGm2JWkRc4vxjbFA,1071
@@ -344,9 +345,9 @@ metaflow/tutorials/07-worldview/README.md,sha256=5vQTrFqulJ7rWN6r20dhot9lI2sVj9W
 metaflow/tutorials/07-worldview/worldview.ipynb,sha256=ztPZPI9BXxvW1QdS2Tfe7LBuVzvFvv0AToDnsDJhLdE,2237
 metaflow/tutorials/08-autopilot/README.md,sha256=GnePFp_q76jPs991lMUqfIIh5zSorIeWznyiUxzeUVE,1039
 metaflow/tutorials/08-autopilot/autopilot.ipynb,sha256=DQoJlILV7Mq9vfPBGW-QV_kNhWPjS5n6SJLqePjFYLY,3191
-metaflow-2.12.18.dist-info/LICENSE,sha256=nl_Lt5v9VvJ-5lWJDT4ddKAG-VZ-2IaLmbzpgYDz2hU,11343
-metaflow-2.12.18.dist-info/METADATA,sha256=YnoqO6tzhh5iHKOGK8jA5DzSl73EcXrjibG238O4I3M,5906
-metaflow-2.12.18.dist-info/WHEEL,sha256=WDDPHYzpiOIm6GP1C2_8y8W6q16ICddAgOHlhTje9Qc,109
-metaflow-2.12.18.dist-info/entry_points.txt,sha256=IKwTN1T3I5eJL3uo_vnkyxVffcgnRdFbKwlghZfn27k,57
-metaflow-2.12.18.dist-info/top_level.txt,sha256=v1pDHoWaSaKeuc5fKTRSfsXCKSdW1zvNVmvA-i0if3o,9
-metaflow-2.12.18.dist-info/RECORD,,
+metaflow-2.12.20.dist-info/LICENSE,sha256=nl_Lt5v9VvJ-5lWJDT4ddKAG-VZ-2IaLmbzpgYDz2hU,11343
+metaflow-2.12.20.dist-info/METADATA,sha256=oHu5racxyaljUQCl6cYXqYOHkI02VB5mR298txO9lAk,5906
+metaflow-2.12.20.dist-info/WHEEL,sha256=AHX6tWk3qWuce7vKLrj7lnulVHEdWoltgauo8bgCXgU,109
+metaflow-2.12.20.dist-info/entry_points.txt,sha256=IKwTN1T3I5eJL3uo_vnkyxVffcgnRdFbKwlghZfn27k,57
+metaflow-2.12.20.dist-info/top_level.txt,sha256=v1pDHoWaSaKeuc5fKTRSfsXCKSdW1zvNVmvA-i0if3o,9
+metaflow-2.12.20.dist-info/RECORD,,

{metaflow-2.12.18.dist-info → metaflow-2.12.20.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (74.0.0)
+Generator: setuptools (75.1.0)
 Root-Is-Purelib: true
 Tag: py2-none-any
 Tag: py3-none-any