metaflow 2.12.18__tar.gz → 2.12.20__tar.gz

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

@@ -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 → 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.18 → 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.18 → 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.18 → 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
@@ -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-2.12.18 → 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.18 → 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.18 → metaflow-2.12.20}/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-2.12.18 → 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.18 → 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.18 → metaflow-2.12.20}/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-2.12.18 → metaflow-2.12.20}/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-2.12.18 → 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.18 → 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)