metaflow 2.12.21__tar.gz → 2.12.23__tar.gz

{metaflow-2.12.21/metaflow.egg-info → metaflow-2.12.23}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: metaflow
-Version: 2.12.21
+Version: 2.12.23
 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.21; extra == "stubs"
+Requires-Dist: metaflow-stubs==2.12.23; 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.21 → metaflow-2.12.23}/metaflow/cli.py

@@ -1,5 +1,5 @@
 import inspect
-import os
+import json
 import sys
 import traceback
 from datetime import datetime
@@ -7,6 +7,7 @@ from functools import wraps
 
 import metaflow.tracing as tracing
 from metaflow._vendor import click
+from metaflow.client.core import get_metadata
 
 from . import decorators, lint, metaflow_version, namespace, parameters, plugins
 from .cli_args import cli_args
@@ -698,15 +699,17 @@ def resume(
     runtime.print_workflow_info()
 
     runtime.persist_constants()
-    write_file(
-        runner_attribute_file,
-        "%s@%s:%s"
-        % (
-            obj.metadata.__class__.TYPE,
-            obj.metadata.__class__.INFO,
-            "/".join((obj.flow.name, runtime.run_id)),
-        ),
-    )
+
+    if runner_attribute_file:
+        with open(runner_attribute_file, "w") as f:
+            json.dump(
+                {
+                    "run_id": runtime.run_id,
+                    "flow_name": obj.flow.name,
+                    "metadata": get_metadata(),
+                },
+                f,
+            )
 
     # We may skip clone-only resume if this is not a resume leader,
     # and clone is already complete.
@@ -774,15 +777,17 @@ def run(
     obj.flow._set_constants(obj.graph, kwargs)
     runtime.print_workflow_info()
     runtime.persist_constants()
-    write_file(
-        runner_attribute_file,
-        "%s@%s:%s"
-        % (
-            obj.metadata.__class__.TYPE,
-            obj.metadata.__class__.INFO,
-            "/".join((obj.flow.name, runtime.run_id)),
-        ),
-    )
+
+    if runner_attribute_file:
+        with open(runner_attribute_file, "w") as f:
+            json.dump(
+                {
+                    "run_id": runtime.run_id,
+                    "flow_name": obj.flow.name,
+                    "metadata": get_metadata(),
+                },
+                f,
+            )
     runtime.execute()
 
 

{metaflow-2.12.21 → metaflow-2.12.23}/metaflow/client/core.py

@@ -2358,7 +2358,7 @@ class Metaflow(object):
         Flow
             Flow with the given name.
         """
-        return Flow(id)
+        return Flow(name)
 
 
 _CLASSES["flow"] = Flow

{metaflow-2.12.21 → metaflow-2.12.23}/metaflow/cmd/develop/stub_generator.py

@@ -115,6 +115,11 @@ class StubGenerator:
         :type members_from_other_modules: List[str]
         """
 
+        # Let metaflow know we are in stubgen mode. This is sometimes useful to skip
+        # some processing like loading libraries, etc. It is used in Metaflow extensions
+        # so do not remove even if you do not see a use for it directly in the code.
+        os.environ["METAFLOW_STUBGEN"] = "1"
+
         self._write_generated_for = include_generated_for
         self._pending_modules = ["metaflow"]  # type: List[str]
         self._pending_modules.extend(get_aliased_modules())
@@ -398,6 +403,18 @@ class StubGenerator:
         name_with_module = self._get_element_name_with_module(clazz.__class__)
         buff.write("metaclass=" + name_with_module + "):\n")
 
+        # Add class docstring
+        if clazz.__doc__:
+            buff.write('%s"""\n' % TAB)
+            my_doc = cast(str, deindent_docstring(clazz.__doc__))
+            init_blank = True
+            for line in my_doc.split("\n"):
+                if init_blank and len(line.strip()) == 0:
+                    continue
+                init_blank = False
+                buff.write("%s%s\n" % (TAB, line.rstrip()))
+            buff.write('%s"""\n' % TAB)
+
         # For NamedTuple, we have __annotations__ but no __init__. In that case,
         # we are going to "create" a __init__ function with the annotations
         # to show what the class takes.

{metaflow-2.12.21 → metaflow-2.12.23}/metaflow/cmd/develop/stubs.py

@@ -170,7 +170,7 @@ def install(ctx: Any, force: bool):
                     "Metaflow stubs are already installed and valid -- use --force to reinstall"
                 )
             return
-    mf_version, _ = get_mf_version()
+    mf_version, _ = get_mf_version(True)
     with tempfile.TemporaryDirectory() as tmp_dir:
         with open(os.path.join(tmp_dir, "setup.py"), "w") as f:
             f.write(
@@ -261,10 +261,10 @@ def split_version(vers: str) -> Tuple[str, Optional[str]]:
     return vers_split[0], vers_split[1]
 
 
-def get_mf_version() -> Tuple[str, Optional[str]]:
+def get_mf_version(public: bool = False) -> Tuple[str, Optional[str]]:
     from metaflow.metaflow_version import get_version
 
-    return split_version(get_version())
+    return split_version(get_version(public))
 
 
 def get_stubs_version(stubs_root_path: Optional[str]) -> Tuple[str, Optional[str]]:

{metaflow-2.12.21 → metaflow-2.12.23}/metaflow/metaflow_version.py

@@ -195,11 +195,14 @@ def get_version(public=False):
         if ext_version is None:
             ext_version = getattr(extension_module, "__version__", "<unk>")
         # Update the package information about reported version for the extension
-        update_package_info(
-            package_name=pkg_name,
-            extension_name=ext_name,
-            package_version=ext_version,
-        )
+        # (only for the full info which is called at least once -- if we update more
+        # it will error out since we can only update_package_info once)
+        if not public:
+            update_package_info(
+                package_name=pkg_name,
+                extension_name=ext_name,
+                package_version=ext_version,
+            )
         ext_versions.append("%s(%s)" % (ext_name, ext_version))
 
     # We now have all the information about extensions so we can form the final string

{metaflow-2.12.21 → metaflow-2.12.23}/metaflow/plugins/argo/argo_client.py

@@ -295,6 +295,8 @@ class ArgoClient(object):
                 "suspend": schedule is None,
                 "schedule": schedule,
                 "timezone": timezone,
+                "failedJobsHistoryLimit": 10000,  # default is unfortunately 1
+                "successfulJobsHistoryLimit": 10000,  # default is unfortunately 3
                 "workflowSpec": {"workflowTemplateRef": {"name": name}},
             },
         }

{metaflow-2.12.21 → metaflow-2.12.23}/metaflow/plugins/argo/argo_workflows.py

@@ -455,11 +455,17 @@ class ArgoWorkflows(object):
                 )
             seen.add(norm)
 
-            if param.kwargs.get("type") == JSONType or isinstance(
-                param.kwargs.get("type"), FilePathClass
-            ):
-                # Special-case this to avoid touching core
+            extra_attrs = {}
+            if param.kwargs.get("type") == JSONType:
+                param_type = str(param.kwargs.get("type").name)
+            elif isinstance(param.kwargs.get("type"), FilePathClass):
                 param_type = str(param.kwargs.get("type").name)
+                extra_attrs["is_text"] = getattr(
+                    param.kwargs.get("type"), "_is_text", True
+                )
+                extra_attrs["encoding"] = getattr(
+                    param.kwargs.get("type"), "_encoding", "utf-8"
+                )
             else:
                 param_type = str(param.kwargs.get("type").__name__)
 
@@ -487,6 +493,7 @@ class ArgoWorkflows(object):
                 type=param_type,
                 description=param.kwargs.get("help"),
                 is_required=is_required,
+                **extra_attrs
             )
         return parameters
 
@@ -1483,7 +1490,11 @@ class ArgoWorkflows(object):
                         # {{foo.bar['param_name']}}.
                         # https://argoproj.github.io/argo-events/tutorials/02-parameterization/
                         # http://masterminds.github.io/sprig/strings.html
-                        "--%s={{workflow.parameters.%s}}"
+                        (
+                            "--%s='{{workflow.parameters.%s}}'"
+                            if parameter["type"] == "JSON"
+                            else "--%s={{workflow.parameters.%s}}"
+                        )
                         % (parameter["name"], parameter["name"])
                         for parameter in self.parameters.values()
                     ]
@@ -2515,10 +2526,29 @@ class ArgoWorkflows(object):
         # Use all the affordances available to _parameters task
         executable = self.environment.executable("_parameters")
         run_id = "argo-{{workflow.name}}"
-        entrypoint = [executable, "-m metaflow.plugins.argo.daemon"]
-        heartbeat_cmds = "{entrypoint} --flow_name {flow_name} --run_id {run_id} {tags} heartbeat".format(
+        script_name = os.path.basename(sys.argv[0])
+        entrypoint = [executable, script_name]
+        # FlowDecorators can define their own top-level options. These might affect run level information
+        # so it is important to pass these to the heartbeat process as well, as it might be the first task to register a run.
+        top_opts_dict = {}
+        for deco in flow_decorators(self.flow):
+            top_opts_dict.update(deco.get_top_level_options())
+
+        top_level = list(dict_to_cli_options(top_opts_dict)) + [
+            "--quiet",
+            "--metadata=%s" % self.metadata.TYPE,
+            "--environment=%s" % self.environment.TYPE,
+            "--datastore=%s" % self.flow_datastore.TYPE,
+            "--datastore-root=%s" % self.flow_datastore.datastore_root,
+            "--event-logger=%s" % self.event_logger.TYPE,
+            "--monitor=%s" % self.monitor.TYPE,
+            "--no-pylint",
+            "--with=argo_workflows_internal:auto-emit-argo-events=%i"
+            % self.auto_emit_argo_events,
+        ]
+        heartbeat_cmds = "{entrypoint} {top_level} argo-workflows heartbeat --run_id {run_id} {tags}".format(
             entrypoint=" ".join(entrypoint),
-            flow_name=self.flow.name,
+            top_level=" ".join(top_level) if top_level else "",
             run_id=run_id,
             tags=" ".join(["--tag %s" % t for t in self.tags]) if self.tags else "",
         )
@@ -2569,12 +2599,16 @@ class ArgoWorkflows(object):
             "METAFLOW_SERVICE_URL": SERVICE_INTERNAL_URL,
             "METAFLOW_SERVICE_HEADERS": json.dumps(SERVICE_HEADERS),
             "METAFLOW_USER": "argo-workflows",
+            "METAFLOW_DATASTORE_SYSROOT_S3": DATASTORE_SYSROOT_S3,
+            "METAFLOW_DATATOOLS_S3ROOT": DATATOOLS_S3ROOT,
             "METAFLOW_DEFAULT_DATASTORE": self.flow_datastore.TYPE,
             "METAFLOW_DEFAULT_METADATA": DEFAULT_METADATA,
+            "METAFLOW_CARD_S3ROOT": CARD_S3ROOT,
             "METAFLOW_KUBERNETES_WORKLOAD": 1,
+            "METAFLOW_KUBERNETES_FETCH_EC2_METADATA": KUBERNETES_FETCH_EC2_METADATA,
             "METAFLOW_RUNTIME_ENVIRONMENT": "kubernetes",
             "METAFLOW_OWNER": self.username,
-            "METAFLOW_PRODUCTION_TOKEN": self.production_token,
+            "METAFLOW_PRODUCTION_TOKEN": self.production_token,  # Used in identity resolving. This affects system tags.
         }
         # support Metaflow sandboxes
         env["METAFLOW_INIT_SCRIPT"] = KUBERNETES_SANDBOX_INIT_SCRIPT
@@ -2597,50 +2631,54 @@ class ArgoWorkflows(object):
         )
         from kubernetes import client as kubernetes_sdk
 
-        return DaemonTemplate("heartbeat-daemon").container(
-            to_camelcase(
-                kubernetes_sdk.V1Container(
-                    name="main",
-                    # TODO: Make the image configurable
-                    image=resources["image"],
-                    command=cmds,
-                    env=[
-                        kubernetes_sdk.V1EnvVar(name=k, value=str(v))
-                        for k, v in env.items()
-                    ],
-                    env_from=[
-                        kubernetes_sdk.V1EnvFromSource(
-                            secret_ref=kubernetes_sdk.V1SecretEnvSource(
-                                name=str(k),
-                                # optional=True
+        return (
+            DaemonTemplate("heartbeat-daemon")
+            .service_account_name(resources["service_account"])
+            .container(
+                to_camelcase(
+                    kubernetes_sdk.V1Container(
+                        name="main",
+                        # TODO: Make the image configurable
+                        image=resources["image"],
+                        command=cmds,
+                        env=[
+                            kubernetes_sdk.V1EnvVar(name=k, value=str(v))
+                            for k, v in env.items()
+                        ],
+                        env_from=[
+                            kubernetes_sdk.V1EnvFromSource(
+                                secret_ref=kubernetes_sdk.V1SecretEnvSource(
+                                    name=str(k),
+                                    # optional=True
+                                )
                             )
-                        )
-                        for k in list(
-                            []
-                            if not resources.get("secrets")
-                            else (
-                                [resources.get("secrets")]
-                                if isinstance(resources.get("secrets"), str)
-                                else resources.get("secrets")
+                            for k in list(
+                                []
+                                if not resources.get("secrets")
+                                else (
+                                    [resources.get("secrets")]
+                                    if isinstance(resources.get("secrets"), str)
+                                    else resources.get("secrets")
+                                )
                             )
-                        )
-                        + KUBERNETES_SECRETS.split(",")
-                        + ARGO_WORKFLOWS_KUBERNETES_SECRETS.split(",")
-                        if k
-                    ],
-                    resources=kubernetes_sdk.V1ResourceRequirements(
-                        # NOTE: base resources for this are kept to a minimum to save on running costs.
-                        # This has an adverse effect on startup time for the daemon, which can be completely
-                        # alleviated by using a base image that has the required dependencies pre-installed
-                        requests={
-                            "cpu": "200m",
-                            "memory": "100Mi",
-                        },
-                        limits={
-                            "cpu": "200m",
-                            "memory": "100Mi",
-                        },
-                    ),
+                            + KUBERNETES_SECRETS.split(",")
+                            + ARGO_WORKFLOWS_KUBERNETES_SECRETS.split(",")
+                            if k
+                        ],
+                        resources=kubernetes_sdk.V1ResourceRequirements(
+                            # NOTE: base resources for this are kept to a minimum to save on running costs.
+                            # This has an adverse effect on startup time for the daemon, which can be completely
+                            # alleviated by using a base image that has the required dependencies pre-installed
+                            requests={
+                                "cpu": "200m",
+                                "memory": "100Mi",
+                            },
+                            limits={
+                                "cpu": "200m",
+                                "memory": "100Mi",
+                            },
+                        ),
+                    )
                 )
             )
         )
@@ -3262,6 +3300,10 @@ class DaemonTemplate(object):
         self.payload["container"] = container
         return self
 
+    def service_account_name(self, service_account_name):
+        self.payload["serviceAccountName"] = service_account_name
+        return self
+
     def to_json(self):
         return self.payload
 

{metaflow-2.12.21 → metaflow-2.12.23}/metaflow/plugins/argo/argo_workflows_cli.py

@@ -4,6 +4,7 @@ import platform
 import re
 import sys
 from hashlib import sha1
+from time import sleep
 
 from metaflow import JSONType, Run, current, decorators, parameters
 from metaflow._vendor import click
@@ -959,6 +960,31 @@ def list_workflow_templates(obj, all=None):
         obj.echo_always(template_name)
 
 
+# Internal CLI command to run a heartbeat daemon in an Argo Workflows Daemon container.
+@argo_workflows.command(hidden=True, help="start heartbeat process for a run")
+@click.option("--run_id", required=True)
+@click.option(
+    "--tag",
+    "tags",
+    multiple=True,
+    default=None,
+    help="Annotate all objects produced by Argo Workflows runs "
+    "with the given tag. You can specify this option multiple "
+    "times to attach multiple tags.",
+)
+@click.pass_obj
+def heartbeat(obj, run_id, tags=None):
+    # Try to register a run in case the start task has not taken care of it yet.
+    obj.metadata.register_run_id(run_id, tags)
+    # Start run heartbeat
+    obj.metadata.start_run_heartbeat(obj.flow.name, run_id)
+    # Keepalive loop
+    while True:
+        # Do not pollute daemon logs with anything unnecessary,
+        # as they might be extremely long running.
+        sleep(10)
+
+
 def validate_run_id(
     workflow_name, token_prefix, authorize, run_id, instructions_fn=None
 ):

{metaflow-2.12.21 → metaflow-2.12.23}/metaflow/plugins/kubernetes/kubernetes_client.py

@@ -121,7 +121,10 @@ class KubernetesClient(object):
         job_api = self._client.BatchV1Api()
         pods = self._find_active_pods(flow_name, run_id, user)
 
+        active_pods = False
+
         def _kill_pod(pod):
+            active_pods = True
             echo("Killing Kubernetes pod %s\n" % pod.metadata.name)
             try:
                 stream(
@@ -155,7 +158,10 @@ class KubernetesClient(object):
                     echo("failed to kill pod %s - %s" % (pod.metadata.name, str(e)))
 
         with ThreadPoolExecutor() as executor:
-            executor.map(_kill_pod, list(pods))
+            executor.map(_kill_pod, pods)
+
+        if not active_pods:
+            echo("No active Kubernetes pods found for run *%s*" % run_id)
 
     def jobset(self, **kwargs):
         return KubernetesJobSet(self, **kwargs)

{metaflow-2.12.21 → metaflow-2.12.23}/metaflow/plugins/kubernetes/kubernetes_decorator.py

@@ -543,7 +543,11 @@ class KubernetesDecorator(StepDecorator):
 
 # TODO: Unify this method with the multi-node setup in @batch
 def _setup_multinode_environment():
-    # FIXME: what about MF_MASTER_PORT
+    # TODO [FIXME SOON]
+    # Even if Kubernetes may deploy control pods before worker pods, there is always a
+    # possibility that the worker pods may start before the control. In the case that this happens,
+    # the worker pods will not be able to resolve the control pod's IP address and this will cause
+    # the worker pods to fail. This function should account for this in the near future.
     import socket
 
     try:

{metaflow-2.12.21 → metaflow-2.12.23}/metaflow/plugins/kubernetes/kubernetes_jobsets.py

@@ -866,7 +866,13 @@ class KubernetesJobSet(object):
             spec=dict(
                 replicatedJobs=[self.control.dump(), self.worker.dump()],
                 suspend=False,
-                startupPolicy=None,
+                startupPolicy=dict(
+                    # We explicitly set an InOrder Startup policy so that
+                    # we can ensure that the control pod starts before the worker pods.
+                    # This is required so that when worker pods try to access the control's IP
+                    # we are able to resolve the control's IP address.
+                    startupPolicyOrder="InOrder"
+                ),
                 successPolicy=None,
                 # The Failure Policy helps setting the number of retries for the jobset.
                 # but we don't rely on it and instead rely on either the local scheduler

{metaflow-2.12.21 → metaflow-2.12.23}/metaflow/plugins/pypi/bootstrap.py

@@ -89,7 +89,7 @@ if __name__ == "__main__":
         # TODO: micromamba installation can be pawned off to micromamba.py
         f"""set -e;
         if ! command -v micromamba >/dev/null 2>&1; then
-            mkdir micromamba;
+            mkdir -p micromamba;
             python -c "import requests, bz2, sys; data = requests.get('https://micro.mamba.pm/api/micromamba/{architecture}/1.5.7').content; sys.stdout.buffer.write(bz2.decompress(data))" | tar -xv -C $(pwd)/micromamba bin/micromamba --strip-components 1;
             export PATH=$PATH:$(pwd)/micromamba;
             if ! command -v micromamba >/dev/null 2>&1; then

{metaflow-2.12.21 → metaflow-2.12.23}/metaflow/plugins/pypi/micromamba.py

@@ -253,7 +253,33 @@ class Micromamba(object):
             try:
                 output = json.loads(e.output)
                 err = []
+                v_pkgs = ["__cuda", "__glibc"]
                 for error in output.get("solver_problems", []):
+                    # raise a specific error message for virtual package related errors
+                    match = next((p for p in v_pkgs if p in error), None)
+                    if match is not None:
+                        vpkg_name = match[2:]
+                        # try to strip version from error msg which are of the format:
+                        # nothing provides <__vpkg> >=2.17,<3.0.a0 needed by <pkg_name>
+                        try:
+                            vpkg_version = error[
+                                len("nothing provides %s " % match) : error.index(
+                                    " needed by"
+                                )
+                            ]
+                        except ValueError:
+                            vpkg_version = None
+                        raise MicromambaException(
+                            "Please set the environment variable CONDA_OVERRIDE_{var} to a specific version{version} of {name}.\n"
+                            "Here is an example of supplying environment variables through the command line -\n\n"
+                            "CONDA_OVERRIDE_{var}=<{name}-version> python flow.py <args>".format(
+                                var=vpkg_name.upper(),
+                                version=(
+                                    "" if not vpkg_version else (" (%s)" % vpkg_version)
+                                ),
+                                name=vpkg_name,
+                            ),
+                        )
                     err.append(error)
                 raise MicromambaException(
                     msg.format(

{metaflow-2.12.21 → metaflow-2.12.23}/metaflow/runner/deployer.py

@@ -6,56 +6,11 @@ import importlib
 import functools
 import tempfile
 
-from subprocess import CalledProcessError
 from typing import Optional, Dict, ClassVar
 
 from metaflow.exception import MetaflowNotFound
-from metaflow.runner.subprocess_manager import CommandManager, SubprocessManager
-from metaflow.runner.utils import read_from_file_when_ready
-
-
-def handle_timeout(
-    tfp_runner_attribute, command_obj: CommandManager, file_read_timeout: int
-):
-    """
-    Handle the timeout for a running subprocess command that reads a file
-    and raises an error with appropriate logs if a TimeoutError occurs.
-
-    Parameters
-    ----------
-    tfp_runner_attribute : NamedTemporaryFile
-        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
-    -------
-    str
-        Content read from the temporary file.
-
-    Raises
-    ------
-    RuntimeError
-        If a TimeoutError occurs, it raises a RuntimeError with the command's
-        stdout and stderr logs.
-    """
-    try:
-        content = read_from_file_when_ready(
-            tfp_runner_attribute.name, command_obj, timeout=file_read_timeout
-        )
-        return content
-    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)
-        error_message = "Error executing: '%s':\n" % command
-        if stdout_log.strip():
-            error_message += "\nStdout:\n%s\n" % stdout_log
-        if stderr_log.strip():
-            error_message += "\nStderr:\n%s\n" % stderr_log
-        raise RuntimeError(error_message) from e
+from metaflow.runner.subprocess_manager import SubprocessManager
+from metaflow.runner.utils import handle_timeout
 
 
 def get_lower_level_group(
@@ -209,7 +164,7 @@ class TriggeredRun(object):
             elif callable(v):
                 setattr(self, k, functools.partial(v, self))
             else:
-                setattr(self.__class__, k, property(fget=lambda _, v=v: v))
+                setattr(self, k, v)
 
     def wait_for_run(self, timeout=None):
         """
@@ -287,7 +242,7 @@ class DeployedFlow(object):
             elif callable(v):
                 setattr(self, k, functools.partial(v, self))
             else:
-                setattr(self.__class__, k, property(fget=lambda _, v=v: v))
+                setattr(self, k, v)
 
 
 class DeployerImpl(object):