metaflow 2.12.19__tar.gz → 2.12.20__tar.gz

{metaflow-2.12.19/metaflow.egg-info → metaflow-2.12.20}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: metaflow
-Version: 2.12.19
+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.19; 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.19 → metaflow-2.12.20}/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-2.12.19 → metaflow-2.12.20}/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-2.12.19 → metaflow-2.12.20}/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-2.12.20/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-2.12.19 → metaflow-2.12.20}/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

{metaflow-2.12.19 → metaflow-2.12.20}/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-2.12.20/metaflow/plugins/kubernetes/kubernetes_client.py

@@ -0,0 +1,164 @@
+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
+
+CLIENT_REFRESH_INTERVAL_SECONDS = 300
+
+
+class KubernetesClientException(MetaflowException):
+    headline = "Kubernetes client error"
+
+
+class KubernetesClient(object):
+    def __init__(self):
+        try:
+            # Kubernetes is a soft dependency.
+            from kubernetes import client, config
+        except (NameError, ImportError):
+            raise KubernetesClientException(
+                "Could not import module 'kubernetes'.\n\nInstall Kubernetes "
+                "Python package (https://pypi.org/project/kubernetes/) first.\n"
+                "You can install the module by executing - "
+                "%s -m pip install kubernetes\n"
+                "or equivalent through your favorite Python package manager."
+                % sys.executable
+            )
+        self._refresh_client()
+        self._namespace = KUBERNETES_NAMESPACE
+
+    def _refresh_client(self):
+        from kubernetes import client, config
+
+        if os.getenv("KUBECONFIG"):
+            # There are cases where we're running inside a pod, but can't use
+            # the kubernetes client for that pod's cluster: for example when
+            # running in Bitbucket Cloud or other CI system.
+            # In this scenario, the user can set a KUBECONFIG environment variable
+            # to load the kubeconfig, regardless of whether we're in a pod or not.
+            config.load_kube_config()
+        elif os.getenv("KUBERNETES_SERVICE_HOST"):
+            # We are inside a pod, authenticate via ServiceAccount assigned to us
+            config.load_incluster_config()
+        else:
+            # Default to using kubeconfig, likely $HOME/.kube/config
+            # TODO (savin):
+            #  1. Support generating kubeconfig on the fly using boto3
+            #  2. Support auth via OIDC - https://docs.aws.amazon.com/eks/latest/userguide/authenticate-oidc-identity-provider.html
+            config.load_kube_config()
+        self._client = client
+        self._client_refresh_timestamp = time.time()
+
+    def get(self):
+        if (
+            time.time() - self._client_refresh_timestamp
+            > CLIENT_REFRESH_INTERVAL_SECONDS
+        ):
+            self._refresh_client()
+
+        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)
+
+    def job(self, **kwargs):
+        return KubernetesJob(self, **kwargs)

{metaflow-2.12.19 → metaflow-2.12.20}/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-2.12.19 → metaflow-2.12.20}/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-2.12.19 → metaflow-2.12.20}/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-2.12.19 → metaflow-2.12.20}/metaflow/plugins/pypi/conda_decorator.py

@@ -353,6 +353,12 @@ class CondaFlowDecorator(FlowDecorator):
     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-2.12.19 → metaflow-2.12.20}/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-2.12.19 → metaflow-2.12.20}/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-2.12.19 → metaflow-2.12.20}/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-2.12.19 → metaflow-2.12.20}/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-2.12.19 → metaflow-2.12.20}/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-2.12.20/metaflow/version.py

@@ -0,0 +1 @@
+metaflow_version = "2.12.20"

{metaflow-2.12.19 → metaflow-2.12.20/metaflow.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: metaflow
-Version: 2.12.19
+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.19; 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)