ob-metaflow 2.12.20.1__py2.py3-none-any.whl → 2.12.23.1__py2.py3-none-any.whl

metaflow/metaflow_version.py

@@ -7,11 +7,15 @@ See the documentation of get_version for more information
 
 # This file is adapted from https://github.com/aebrahim/python-git-version
 
-from subprocess import check_output, CalledProcessError
-from os import path, name, devnull, environ, listdir
-import json
+import subprocess
+from os import path, name, environ, listdir
 
-from metaflow import CURRENT_DIRECTORY, INFO_FILE
+from metaflow.extension_support import update_package_info
+from metaflow.info_file import CURRENT_DIRECTORY, read_info_file
+
+
+# True/False correspond to the value `public`` in get_version
+_version_cache = {True: None, False: None}
 
 __all__ = ("get_version",)
 
@@ -57,87 +61,152 @@ if name == "nt":
     GIT_COMMAND = find_git_on_windows()
 
 
-def call_git_describe(abbrev=7):
+def call_git_describe(file_to_check, abbrev=7):
     """return the string output of git describe"""
     try:
-
-        # first, make sure we are actually in a Metaflow repo,
-        # not some other repo
-        with open(devnull, "w") as fnull:
-            arguments = [GIT_COMMAND, "rev-parse", "--show-toplevel"]
-            reponame = (
-                check_output(arguments, cwd=CURRENT_DIRECTORY, stderr=fnull)
-                .decode("ascii")
-                .strip()
-            )
-            if path.basename(reponame) != "metaflow":
-                return None
-
-        with open(devnull, "w") as fnull:
-            arguments = [GIT_COMMAND, "describe", "--tags", "--abbrev=%d" % abbrev]
-            return (
-                check_output(arguments, cwd=CURRENT_DIRECTORY, stderr=fnull)
-                .decode("ascii")
-                .strip()
-            )
-
-    except (OSError, CalledProcessError):
+        wd = path.dirname(file_to_check)
+        filename = path.basename(file_to_check)
+
+        # First check if the file is tracked in the GIT repository we are in
+        # We do this because in some setups and for some bizarre reason, python files
+        # are installed directly into a git repository (I am looking at you brew). We
+        # don't want to consider this a GIT install in that case.
+        args = [GIT_COMMAND, "ls-files", "--error-unmatch", filename]
+        git_return_code = subprocess.run(
+            args,
+            cwd=wd,
+            stderr=subprocess.DEVNULL,
+            stdout=subprocess.DEVNULL,
+            check=False,
+        ).returncode
+
+        if git_return_code != 0:
+            return None
+
+        args = [
+            GIT_COMMAND,
+            "describe",
+            "--tags",
+            "--dirty",
+            "--long",
+            "--abbrev=%d" % abbrev,
+        ]
+        return (
+            subprocess.check_output(args, cwd=wd, stderr=subprocess.DEVNULL)
+            .decode("ascii")
+            .strip()
+        )
+
+    except (OSError, subprocess.CalledProcessError):
         return None
 
 
-def format_git_describe(git_str, pep440=False):
+def format_git_describe(git_str, public=False):
     """format the result of calling 'git describe' as a python version"""
     if git_str is None:
         return None
-    if "-" not in git_str:  # currently at a tag
-        return git_str
+    splits = git_str.split("-")
+    if len(splits) == 4:
+        # Formatted as <tag>-<post>-<hash>-dirty
+        tag, post, h = splits[:3]
+        dirty = "-" + splits[3]
     else:
-        # formatted as version-N-githash
-        # want to convert to version.postN-githash
-        git_str = git_str.replace("-", ".post", 1)
-        if pep440:  # does not allow git hash afterwards
-            return git_str.split("-")[0]
-        else:
-            return git_str.replace("-g", "+git")
+        # Formatted as <tag>-<post>-<hash>
+        tag, post, h = splits
+        dirty = ""
+    if post == "0":
+        if public:
+            return tag
+        return tag + dirty
+
+    if public:
+        return "%s.post%s" % (tag, post)
+
+    return "%s.post%s-git%s%s" % (tag, post, h[1:], dirty)
 
 
 def read_info_version():
     """Read version information from INFO file"""
-    try:
-        with open(INFO_FILE, "r") as contents:
-            return json.load(contents).get("metaflow_version")
-    except IOError:
-        return None
+    info_file = read_info_file()
+    if info_file:
+        return info_file.get("metaflow_version")
+    return None
 
 
-def get_version(pep440=False):
+def get_version(public=False):
     """Tracks the version number.
 
-    pep440: bool
-        When True, this function returns a version string suitable for
-        a release as defined by PEP 440. When False, the githash (if
-        available) will be appended to the version string.
+    public: bool
+        When True, this function returns a *public* version specification which
+        doesn't include any local information (dirtiness or hash). See
+        https://packaging.python.org/en/latest/specifications/version-specifiers/#version-scheme
 
-    If the script is located within an active git repository,
-    git-describe is used to get the version information.
+    We first check the INFO file to see if we recorded a version of Metaflow. If there
+    is none, we check if we are in a GIT repository and if so, form the version
+    from that.
 
-    Otherwise, the version logged by package installer is returned.
-
-    If even that information isn't available (likely when executing on a
-    remote cloud instance), the version information is returned from INFO file
-    in the current directory.
+    Otherwise, we return the version of Metaflow that was installed.
 
     """
 
-    version = format_git_describe(call_git_describe(), pep440=pep440)
-    version_addl = None
-    if version is None:  # not a git repository
-        import metaflow
-
+    global _version_cache
+
+    # To get the version we do the following:
+    #  - Check if we have a cached version. If so, return that
+    #  - Then check if we have an INFO file present. If so, use that as it is
+    #    the most reliable way to get the version. In particular, when running remotely,
+    #    metaflow is installed in a directory and if any extension is using distutils to
+    #    determine its version, this would return None and querying the version directly
+    #    from the extension would fail to produce the correct result
+    #  - Then if we are in the GIT repository and if so, use the git describe
+    #  - If we don't have an INFO file, we look at the version information that is
+    #    populated by metaflow and the extensions.
+
+    if _version_cache[public] is not None:
+        return _version_cache[public]
+
+    version = (
+        read_info_version()
+    )  # Version info is cached in INFO file; includes extension info
+
+    if version:
+        _version_cache[public] = version
+        return version
+
+    # Get the version for Metaflow, favor the GIT version
+    import metaflow
+
+    version = format_git_describe(
+        call_git_describe(file_to_check=metaflow.__file__), public=public
+    )
+    if version is None:
         version = metaflow.__version__
-        version_addl = metaflow.__version_addl__
-    if version is None:  # not a proper python package
-        return read_info_version()
-    if version_addl:
-        return "+".join([version, version_addl])
+
+    # Look for extensions and compute their versions. Properly formed extensions have
+    # a toplevel file which will contain a __mf_extensions__ value and a __version__
+    # value. We already saved the properly formed modules when loading metaflow in
+    # __ext_tl_modules__.
+    ext_versions = []
+    for pkg_name, extension_module in metaflow.__ext_tl_modules__:
+        ext_name = getattr(extension_module, "__mf_extensions__", "<unk>")
+        ext_version = format_git_describe(
+            call_git_describe(file_to_check=extension_module.__file__), public=public
+        )
+        if ext_version is None:
+            ext_version = getattr(extension_module, "__version__", "<unk>")
+        # Update the package information about reported version for the extension
+        # (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
+    if ext_versions:
+        version = version + "+" + ";".join(ext_versions)
+    _version_cache[public] = version
     return version

metaflow/package.py

@@ -10,7 +10,8 @@ from .extension_support import EXT_PKG, package_mfext_all
 from .metaflow_config import DEFAULT_PACKAGE_SUFFIXES
 from .exception import MetaflowException
 from .util import to_unicode
-from . import R, INFO_FILE
+from . import R
+from .info_file import INFO_FILE
 
 DEFAULT_SUFFIXES_LIST = DEFAULT_PACKAGE_SUFFIXES.split(",")
 METAFLOW_SUFFIXES_LIST = [".py", ".html", ".css", ".js"]

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/plugins/argo/argo_workflows.py

@@ -456,11 +456,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__)
 
@@ -488,6 +494,7 @@ class ArgoWorkflows(object):
                 type=param_type,
                 description=param.kwargs.get("help"),
                 is_required=is_required,
+                **extra_attrs
             )
         return parameters
 
@@ -1484,7 +1491,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()
                     ]
@@ -2524,10 +2535,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 "",
         )
@@ -2578,12 +2608,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
@@ -2606,50 +2640,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",
+                            },
+                        ),
+                    )
                 )
             )
         )
@@ -3271,6 +3309,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/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/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/plugins/kubernetes/kubernetes_decorator.py

@@ -558,7 +558,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/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/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/plugins/pypi/conda_decorator.py

@@ -12,7 +12,7 @@ from metaflow.metadata import MetaDatum
 from metaflow.metaflow_environment import InvalidEnvironmentException
 from metaflow.util import get_metaflow_root
 
-from ... import INFO_FILE
+from ...info_file import INFO_FILE
 
 
 class CondaStepDecorator(StepDecorator):

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(