metaflow 2.12.20__py2.py3-none-any.whl → 2.12.22__py2.py3-none-any.whl

metaflow/__init__.py

@@ -42,14 +42,8 @@ If you have any questions, feel free to post a bug report/question on the
 Metaflow GitHub page.
 """
 
-import importlib
+import os
 import sys
-import types
-
-from os import path
-
-CURRENT_DIRECTORY = path.dirname(path.abspath(__file__))
-INFO_FILE = path.join(path.dirname(CURRENT_DIRECTORY), "INFO")
 
 from metaflow.extension_support import (
     alias_submodules,
@@ -61,7 +55,6 @@ from metaflow.extension_support import (
     _ext_debug,
 )
 
-
 # We load the module overrides *first* explicitly. Non overrides can be loaded
 # in toplevel as well but these can be loaded first if needed. Note that those
 # modules should be careful not to include anything in Metaflow at their top-level
@@ -79,9 +72,14 @@ try:
             )
         tl_module = m.module.__dict__.get("toplevel", None)
         if tl_module is not None:
-            _tl_modules.append(".".join([EXT_PKG, m.tl_package, "toplevel", tl_module]))
+            _tl_modules.append(
+                (
+                    m.package_name,
+                    ".".join([EXT_PKG, m.tl_package, "toplevel", tl_module]),
+                )
+            )
     _ext_debug("Got overrides to load: %s" % _override_modules)
-    _ext_debug("Got top-level imports: %s" % _tl_modules)
+    _ext_debug("Got top-level imports: %s" % str(_tl_modules))
 except Exception as e:
     _ext_debug("Error in importing toplevel/overrides: %s" % e)
 
@@ -153,9 +151,9 @@ if sys.version_info >= (3, 7):
     from .runner.deployer import Deployer
     from .runner.nbdeploy import NBDeployer
 
-__version_addl__ = []
+__ext_tl_modules__ = []
 _ext_debug("Loading top-level modules")
-for m in _tl_modules:
+for pkg_name, m in _tl_modules:
     extension_module = load_module(m)
     if extension_module:
         tl_package = m.split(".")[1]
@@ -163,15 +161,7 @@ for m in _tl_modules:
         lazy_load_aliases(
             alias_submodules(extension_module, tl_package, None, extra_indent=True)
         )
-        version_info = getattr(extension_module, "__mf_extensions__", "<unk>")
-        if extension_module.__version__:
-            version_info = "%s(%s)" % (version_info, extension_module.__version__)
-        __version_addl__.append(version_info)
-
-if __version_addl__:
-    __version_addl__ = ";".join(__version_addl__)
-else:
-    __version_addl__ = None
+        __ext_tl_modules__.append((pkg_name, extension_module))
 
 # Erase all temporary names to avoid leaking things
 for _n in [

metaflow/client/core.py

@@ -34,7 +34,7 @@ from metaflow.plugins import ENVIRONMENTS, METADATA_PROVIDERS
 from metaflow.unbounded_foreach import CONTROL_TASK_TAG
 from metaflow.util import cached_property, is_stringish, resolve_identity, to_unicode
 
-from .. import INFO_FILE
+from ..info_file import INFO_FILE
 from .filecache import FileCache
 
 try:

metaflow/cmd/main_cli.py

@@ -84,12 +84,13 @@ def start(ctx):
 
     import metaflow
 
+    version = get_version()
     echo("Metaflow ", fg="magenta", bold=True, nl=False)
 
     if ctx.invoked_subcommand is None:
-        echo("(%s): " % get_version(), fg="magenta", bold=False, nl=False)
+        echo("(%s): " % version, fg="magenta", bold=False, nl=False)
     else:
-        echo("(%s)\n" % get_version(), fg="magenta", bold=False)
+        echo("(%s)\n" % version, fg="magenta", bold=False)
 
     if ctx.invoked_subcommand is None:
         echo("More data science, less engineering\n", fg="magenta")

metaflow/extension_support/__init__.py

@@ -1,7 +1,6 @@
 from __future__ import print_function
 
 import importlib
-import json
 import os
 import re
 import sys
@@ -11,6 +10,10 @@ from collections import defaultdict, namedtuple
 
 from importlib.abc import MetaPathFinder, Loader
 from itertools import chain
+from pathlib import Path
+
+from metaflow.info_file import read_info_file
+
 
 #
 # This file provides the support for Metaflow's extension mechanism which allows
@@ -59,6 +62,9 @@ __all__ = (
     "load_module",
     "get_modules",
     "dump_module_info",
+    "get_extensions_in_dir",
+    "extension_info",
+    "update_package_info",
     "get_aliased_modules",
     "package_mfext_package",
     "package_mfext_all",
@@ -80,9 +86,14 @@ EXT_EXCLUDE_SUFFIXES = [".pyc"]
 # To get verbose messages, set METAFLOW_DEBUG_EXT to 1
 DEBUG_EXT = os.environ.get("METAFLOW_DEBUG_EXT", False)
 
+# This is extracted only from environment variable and here separately from
+# metaflow_config to prevent nasty circular dependencies
+EXTENSIONS_SEARCH_DIRS = os.environ.get("METAFLOW_EXTENSIONS_SEARCH_DIRS", "").split(
+    os.pathsep
+)
 
 MFExtPackage = namedtuple("MFExtPackage", "package_name tl_package config_module")
-MFExtModule = namedtuple("MFExtModule", "tl_package module")
+MFExtModule = namedtuple("MFExtModule", "package_name tl_package module")
 
 
 def load_module(module_name):
@@ -113,17 +124,64 @@ def get_modules(extension_point):
     return modules_to_load
 
 
-def dump_module_info():
-    _filter_files_all()
+def dump_module_info(all_packages=None, pkgs_per_extension_point=None):
+    if all_packages is None:
+        all_packages = _all_packages
+    if pkgs_per_extension_point is None:
+        pkgs_per_extension_point = _pkgs_per_extension_point
+
+    _filter_files_all(all_packages)
     sanitized_all_packages = dict()
     # Strip out root_paths (we don't need it and no need to expose user's dir structure)
-    for k, v in _all_packages.items():
+    for k, v in all_packages.items():
         sanitized_all_packages[k] = {
             "root_paths": None,
             "meta_module": v["meta_module"],
             "files": v["files"],
+            "version": v["version"],
+            "package_version": v.get("package_version", "<unk>"),
+            "extension_name": v.get("extension_name", "<unk>"),
         }
-    return "ext_info", [sanitized_all_packages, _pkgs_per_extension_point]
+    return "ext_info", [sanitized_all_packages, pkgs_per_extension_point]
+
+
+def get_extensions_in_dir(d):
+    if not _mfext_supported:
+        _ext_debug("Not supported for your Python version -- 3.4+ is needed")
+        return None, None
+    return _get_extension_packages(ignore_info_file=True, restrict_to_directories=[d])
+
+
+def extension_info(packages=None):
+    if packages is None:
+        packages = _all_packages
+    # Returns information about installed extensions so it it can be stored in
+    # _graph_info.
+    return {
+        "installed": {
+            k: {
+                "dist_version": v["version"],
+                "package_version": v.get("package_version", "<unk>"),
+                "extension_name": v.get("extension_name", "<unk>"),
+            }
+            for k, v in packages.items()
+        },
+    }
+
+
+def update_package_info(pkg_to_update=None, package_name=None, **kwargs):
+    pkg = None
+    if pkg_to_update:
+        pkg = pkg_to_update
+    elif package_name:
+        pkg = _all_packages.get(package_name)
+    for k, v in kwargs.items():
+        if k in pkg:
+            raise ValueError(
+                "Trying to overwrite existing key '%s' for package %s" % (k, str(pkg))
+            )
+        pkg[k] = v
+    return pkg
 
 
 def get_aliased_modules():
@@ -134,8 +192,8 @@ def package_mfext_package(package_name):
     from metaflow.util import to_unicode
 
     _ext_debug("Packaging '%s'" % package_name)
-    _filter_files_package(package_name)
     pkg_info = _all_packages.get(package_name, None)
+    _filter_files_package(pkg_info)
     if pkg_info and pkg_info.get("root_paths", None):
         single_path = len(pkg_info["root_paths"]) == 1
         for p in pkg_info["root_paths"]:
@@ -296,7 +354,7 @@ def _ext_debug(*args, **kwargs):
         print(init_str, *args, **kwargs)
 
 
-def _get_extension_packages():
+def _get_extension_packages(ignore_info_file=False, restrict_to_directories=None):
     if not _mfext_supported:
         _ext_debug("Not supported for your Python version -- 3.4+ is needed")
         return {}, {}
@@ -304,20 +362,20 @@ def _get_extension_packages():
     # If we have an INFO file with the appropriate information (if running from a saved
     # code package for example), we use that directly
     # Pre-compute on _extension_points
-    from metaflow import INFO_FILE
-
-    try:
-        with open(INFO_FILE, encoding="utf-8") as contents:
-            all_pkg, ext_to_pkg = json.load(contents).get("ext_info", (None, None))
-            if all_pkg is not None and ext_to_pkg is not None:
-                _ext_debug("Loading pre-computed information from INFO file")
-                # We need to properly convert stuff in ext_to_pkg
-                for k, v in ext_to_pkg.items():
-                    v = [MFExtPackage(*d) for d in v]
-                    ext_to_pkg[k] = v
-                return all_pkg, ext_to_pkg
-    except IOError:
-        pass
+    info_content = read_info_file()
+    if not ignore_info_file and info_content:
+        all_pkg, ext_to_pkg = info_content.get("ext_info", (None, None))
+        if all_pkg is not None and ext_to_pkg is not None:
+            _ext_debug("Loading pre-computed information from INFO file")
+            # We need to properly convert stuff in ext_to_pkg
+            for k, v in ext_to_pkg.items():
+                v = [MFExtPackage(*d) for d in v]
+                ext_to_pkg[k] = v
+            return all_pkg, ext_to_pkg
+
+    # Late import to prevent some circular nastiness
+    if restrict_to_directories is None and EXTENSIONS_SEARCH_DIRS != [""]:
+        restrict_to_directories = EXTENSIONS_SEARCH_DIRS
 
     # Check if we even have extensions
     try:
@@ -331,6 +389,11 @@ def _get_extension_packages():
                 raise
         return {}, {}
 
+    if restrict_to_directories:
+        restrict_to_directories = [
+            Path(p).resolve().as_posix() for p in restrict_to_directories
+        ]
+
     # There are two "types" of packages:
     #   - those installed on the system (distributions)
     #   - those present in the PYTHONPATH
@@ -343,8 +406,12 @@ def _get_extension_packages():
     # At this point, we look at all the paths and create a set. As we find distributions
     # that match it, we will remove from the set and then will be left with any
     # PYTHONPATH "packages"
-    all_paths = set(extensions_module.__path__)
+    all_paths = set(Path(p).resolve().as_posix() for p in extensions_module.__path__)
     _ext_debug("Found packages present at %s" % str(all_paths))
+    if restrict_to_directories:
+        _ext_debug(
+            "Processed packages will be restricted to %s" % str(restrict_to_directories)
+        )
 
     list_ext_points = [x.split(".") for x in _extension_points]
     init_ext_points = [x[0] for x in list_ext_points]
@@ -391,9 +458,20 @@ def _get_extension_packages():
             # This is not 100% accurate because it is possible that at the same
             # location there is a package and a non-package, but this is extremely
             # unlikely so we are going to ignore this case.
-            dist_root = dist.locate_file(EXT_PKG).as_posix()
+            dist_root = dist.locate_file(EXT_PKG).resolve().as_posix()
             all_paths.discard(dist_root)
             dist_name = dist.metadata["Name"]
+            dist_version = dist.metadata["Version"]
+            if restrict_to_directories:
+                parent_dirs = list(
+                    p.as_posix() for p in Path(dist_root).resolve().parents
+                )
+                if all(p not in parent_dirs for p in restrict_to_directories):
+                    _ext_debug(
+                        "Ignoring package at %s as it is not in the considered directories"
+                        % dist_root
+                    )
+                    continue
             if dist_name in mf_ext_packages:
                 _ext_debug(
                     "Ignoring duplicate package '%s' (duplicate paths in sys.path? (%s))"
@@ -537,6 +615,7 @@ def _get_extension_packages():
                 "root_paths": [dist_root],
                 "meta_module": meta_module,
                 "files": files_to_include,
+                "version": dist_version,
             }
     # At this point, we have all the packages that contribute to EXT_PKG,
     # we now check to see if there is an order to respect based on dependencies. We will
@@ -605,6 +684,16 @@ def _get_extension_packages():
     if len(all_paths_list) > 0:
         _ext_debug("Non installed packages present at %s" % str(all_paths))
         for package_count, package_path in enumerate(all_paths_list):
+            if restrict_to_directories:
+                parent_dirs = list(
+                    p.as_posix() for p in Path(package_path).resolve().parents
+                )
+                if all(p not in parent_dirs for p in restrict_to_directories):
+                    _ext_debug(
+                        "Ignoring non-installed package at %s as it is not in "
+                        "the considered directories" % package_path
+                    )
+                    continue
             # We give an alternate name for the visible package name. It is
             # not exposed to the end user but used to refer to the package, and it
             # doesn't provide much additional information to have the full path
@@ -740,6 +829,7 @@ def _get_extension_packages():
                 "root_paths": [package_path],
                 "meta_module": meta_module,
                 "files": files_to_include,
+                "version": "_local_",
             }
 
     # Sanity check that we only have one package per configuration file.
@@ -868,12 +958,13 @@ def _get_extension_config(distribution_name, tl_pkg, extension_point, config_mod
             _ext_debug("Package '%s' is rooted at %s" % (distribution_name, root_paths))
             _all_packages[distribution_name]["root_paths"] = root_paths
 
-        return MFExtModule(tl_package=tl_pkg, module=extension_module)
+        return MFExtModule(
+            package_name=distribution_name, tl_package=tl_pkg, module=extension_module
+        )
     return None
 
 
-def _filter_files_package(package_name):
-    pkg = _all_packages.get(package_name)
+def _filter_files_package(pkg):
     if pkg and pkg["root_paths"] and pkg["meta_module"]:
         meta_module = _attempt_load_module(pkg["meta_module"])
         if meta_module:
@@ -902,8 +993,8 @@ def _filter_files_package(package_name):
             pkg["files"] = new_files
 
 
-def _filter_files_all():
-    for p in _all_packages:
+def _filter_files_all(all_packages):
+    for p in all_packages.values():
         _filter_files_package(p)
 
 

metaflow/flowspec.py

@@ -15,6 +15,9 @@ from .exception import (
     MissingInMergeArtifactsException,
     UnhandledInMergeArtifactsException,
 )
+
+from .extension_support import extension_info
+
 from .graph import FlowGraph
 from .unbounded_foreach import UnboundedForeachInput
 from .util import to_pod
@@ -208,6 +211,7 @@ class FlowSpec(metaclass=_FlowSpecMeta):
                 for deco in flow_decorators(self)
                 if not deco.name.startswith("_")
             ],
+            "extensions": extension_info(),
         }
         self._graph_info = graph_info
 

metaflow/info_file.py

@@ -0,0 +1,25 @@
+import json
+
+from os import path
+
+CURRENT_DIRECTORY = path.dirname(path.abspath(__file__))
+INFO_FILE = path.join(path.dirname(CURRENT_DIRECTORY), "INFO")
+
+_info_file_content = None
+_info_file_present = None
+
+
+def read_info_file():
+    global _info_file_content
+    global _info_file_present
+    if _info_file_present is None:
+        _info_file_present = path.exists(INFO_FILE)
+        if _info_file_present:
+            try:
+                with open(INFO_FILE, "r", encoding="utf-8") as contents:
+                    _info_file_content = json.load(contents)
+            except IOError:
+                pass
+    if _info_file_present:
+        return _info_file_content
+    return None

metaflow/metaflow_config.py

@@ -25,7 +25,6 @@ DATASTORE_LOCAL_DIR = ".metaflow"
 # Local configuration file (in .metaflow) containing overrides per-project
 LOCAL_CONFIG_FILE = "config.json"
 
-
 ###
 # Default configuration
 ###

metaflow/metaflow_environment.py

@@ -9,8 +9,6 @@ from metaflow.extension_support import dump_module_info
 from metaflow.mflog import BASH_MFLOG
 from . import R
 
-version_cache = None
-
 
 class InvalidEnvironmentException(MetaflowException):
     headline = "Incompatible environment"
@@ -180,10 +178,6 @@ class MetaflowEnvironment(object):
         return cmds
 
     def get_environment_info(self, include_ext_info=False):
-        global version_cache
-        if version_cache is None:
-            version_cache = metaflow_version.get_version()
-
         # note that this dict goes into the code package
         # so variables here should be relatively stable (no
         # timestamps) so the hash won't change all the time
@@ -197,7 +191,7 @@ class MetaflowEnvironment(object):
             "use_r": R.use_r(),
             "python_version": sys.version,
             "python_version_code": "%d.%d.%d" % sys.version_info[:3],
-            "metaflow_version": version_cache,
+            "metaflow_version": metaflow_version.get_version(),
             "script": os.path.basename(os.path.abspath(sys.argv[0])),
         }
         if R.use_r():

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,149 @@ 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
+        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_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
 
@@ -2515,10 +2522,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 +2595,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

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/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/version.py

@@ -1 +1 @@
-metaflow_version = "2.12.20"
+metaflow_version = "2.12.22"

{metaflow-2.12.20.dist-info → metaflow-2.12.22.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: metaflow
-Version: 2.12.20
+Version: 2.12.22
 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.20; extra == "stubs"
+Requires-Dist: metaflow-stubs==2.12.22; 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.20.dist-info → metaflow-2.12.22.dist-info}/RECORD

@@ -1,5 +1,5 @@
 metaflow/R.py,sha256=CqVfIatvmjciuICNnoyyNGrwE7Va9iXfLdFbQa52hwA,3958
-metaflow/__init__.py,sha256=1TJeSmvQVELJpVGh9s9GoB6aBIK0_zEl9mpW4eSz64Q,5971
+metaflow/__init__.py,sha256=9pnkRH00bbtzaQDSxEspANaB3r6YDC37_EjX7GxtJug,5641
 metaflow/cards.py,sha256=tP1_RrtmqdFh741pqE4t98S7SA0MtGRlGvRICRZF1Mg,426
 metaflow/cli.py,sha256=fa6dx6F2PKtq0oeMxRM2efA9LyE108jGZQvYCtiUS94,34274
 metaflow/cli_args.py,sha256=lcgBGNTvfaiPxiUnejAe60Upt9swG6lRy1_3OqbU6MY,2616
@@ -10,20 +10,21 @@ metaflow/decorators.py,sha256=hbJwRlZuLbl6t7fOsTiH7Sux4Lx6zgEEpldIXoM5TQc,21540
 metaflow/event_logger.py,sha256=joTVRqZPL87nvah4ZOwtqWX8NeraM_CXKXXGVpKGD8o,780
 metaflow/events.py,sha256=ahjzkSbSnRCK9RZ-9vTfUviz_6gMvSO9DGkJ86X80-k,5300
 metaflow/exception.py,sha256=KC1LHJQzzYkWib0DeQ4l_A2r8VaudywsSqIQuq1RDZU,4954
-metaflow/flowspec.py,sha256=0-NnCKiX4oJ21Spb5PJkl2IvHoJxBcpi8CQrZEGB2Ag,27178
+metaflow/flowspec.py,sha256=4w7Hm1aFSmVBOq0x9-LtFAq4eDjFb5F6Rdyh_NW41BQ,27270
 metaflow/graph.py,sha256=HFJ7V_bPSht_NHIm8BejrSqOX2fyBQpVOczRCliRw08,11975
 metaflow/includefile.py,sha256=yHczcZ_U0SrasxSNhZb3DIBzx8UZnrJCl3FzvpEQLOA,19753
+metaflow/info_file.py,sha256=wtf2_F0M6dgiUu74AFImM8lfy5RrUw5Yj7Rgs2swKRY,686
 metaflow/integrations.py,sha256=LlsaoePRg03DjENnmLxZDYto3NwWc9z_PtU6nJxLldg,1480
 metaflow/lint.py,sha256=5rj1MlpluxyPTSINjtMoJ7viotyNzfjtBJSAihlAwMU,10870
-metaflow/metaflow_config.py,sha256=4PUd2-JWJs35SaobDgMg4RdpZaKjhEqPUFh2f0pjqnU,22853
+metaflow/metaflow_config.py,sha256=cFd_RaZnxAWQX1kNEA-Y6HlxLvTpL-OM-d1M0CQoz_M,22852
 metaflow/metaflow_config_funcs.py,sha256=5GlvoafV6SxykwfL8D12WXSfwjBN_NsyuKE_Q3gjGVE,6738
 metaflow/metaflow_current.py,sha256=pC-EMnAsnvBLvLd61W6MvfiCKcboryeui9f6r8z_sg8,7161
-metaflow/metaflow_environment.py,sha256=D7zHYqe8aLZVwJ20nx19vmsNW29Kf3PVE8hbBjVQin8,8115
+metaflow/metaflow_environment.py,sha256=rojFyGdyY56sN1HaEb1-0XX53Q3XPNnl0SaH-8xXZ8w,7987
 metaflow/metaflow_profile.py,sha256=jKPEW-hmAQO-htSxb9hXaeloLacAh41A35rMZH6G8pA,418
-metaflow/metaflow_version.py,sha256=mPQ6g_3XjNdi0NrxDzwlW8ZH0nMyYpwqmJ04P7TIdP0,4774
+metaflow/metaflow_version.py,sha256=lOsrhMfHdKXBPS-673sZb48dK2FJhr34rIwLqMFeG9I,7291
 metaflow/monitor.py,sha256=T0NMaBPvXynlJAO_avKtk8OIIRMyEuMAyF8bIp79aZU,5323
 metaflow/multicore_utils.py,sha256=vdTNgczVLODifscUbbveJbuSDOl3Y9pAxhr7sqYiNf4,4760
-metaflow/package.py,sha256=sOvRpnvqVaQa6eR8lwcfb5HYCGqmpYFPm-cLgOEdl04,7377
+metaflow/package.py,sha256=QutDP6WzjwGk1UCKXqBfXa9F10Q--FlRr0J7fwlple0,7399
 metaflow/parameters.py,sha256=l8qnhBG9C4wf_FkXWjq5sapUA6npLdR7pyB0PPQ-KF0,15712
 metaflow/procpoll.py,sha256=U2tE4iK_Mwj2WDyVTx_Uglh6xZ-jixQOo4wrM9OOhxg,2859
 metaflow/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -35,7 +36,7 @@ metaflow/tuple_util.py,sha256=_G5YIEhuugwJ_f6rrZoelMFak3DqAR2tt_5CapS1XTY,830
 metaflow/unbounded_foreach.py,sha256=p184WMbrMJ3xKYHwewj27ZhRUsSj_kw1jlye5gA9xJk,387
 metaflow/util.py,sha256=olAvJK3y1it_k99MhLulTaAJo7OFVt5rnrD-ulIFLCU,13616
 metaflow/vendor.py,sha256=FchtA9tH22JM-eEtJ2c9FpUdMn8sSb1VHuQS56EcdZk,5139
-metaflow/version.py,sha256=kleNZLjxgidWFnFZyqIbSBKOHUU3wMG2-4RRKGiPktc,29
+metaflow/version.py,sha256=rnbTumj_lXNCrZgYvwtVWNwWj0XG3M0qQhmivZyGWG4,29
 metaflow/_vendor/__init__.py,sha256=y_CiwUD3l4eAKvTVDZeqgVujMy31cAM1qjAB-HfI-9s,353
 metaflow/_vendor/typing_extensions.py,sha256=0nUs5p1A_UrZigrAVBoOEM6TxU37zzPDUtiij1ZwpNc,110417
 metaflow/_vendor/zipp.py,sha256=ajztOH-9I7KA_4wqDYygtHa6xUBVZgFpmZ8FE74HHHI,8425
@@ -110,11 +111,11 @@ metaflow/_vendor/v3_6/importlib_metadata/_meta.py,sha256=_F48Hu_jFxkfKWz5wcYS8vO
 metaflow/_vendor/v3_6/importlib_metadata/_text.py,sha256=HCsFksZpJLeTP3NEk_ngrAeXVRRtTrtyh9eOABoRP4A,2166
 metaflow/_vendor/v3_6/importlib_metadata/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 metaflow/client/__init__.py,sha256=1GtQB4Y_CBkzaxg32L1syNQSlfj762wmLrfrDxGi1b8,226
-metaflow/client/core.py,sha256=L4COrMyQgSAWkXIPoXFFQ3OqGnoHfePsBZdJQvShU1E,74162
+metaflow/client/core.py,sha256=jnFTP_XLTLJCHRqaxGKYDn1ZwB_v6a1AAkULlLmJGOk,74171
 metaflow/client/filecache.py,sha256=Wy0yhhCqC1JZgebqi7z52GCwXYnkAqMZHTtxThvwBgM,15229
 metaflow/cmd/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
 metaflow/cmd/configure_cmd.py,sha256=o-DKnUf2FBo_HiMVyoyzQaGBSMtpbEPEdFTQZ0hkU-k,33396
-metaflow/cmd/main_cli.py,sha256=Yunxh0QFUPgGYMe0eeGrbmxtZGVdJPCgkd4P1x_j8B4,2950
+metaflow/cmd/main_cli.py,sha256=E546zT_jYQKysmjwfpEgzZd5QMsyirs28M2s0OPU93E,2966
 metaflow/cmd/tutorials_cmd.py,sha256=8FdlKkicTOhCIDKcBR5b0Oz6giDvS-EMY3o9skIrRqw,5156
 metaflow/cmd/util.py,sha256=jS_0rUjOnGGzPT65fzRLdGjrYAOOLA4jU2S0HJLV0oc,406
 metaflow/cmd/develop/__init__.py,sha256=p1Sy8yU1MEKSrH5ttOWOZvNcI1qYu6J6jghdTHwPgOw,689
@@ -128,7 +129,7 @@ metaflow/datastore/exceptions.py,sha256=r7Ab5FvHIzyFh6kwiptA1lO5nLqWg0xRBoeYGefv
 metaflow/datastore/flow_datastore.py,sha256=kbJcOLYnvPHgJfZ_WWkD9LJSX1PHI1K6f9oVUu08A9U,10235
 metaflow/datastore/inputs.py,sha256=i43dXr2xvgtsgKMO9allgCR18bk80GeayeQFyUTH36w,449
 metaflow/datastore/task_datastore.py,sha256=-aDvRUUEaWq0chiGzCCSYmsTP8T-gjs0eowvy2Xyaug,36167
-metaflow/extension_support/__init__.py,sha256=FF-GdIkxiHSqdwCN5xkW5A0DQn3B8GpnphIFDBpLEr0,49432
+metaflow/extension_support/__init__.py,sha256=2z0c4R8zsVmEFOMGT2Jujsl6xveDVa9KLll7moL58NE,52984
 metaflow/extension_support/_empty_file.py,sha256=HENjnM4uAfeNygxMB_feCCWORFoSat9n_QwzSx2oXPw,109
 metaflow/extension_support/cmd.py,sha256=hk8iBUUINqvKCDxInKgWpum8ThiRZtHSJP7qBASHzl8,5711
 metaflow/extension_support/integrations.py,sha256=AWAh-AZ-vo9IxuAVEjGw3s8p_NMm2DKHYx10oC51gPU,5506
@@ -174,12 +175,11 @@ metaflow/plugins/airflow/sensors/s3_sensor.py,sha256=iDReG-7FKnumrtQg-HY6cCUAAqN
 metaflow/plugins/argo/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 metaflow/plugins/argo/argo_client.py,sha256=MKKhMCbWOPzf6z5zQQiyDRHHkAXcO7ipboDZDqAAvOk,15849
 metaflow/plugins/argo/argo_events.py,sha256=_C1KWztVqgi3zuH57pInaE9OzABc2NnncC-zdwOMZ-w,5909
-metaflow/plugins/argo/argo_workflows.py,sha256=8YajGVnvBjGC22L308p8zMXzvfJP0fI9RsxGl46nO1k,170748
-metaflow/plugins/argo/argo_workflows_cli.py,sha256=X2j_F0xF8-K30ebM4dSLOTteDKXbr-jMN18oMpl5S6Y,36313
+metaflow/plugins/argo/argo_workflows.py,sha256=W8lrMrtJQUCvJbzNyGcT3UVc1CJyMfGw1mrIm7U5ec4,172389
+metaflow/plugins/argo/argo_workflows_cli.py,sha256=E_PyhOtxuS2F8DwhBANsRZCMxpeZ5rfID8eksfSOPm8,37231
 metaflow/plugins/argo/argo_workflows_decorator.py,sha256=yprszMdbE3rBTcEA9VR0IEnPjTprUauZBc4SBb-Q7sA,7878
 metaflow/plugins/argo/argo_workflows_deployer.py,sha256=wSSZtThn_VPvE_Wu6NB1L0Q86LmBJh9g009v_lpvBPM,8125
 metaflow/plugins/argo/capture_error.py,sha256=Ys9dscGrTpW-ZCirLBU0gD9qBM0BjxyxGlUMKcwewQc,1852
-metaflow/plugins/argo/daemon.py,sha256=dJOS_UUISXBYffi3oGVKPwq4Pa4P_nGBGL15piPaPto,1776
 metaflow/plugins/argo/generate_input_paths.py,sha256=loYsI6RFX9LlFsHb7Fe-mzlTTtRdySoOu7sYDy-uXK0,881
 metaflow/plugins/argo/jobset_input_paths.py,sha256=_JhZWngA6p9Q_O2fx3pdzKI0WE-HPRHz_zFvY2pHPTQ,525
 metaflow/plugins/aws/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -291,7 +291,7 @@ metaflow/plugins/metadata/local.py,sha256=YhLJC5zjVJrvQFIyQ92ZBByiUmhCC762RUX7IT
 metaflow/plugins/metadata/service.py,sha256=ihq5F7KQZlxvYwzH_-jyP2aWN_I96i2vp92j_d697s8,20204
 metaflow/plugins/pypi/__init__.py,sha256=0YFZpXvX7HCkyBFglatual7XGifdA1RwC3U4kcizyak,1037
 metaflow/plugins/pypi/bootstrap.py,sha256=LSRjpDhtpvYqCC0lDPIWvZsrHW7PXx626NDzdYjFeVM,5224
-metaflow/plugins/pypi/conda_decorator.py,sha256=sIAfvkjhUfTMK5yLrPOkqxeLUhlC5Buqmc9aITaf22E,15686
+metaflow/plugins/pypi/conda_decorator.py,sha256=fPeXxvmg51oSFTnlguNlcWUIdXHA9OuMnp9ElaxQPFo,15695
 metaflow/plugins/pypi/conda_environment.py,sha256=--q-8lypKupCdGsASpqABNpNqRxtQi6UCDgq8iHDFe4,19476
 metaflow/plugins/pypi/micromamba.py,sha256=67FiIZZz0Kig9EcN7bZLObsE6Z1MFyo4Dp93fd3Grcc,12178
 metaflow/plugins/pypi/pip.py,sha256=7B06mPOs5MvY33xbzPVYZlBr1iKMYaN-n8uulL9zSVg,13649
@@ -345,9 +345,9 @@ metaflow/tutorials/07-worldview/README.md,sha256=5vQTrFqulJ7rWN6r20dhot9lI2sVj9W
 metaflow/tutorials/07-worldview/worldview.ipynb,sha256=ztPZPI9BXxvW1QdS2Tfe7LBuVzvFvv0AToDnsDJhLdE,2237
 metaflow/tutorials/08-autopilot/README.md,sha256=GnePFp_q76jPs991lMUqfIIh5zSorIeWznyiUxzeUVE,1039
 metaflow/tutorials/08-autopilot/autopilot.ipynb,sha256=DQoJlILV7Mq9vfPBGW-QV_kNhWPjS5n6SJLqePjFYLY,3191
-metaflow-2.12.20.dist-info/LICENSE,sha256=nl_Lt5v9VvJ-5lWJDT4ddKAG-VZ-2IaLmbzpgYDz2hU,11343
-metaflow-2.12.20.dist-info/METADATA,sha256=oHu5racxyaljUQCl6cYXqYOHkI02VB5mR298txO9lAk,5906
-metaflow-2.12.20.dist-info/WHEEL,sha256=AHX6tWk3qWuce7vKLrj7lnulVHEdWoltgauo8bgCXgU,109
-metaflow-2.12.20.dist-info/entry_points.txt,sha256=IKwTN1T3I5eJL3uo_vnkyxVffcgnRdFbKwlghZfn27k,57
-metaflow-2.12.20.dist-info/top_level.txt,sha256=v1pDHoWaSaKeuc5fKTRSfsXCKSdW1zvNVmvA-i0if3o,9
-metaflow-2.12.20.dist-info/RECORD,,
+metaflow-2.12.22.dist-info/LICENSE,sha256=nl_Lt5v9VvJ-5lWJDT4ddKAG-VZ-2IaLmbzpgYDz2hU,11343
+metaflow-2.12.22.dist-info/METADATA,sha256=XI-AUDFfekPRYdOciqofewvcuPCSZx7a5lZd58nU2Rw,5906
+metaflow-2.12.22.dist-info/WHEEL,sha256=AHX6tWk3qWuce7vKLrj7lnulVHEdWoltgauo8bgCXgU,109
+metaflow-2.12.22.dist-info/entry_points.txt,sha256=IKwTN1T3I5eJL3uo_vnkyxVffcgnRdFbKwlghZfn27k,57
+metaflow-2.12.22.dist-info/top_level.txt,sha256=v1pDHoWaSaKeuc5fKTRSfsXCKSdW1zvNVmvA-i0if3o,9
+metaflow-2.12.22.dist-info/RECORD,,

metaflow/plugins/argo/daemon.py

@@ -1,59 +0,0 @@
-from collections import namedtuple
-from time import sleep
-from metaflow.metaflow_config import DEFAULT_METADATA
-from metaflow.metaflow_environment import MetaflowEnvironment
-from metaflow.plugins import METADATA_PROVIDERS
-from metaflow._vendor import click
-
-
-class CliState:
-    pass
-
-
-@click.group()
-@click.option("--flow_name", required=True)
-@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_context
-def cli(ctx, flow_name, run_id, tags=None):
-    ctx.obj = CliState()
-    ctx.obj.flow_name = flow_name
-    ctx.obj.run_id = run_id
-    ctx.obj.tags = tags
-    # Use a dummy flow to initialize the environment and metadata service,
-    # as we only need a name for the flow object.
-    flow = namedtuple("DummyFlow", "name")
-    dummyflow = flow(flow_name)
-
-    # Initialize a proper metadata service instance
-    environment = MetaflowEnvironment(dummyflow)
-
-    ctx.obj.metadata = [m for m in METADATA_PROVIDERS if m.TYPE == DEFAULT_METADATA][0](
-        environment, dummyflow, None, None
-    )
-
-
-@cli.command(help="start heartbeat process for a run")
-@click.pass_obj
-def heartbeat(obj):
-    # Try to register a run in case the start task has not taken care of it yet.
-    obj.metadata.register_run_id(obj.run_id, obj.tags)
-    # Start run heartbeat
-    obj.metadata.start_run_heartbeat(obj.flow_name, obj.run_id)
-    # Keepalive loop
-    while True:
-        # Do not pollute daemon logs with anything unnecessary,
-        # as they might be extremely long running.
-        sleep(10)
-
-
-if __name__ == "__main__":
-    cli()