metaflow 2.12.33__py2.py3-none-any.whl → 2.12.35__py2.py3-none-any.whl

metaflow/metaflow_config.py

@@ -377,6 +377,8 @@ KUBERNETES_PORT = from_conf("KUBERNETES_PORT", None)
 KUBERNETES_CPU = from_conf("KUBERNETES_CPU", None)
 KUBERNETES_MEMORY = from_conf("KUBERNETES_MEMORY", None)
 KUBERNETES_DISK = from_conf("KUBERNETES_DISK", None)
+# Default kubernetes QoS class
+KUBERNETES_QOS = from_conf("KUBERNETES_QOS", "burstable")
 
 ARGO_WORKFLOWS_KUBERNETES_SECRETS = from_conf("ARGO_WORKFLOWS_KUBERNETES_SECRETS", "")
 ARGO_WORKFLOWS_ENV_VARS_TO_SKIP = from_conf("ARGO_WORKFLOWS_ENV_VARS_TO_SKIP", "")

metaflow/multicore_utils.py

@@ -6,7 +6,18 @@ from tempfile import NamedTemporaryFile
 import time
 import metaflow.tracing as tracing
 
-from typing import Any, Callable, Iterable, Iterator, List, Optional
+from typing import (
+    Any,
+    Callable,
+    Iterable,
+    Iterator,
+    List,
+    Optional,
+    NoReturn,
+    Tuple,
+    TypeVar,
+    Union,
+)
 
 try:
     # Python 2
@@ -30,7 +41,13 @@ class MulticoreException(Exception):
     pass
 
 
-def _spawn(func, arg, dir):
+_A = TypeVar("_A")
+_R = TypeVar("_R")
+
+
+def _spawn(
+    func: Callable[[_A], _R], arg: _A, dir: Optional[str]
+) -> Union[Tuple[int, str], NoReturn]:
     with NamedTemporaryFile(prefix="parallel_map_", dir=dir, delete=False) as tmpfile:
         output_file = tmpfile.name
 
@@ -63,11 +80,11 @@ def _spawn(func, arg, dir):
 
 
 def parallel_imap_unordered(
-    func: Callable[[Any], Any],
-    iterable: Iterable[Any],
+    func: Callable[[_A], _R],
+    iterable: Iterable[_A],
     max_parallel: Optional[int] = None,
     dir: Optional[str] = None,
-) -> Iterator[Any]:
+) -> Iterator[_R]:
     """
     Parallelizes execution of a function using multiprocessing. The result
     order is not guaranteed.
@@ -79,9 +96,9 @@ def parallel_imap_unordered(
     iterable : Iterable[Any]
         Iterable over arguments to pass to fun
     max_parallel int, optional, default None
-        Maximum parallelism. If not specified, uses the number of CPUs
+        Maximum parallelism. If not specified, it uses the number of CPUs
     dir : str, optional, default None
-        If specified, directory where temporary files are created
+        If specified, it's the directory where temporary files are created
 
     Yields
     ------
@@ -121,14 +138,14 @@ def parallel_imap_unordered(
 
 
 def parallel_map(
-    func: Callable[[Any], Any],
-    iterable: Iterable[Any],
+    func: Callable[[_A], _R],
+    iterable: Iterable[_A],
     max_parallel: Optional[int] = None,
     dir: Optional[str] = None,
-) -> List[Any]:
+) -> List[_R]:
     """
     Parallelizes execution of a function using multiprocessing. The result
-    order is that of the arguments in `iterable`
+    order is that of the arguments in `iterable`.
 
     Parameters
     ----------
@@ -137,9 +154,9 @@ def parallel_map(
     iterable : Iterable[Any]
         Iterable over arguments to pass to fun
     max_parallel int, optional, default None
-        Maximum parallelism. If not specified, uses the number of CPUs
+        Maximum parallelism. If not specified, it uses the number of CPUs
     dir : str, optional, default None
-        If specified, directory where temporary files are created
+        If specified, it's the directory where temporary files are created
 
     Returns
     -------
@@ -155,4 +172,4 @@ def parallel_map(
     res = parallel_imap_unordered(
         wrapper, enumerate(iterable), max_parallel=max_parallel, dir=dir
     )
-    return [r for idx, r in sorted(res)]
+    return [r for _, r in sorted(res)]

metaflow/parameters.py

@@ -151,6 +151,7 @@ class DeployTimeField(object):
             return self._check_type(val, deploy_time)
 
     def _check_type(self, val, deploy_time):
+
         # it is easy to introduce a deploy-time function that accidentally
         # returns a value whose type is not compatible with what is defined
         # in Parameter. Let's catch those mistakes early here, instead of
@@ -158,7 +159,7 @@ class DeployTimeField(object):
 
         # note: this doesn't work with long in Python2 or types defined as
         # click types, e.g. click.INT
-        TYPES = {bool: "bool", int: "int", float: "float", list: "list"}
+        TYPES = {bool: "bool", int: "int", float: "float", list: "list", dict: "dict"}
 
         msg = (
             "The value returned by the deploy-time function for "
@@ -166,7 +167,12 @@ class DeployTimeField(object):
             % (self.parameter_name, self.field)
         )
 
-        if self.parameter_type in TYPES:
+        if isinstance(self.parameter_type, list):
+            if not any(isinstance(val, x) for x in self.parameter_type):
+                msg += "Expected one of the following %s." % TYPES[self.parameter_type]
+                raise ParameterFieldTypeMismatch(msg)
+            return str(val) if self.return_str else val
+        elif self.parameter_type in TYPES:
             if type(val) != self.parameter_type:
                 msg += "Expected a %s." % TYPES[self.parameter_type]
                 raise ParameterFieldTypeMismatch(msg)

metaflow/plugins/airflow/airflow.py

@@ -46,6 +46,7 @@ from metaflow.parameters import (
 # TODO: Move chevron to _vendor
 from metaflow.plugins.cards.card_modules import chevron
 from metaflow.plugins.kubernetes.kubernetes import Kubernetes
+from metaflow.plugins.kubernetes.kube_utils import qos_requests_and_limits
 from metaflow.plugins.timeout_decorator import get_run_time_limit_for_task
 from metaflow.util import compress_list, dict_to_cli_options, get_username
 
@@ -428,25 +429,25 @@ class Airflow(object):
             if k8s_deco.attributes["namespace"] is not None
             else "default"
         )
-
+        qos_requests, qos_limits = qos_requests_and_limits(
+            k8s_deco.attributes["qos"],
+            k8s_deco.attributes["cpu"],
+            k8s_deco.attributes["memory"],
+            k8s_deco.attributes["disk"],
+        )
         resources = dict(
-            requests={
-                "cpu": k8s_deco.attributes["cpu"],
-                "memory": "%sM" % str(k8s_deco.attributes["memory"]),
-                "ephemeral-storage": str(k8s_deco.attributes["disk"]),
-            }
+            requests=qos_requests,
+            limits={
+                **qos_limits,
+                **{
+                    "%s.com/gpu".lower()
+                    % k8s_deco.attributes["gpu_vendor"]: str(k8s_deco.attributes["gpu"])
+                    for k in [0]
+                    # Don't set GPU limits if gpu isn't specified.
+                    if k8s_deco.attributes["gpu"] is not None
+                },
+            },
         )
-        if k8s_deco.attributes["gpu"] is not None:
-            resources.update(
-                dict(
-                    limits={
-                        "%s.com/gpu".lower()
-                        % k8s_deco.attributes["gpu_vendor"]: str(
-                            k8s_deco.attributes["gpu"]
-                        )
-                    }
-                )
-            )
 
         annotations = {
             "metaflow/production_token": self.production_token,

metaflow/plugins/argo/argo_workflows.py

@@ -53,6 +53,7 @@ from metaflow.metaflow_config import (
 from metaflow.metaflow_config_funcs import config_values
 from metaflow.mflog import BASH_SAVE_LOGS, bash_capture_logs, export_mflog_env_vars
 from metaflow.parameters import deploy_time_eval
+from metaflow.plugins.kubernetes.kube_utils import qos_requests_and_limits
 from metaflow.plugins.kubernetes.kubernetes import (
     parse_kube_keyvalue_list,
     validate_kube_labels,
@@ -522,7 +523,9 @@ class ArgoWorkflows(object):
             params = set(
                 [param.name.lower() for var, param in self.flow._get_parameters()]
             )
-            for event in self.flow._flow_decorators.get("trigger")[0].triggers:
+            trigger_deco = self.flow._flow_decorators.get("trigger")[0]
+            trigger_deco.format_deploytime_value()
+            for event in trigger_deco.triggers:
                 parameters = {}
                 # TODO: Add a check to guard against names starting with numerals(?)
                 if not re.match(r"^[A-Za-z0-9_.-]+$", event["name"]):
@@ -562,9 +565,11 @@ class ArgoWorkflows(object):
 
         # @trigger_on_finish decorator
         if self.flow._flow_decorators.get("trigger_on_finish"):
-            for event in self.flow._flow_decorators.get("trigger_on_finish")[
-                0
-            ].triggers:
+            trigger_on_finish_deco = self.flow._flow_decorators.get(
+                "trigger_on_finish"
+            )[0]
+            trigger_on_finish_deco.format_deploytime_value()
+            for event in trigger_on_finish_deco.triggers:
                 # Actual filters are deduced here since we don't have access to
                 # the current object in the @trigger_on_finish decorator.
                 triggers.append(
@@ -1838,6 +1843,13 @@ class ArgoWorkflows(object):
             if tmpfs_enabled and tmpfs_tempdir:
                 env["METAFLOW_TEMPDIR"] = tmpfs_path
 
+            qos_requests, qos_limits = qos_requests_and_limits(
+                resources["qos"],
+                resources["cpu"],
+                resources["memory"],
+                resources["disk"],
+            )
+
             # Create a ContainerTemplate for this node. Ideally, we would have
             # liked to inline this ContainerTemplate and avoid scanning the workflow
             # twice, but due to issues with variable substitution, we will have to
@@ -1901,6 +1913,7 @@ class ArgoWorkflows(object):
                     persistent_volume_claims=resources["persistent_volume_claims"],
                     shared_memory=shared_memory,
                     port=port,
+                    qos=resources["qos"],
                 )
 
                 for k, v in env.items():
@@ -2086,17 +2099,17 @@ class ArgoWorkflows(object):
                                 image=resources["image"],
                                 image_pull_policy=resources["image_pull_policy"],
                                 resources=kubernetes_sdk.V1ResourceRequirements(
-                                    requests={
-                                        "cpu": str(resources["cpu"]),
-                                        "memory": "%sM" % str(resources["memory"]),
-                                        "ephemeral-storage": "%sM"
-                                        % str(resources["disk"]),
-                                    },
+                                    requests=qos_requests,
                                     limits={
-                                        "%s.com/gpu".lower()
-                                        % resources["gpu_vendor"]: str(resources["gpu"])
-                                        for k in [0]
-                                        if resources["gpu"] is not None
+                                        **qos_limits,
+                                        **{
+                                            "%s.com/gpu".lower()
+                                            % resources["gpu_vendor"]: str(
+                                                resources["gpu"]
+                                            )
+                                            for k in [0]
+                                            if resources["gpu"] is not None
+                                        },
                                     },
                                 ),
                                 # Configure secrets
@@ -2333,7 +2346,7 @@ class ArgoWorkflows(object):
                                 "memory": "500Mi",
                             },
                         ),
-                    )
+                    ).to_dict()
                 )
             ),
             Template("capture-error-hook-fn-preflight").steps(
@@ -2684,7 +2697,7 @@ class ArgoWorkflows(object):
                             },
                         ),
                     )
-                )
+                ).to_dict()
             )
         )
 
@@ -2854,7 +2867,7 @@ class ArgoWorkflows(object):
                                         "memory": "250Mi",
                                     },
                                 ),
-                            )
+                            ).to_dict()
                         )
                     )
                     .service_account_name(ARGO_EVENTS_SERVICE_ACCOUNT)

metaflow/plugins/argo/argo_workflows_deployer_objects.py

@@ -10,7 +10,7 @@ from metaflow.metaflow_config import KUBERNETES_NAMESPACE
 from metaflow.plugins.argo.argo_workflows import ArgoWorkflows
 from metaflow.runner.deployer import Deployer, DeployedFlow, TriggeredRun
 
-from metaflow.runner.utils import get_lower_level_group, handle_timeout
+from metaflow.runner.utils import get_lower_level_group, handle_timeout, temporary_fifo
 
 
 def generate_fake_flow_file_contents(
@@ -341,18 +341,14 @@ class ArgoWorkflowsDeployedFlow(DeployedFlow):
         Exception
             If there is an error during the trigger process.
         """
-        with tempfile.TemporaryDirectory() as temp_dir:
-            tfp_runner_attribute = tempfile.NamedTemporaryFile(
-                dir=temp_dir, delete=False
-            )
-
+        with temporary_fifo() as (attribute_file_path, attribute_file_fd):
             # every subclass needs to have `self.deployer_kwargs`
             command = get_lower_level_group(
                 self.deployer.api,
                 self.deployer.top_level_kwargs,
                 self.deployer.TYPE,
                 self.deployer.deployer_kwargs,
-            ).trigger(deployer_attribute_file=tfp_runner_attribute.name, **kwargs)
+            ).trigger(deployer_attribute_file=attribute_file_path, **kwargs)
 
             pid = self.deployer.spm.run_command(
                 [sys.executable, *command],
@@ -363,7 +359,7 @@ class ArgoWorkflowsDeployedFlow(DeployedFlow):
 
             command_obj = self.deployer.spm.get(pid)
             content = handle_timeout(
-                tfp_runner_attribute, command_obj, self.deployer.file_read_timeout
+                attribute_file_fd, command_obj, self.deployer.file_read_timeout
             )
 
             if command_obj.process.returncode == 0:

metaflow/plugins/aws/step_functions/step_functions_deployer_objects.py

@@ -6,7 +6,7 @@ from typing import ClassVar, Optional, List
 from metaflow.plugins.aws.step_functions.step_functions import StepFunctions
 from metaflow.runner.deployer import DeployedFlow, TriggeredRun
 
-from metaflow.runner.utils import get_lower_level_group, handle_timeout
+from metaflow.runner.utils import get_lower_level_group, handle_timeout, temporary_fifo
 
 
 class StepFunctionsTriggeredRun(TriggeredRun):
@@ -196,18 +196,14 @@ class StepFunctionsDeployedFlow(DeployedFlow):
         Exception
             If there is an error during the trigger process.
         """
-        with tempfile.TemporaryDirectory() as temp_dir:
-            tfp_runner_attribute = tempfile.NamedTemporaryFile(
-                dir=temp_dir, delete=False
-            )
-
+        with temporary_fifo() as (attribute_file_path, attribute_file_fd):
             # every subclass needs to have `self.deployer_kwargs`
             command = get_lower_level_group(
                 self.deployer.api,
                 self.deployer.top_level_kwargs,
                 self.deployer.TYPE,
                 self.deployer.deployer_kwargs,
-            ).trigger(deployer_attribute_file=tfp_runner_attribute.name, **kwargs)
+            ).trigger(deployer_attribute_file=attribute_file_path, **kwargs)
 
             pid = self.deployer.spm.run_command(
                 [sys.executable, *command],
@@ -218,7 +214,7 @@ class StepFunctionsDeployedFlow(DeployedFlow):
 
             command_obj = self.deployer.spm.get(pid)
             content = handle_timeout(
-                tfp_runner_attribute, command_obj, self.deployer.file_read_timeout
+                attribute_file_fd, command_obj, self.deployer.file_read_timeout
             )
 
             if command_obj.process.returncode == 0:

metaflow/plugins/datatools/s3/s3.py

@@ -600,7 +600,9 @@ class S3(object):
         # returned are Unicode.
         key = getattr(key_value, "key", key_value)
         if self._s3root is None:
-            parsed = urlparse(to_unicode(key))
+            # NOTE: S3 allows fragments as part of object names, e.g. /dataset #1/data.txt
+            # Without allow_fragments=False the parsed.path for an object name with fragments is incomplete.
+            parsed = urlparse(to_unicode(key), allow_fragments=False)
             if parsed.scheme == "s3" and parsed.path:
                 return key
             else:
@@ -765,7 +767,9 @@ class S3(object):
         """
 
         url = self._url(key)
-        src = urlparse(url)
+        # NOTE: S3 allows fragments as part of object names, e.g. /dataset #1/data.txt
+        # Without allow_fragments=False the parsed src.path for an object name with fragments is incomplete.
+        src = urlparse(url, allow_fragments=False)
 
         def _info(s3, tmp):
             resp = s3.head_object(Bucket=src.netloc, Key=src.path.lstrip('/"'))
@@ -891,7 +895,9 @@ class S3(object):
         DOWNLOAD_MAX_CHUNK = 2 * 1024 * 1024 * 1024 - 1
 
         url, r = self._url_and_range(key)
-        src = urlparse(url)
+        # NOTE: S3 allows fragments as part of object names, e.g. /dataset #1/data.txt
+        # Without allow_fragments=False the parsed src.path for an object name with fragments is incomplete.
+        src = urlparse(url, allow_fragments=False)
 
         def _download(s3, tmp):
             if r:
@@ -1173,7 +1179,9 @@ class S3(object):
         blob.close = lambda: None
 
         url = self._url(key)
-        src = urlparse(url)
+        # NOTE: S3 allows fragments as part of object names, e.g. /dataset #1/data.txt
+        # Without allow_fragments=False the parsed src.path for an object name with fragments is incomplete.
+        src = urlparse(url, allow_fragments=False)
         extra_args = None
         if content_type or metadata or self._encryption:
             extra_args = {}