ob-metaflow 2.11.13.1__py2.py3-none-any.whl → 2.19.7.1rc0__py2.py3-none-any.whl

metaflow/metaflow_environment.py

@@ -1,15 +1,17 @@
+import json
 import os
 import platform
 import sys
 
 from .util import get_username
 from . import metaflow_version
+from . import metaflow_git
 from metaflow.exception import MetaflowException
 from metaflow.extension_support import dump_module_info
-from metaflow.mflog import BASH_MFLOG
-from . import R
+from metaflow.mflog import BASH_MFLOG, BASH_FLUSH_LOGS
+from metaflow.package import MetaflowPackage
 
-version_cache = None
+from . import R
 
 
 class InvalidEnvironmentException(MetaflowException):
@@ -50,8 +52,36 @@ class MetaflowEnvironment(object):
 
     def add_to_package(self):
         """
-        A list of tuples (file, arcname) to add to the job package.
-        `arcname` is an alternative name for the file in the job package.
+        Called to add custom files needed for this environment. This hook will be
+        called in the `MetaflowPackage` class where metaflow compiles the code package
+        tarball. This hook can return one of two things (the first is for backwards
+        compatibility -- move to the second):
+          - a generator yielding a tuple of `(file_path, arcname)` to add files to
+            the code package. `file_path` is the path to the file on the local filesystem
+            and `arcname` is the path relative to the packaged code.
+          - a generator yielding a tuple of `(content, arcname, type)` where:
+            - type is one of
+            ContentType.{USER_CONTENT, CODE_CONTENT, MODULE_CONTENT, OTHER_CONTENT}
+            - for USER_CONTENT:
+              - the file will be included relative to the directory containing the
+                user's flow file.
+              - content: path to the file to include
+              - arcname: path relative to the directory containing the user's flow file
+            - for CODE_CONTENT:
+              - the file will be included relative to the code directory in the package.
+                This will be the directory containing `metaflow`.
+              - content: path to the file to include
+              - arcname: path relative to the code directory in the package
+            - for MODULE_CONTENT:
+              - the module will be added to the code package as a python module. It will
+                be accessible as usual (import <module_name>)
+              - content: name of the module
+              - arcname: None (ignored)
+            - for OTHER_CONTENT:
+              - the file will be included relative to any other configuration/metadata
+                files for the flow
+              - content: path to the file to include
+              - arcname: path relative to the config directory in the package
         """
         return []
 
@@ -89,10 +119,18 @@ class MetaflowEnvironment(object):
         It should work silently if everything goes well.
         """
         if datastore_type == "s3":
-            return (
-                '%s -m awscli ${METAFLOW_S3_ENDPOINT_URL:+--endpoint-url=\\"${METAFLOW_S3_ENDPOINT_URL}\\"} '
-                + "s3 cp %s job.tar"
-            ) % (self._python(), code_package_url)
+            from .plugins.aws.aws_utils import parse_s3_full_path
+
+            bucket, s3_object = parse_s3_full_path(code_package_url)
+            # NOTE: the script quoting is extremely sensitive due to the way shlex.split operates and this being inserted
+            # into a quoted command elsewhere.
+            # NOTE: Reason for the extra conditionals in the script are because
+            # Boto3 does not play well with passing None or an empty string to endpoint_url
+            return "{python} -c '{script}'".format(
+                python=self._python(),
+                script='import boto3, os; ep=os.getenv(\\"METAFLOW_S3_ENDPOINT_URL\\"); boto3.client(\\"s3\\", **({\\"endpoint_url\\":ep} if ep else {})).download_file(\\"%s\\", \\"%s\\", \\"job.tar\\")'
+                % (bucket, s3_object),
+            )
         elif datastore_type == "azure":
             from .plugins.azure.azure_utils import parse_azure_full_path
 
@@ -119,54 +157,98 @@ class MetaflowEnvironment(object):
             )
 
     def _get_install_dependencies_cmd(self, datastore_type):
-        cmds = ["%s -m pip install requests" % self._python()]
-        if datastore_type == "s3":
-            cmds.append("%s -m pip install awscli boto3" % self._python())
-        elif datastore_type == "azure":
-            cmds.append(
-                "%s -m pip install azure-identity azure-storage-blob simple-azure-blob-downloader -qqq"
-                % self._python()
+        base_cmd = "{} -m pip install -qqq --no-compile --no-cache-dir --disable-pip-version-check".format(
+            self._python()
+        )
+
+        datastore_packages = {
+            "s3": ["boto3"],
+            "azure": [
+                "azure-identity",
+                "azure-storage-blob",
+                "azure-keyvault-secrets",
+                "simple-azure-blob-downloader",
+            ],
+            "gs": [
+                "google-cloud-storage",
+                "google-auth",
+                "simple-gcp-object-downloader",
+                "google-cloud-secret-manager",
+                "packaging",
+            ],
+        }
+
+        if datastore_type not in datastore_packages:
+            raise NotImplementedError(
+                "Unknown datastore type: {}".format(datastore_type)
             )
-        elif datastore_type == "gs":
-            cmds.append(
-                "%s -m pip install google-cloud-storage google-auth simple-gcp-object-downloader google-cloud-secret-manager -qqq"
-                % self._python()
+
+        cmd = "{} {}".format(
+            base_cmd, " ".join(datastore_packages[datastore_type] + ["requests"])
+        )
+        # skip pip installs if we know that packages might already be available
+        return "if [ -z $METAFLOW_SKIP_INSTALL_DEPENDENCIES ]; then {}; fi".format(cmd)
+
+    def get_package_commands(
+        self, code_package_url, datastore_type, code_package_metadata=None
+    ):
+        # HACK: We want to keep forward compatibility with compute layers so that
+        # they can still call get_package_commands and NOT pass any metadata. If
+        # there is no additional information, we *assume* that it is the default
+        # used.
+        if code_package_metadata is None:
+            code_package_metadata = json.dumps(
+                {
+                    "version": 0,
+                    "archive_format": "tgz",
+                    "mfcontent_version": 1,
+                }
             )
-        else:
-            raise NotImplementedError(
-                "We don't know how to generate an install dependencies cmd for datastore %s"
-                % datastore_type
+
+        extra_exports = []
+        for k, v in MetaflowPackage.get_post_extract_env_vars(
+            code_package_metadata, dest_dir="$(pwd)"
+        ).items():
+            if k.endswith(":"):
+                # If the value ends with a colon, we override the existing value
+                extra_exports.append("export %s=%s" % (k[:-1], v))
+            else:
+                extra_exports.append(
+                    "export %s=%s:$(printenv %s)" % (k, v.replace('"', '\\"'), k)
+                )
+
+        cmds = (
+            [
+                BASH_MFLOG,
+                BASH_FLUSH_LOGS,
+                "mflog 'Setting up task environment.'",
+                self._get_install_dependencies_cmd(datastore_type),
+                "mkdir metaflow",
+                "cd metaflow",
+                "mkdir .metaflow",  # mute local datastore creation log
+                "i=0; while [ $i -le 5 ]; do "
+                "mflog 'Downloading code package...'; "
+                + self._get_download_code_package_cmd(code_package_url, datastore_type)
+                + " && mflog 'Code package downloaded.' && break; "
+                "sleep 10; i=$((i+1)); "
+                "done",
+                "if [ $i -gt 5 ]; then "
+                "mflog 'Failed to download code package from %s "
+                "after 6 tries. Exiting...' && exit 1; "
+                "fi" % code_package_url,
+            ]
+            + MetaflowPackage.get_extract_commands(
+                code_package_metadata, "job.tar", dest_dir="."
             )
-        return " && ".join(cmds)
-
-    def get_package_commands(self, code_package_url, datastore_type):
-        cmds = [
-            BASH_MFLOG,
-            "mflog 'Setting up task environment.'",
-            self._get_install_dependencies_cmd(datastore_type),
-            "mkdir metaflow",
-            "cd metaflow",
-            "mkdir .metaflow",  # mute local datastore creation log
-            "i=0; while [ $i -le 5 ]; do "
-            "mflog 'Downloading code package...'; "
-            + self._get_download_code_package_cmd(code_package_url, datastore_type)
-            + " && mflog 'Code package downloaded.' && break; "
-            "sleep 10; i=$((i+1)); "
-            "done",
-            "if [ $i -gt 5 ]; then "
-            "mflog 'Failed to download code package from %s "
-            "after 6 tries. Exiting...' && exit 1; "
-            "fi" % code_package_url,
-            "TAR_OPTIONS='--warning=no-timestamp' tar xf job.tar",
-            "mflog 'Task is starting.'",
-        ]
+            + extra_exports
+            + [
+                "mflog 'Task is starting.'",
+                "flush_mflogs",
+            ]
+        )
         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
@@ -180,8 +262,12 @@ 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])),
+            # Add git info
+            **metaflow_git.get_repository_info(
+                path=os.path.dirname(os.path.abspath(sys.argv[0]))
+            ),
         }
         if R.use_r():
             env["metaflow_r_version"] = R.metaflow_r_version()
@@ -191,7 +277,7 @@ class MetaflowEnvironment(object):
             # Information about extension modules (to load them in the proper order)
             ext_key, ext_val = dump_module_info()
             env[ext_key] = ext_val
-        return env
+        return {k: v for k, v in env.items() if v is not None and v != ""}
 
     def executable(self, step_name, default=None):
         if default is not None:

metaflow/metaflow_git.py

@@ -0,0 +1,115 @@
+#!/usr/bin/env python
+"""Get git repository information for the package
+
+Functions to retrieve git repository details like URL, branch name, 
+and commit SHA for Metaflow code provenance tracking.
+"""
+
+import os
+import subprocess
+from typing import Dict, List, Optional, Tuple, Union
+
+# Cache for git information to avoid repeated subprocess calls
+_git_info_cache = None
+
+__all__ = ("get_repository_info",)
+
+
+def _call_git(
+    args: List[str], path=Union[str, os.PathLike]
+) -> Tuple[Optional[str], Optional[int], bool]:
+    """
+    Call git with provided args.
+
+    Returns
+    -------
+        tuple : Tuple containing
+            (stdout, exitcode, failure) of the call
+    """
+    try:
+        result = subprocess.run(
+            ["git", *args],
+            cwd=path,
+            capture_output=True,
+            text=True,
+            check=False,
+        )
+        return result.stdout.strip(), result.returncode, False
+    except (OSError, subprocess.SubprocessError):
+        # Covers subprocess timeouts and other errors which would not lead to an exit code
+        return None, None, True
+
+
+def _get_repo_url(path: Union[str, os.PathLike]) -> Optional[str]:
+    """Get the repository URL from git config"""
+    stdout, returncode, _failed = _call_git(
+        ["config", "--get", "remote.origin.url"], path
+    )
+    if returncode == 0:
+        url = stdout
+        # Convert SSH URLs to HTTPS for clickable links
+        if url.startswith("git@"):
+            parts = url.split(":", 1)
+            if len(parts) == 2:
+                domain = parts[0].replace("git@", "")
+                repo_path = parts[1]
+                url = f"https://{domain}/{repo_path}"
+        return url
+    return None
+
+
+def _get_branch_name(path: Union[str, os.PathLike]) -> Optional[str]:
+    """Get the current git branch name"""
+    stdout, returncode, _failed = _call_git(["rev-parse", "--abbrev-ref", "HEAD"], path)
+    return stdout if returncode == 0 else None
+
+
+def _get_commit_sha(path: Union[str, os.PathLike]) -> Optional[str]:
+    """Get the current git commit SHA"""
+    stdout, returncode, _failed = _call_git(["rev-parse", "HEAD"], path)
+    return stdout if returncode == 0 else None
+
+
+def _is_in_git_repo(path: Union[str, os.PathLike]) -> bool:
+    """Check if we're currently in a git repository"""
+    stdout, returncode, _failed = _call_git(
+        ["rev-parse", "--is-inside-work-tree"], path
+    )
+    return returncode == 0 and stdout == "true"
+
+
+def _has_uncommitted_changes(path: Union[str, os.PathLike]) -> Optional[bool]:
+    """Check if the git repository has uncommitted changes"""
+    _stdout, returncode, failed = _call_git(
+        ["diff-index", "--quiet", "HEAD", "--"], path
+    )
+    if failed:
+        return None
+    return returncode != 0
+
+
+def get_repository_info(path: Union[str, os.PathLike]) -> Dict[str, Union[str, bool]]:
+    """Get git repository information for a path
+
+    Returns:
+        dict: Dictionary containing:
+            repo_url: Repository URL (converted to HTTPS if from SSH)
+            branch_name: Current branch name
+            commit_sha: Current commit SHA
+            has_uncommitted_changes: Boolean indicating if there are uncommitted changes
+    """
+    global _git_info_cache
+
+    if _git_info_cache is not None:
+        return _git_info_cache
+
+    _git_info_cache = {}
+    if _is_in_git_repo(path):
+        _git_info_cache = {
+            "repo_url": _get_repo_url(path),
+            "branch_name": _get_branch_name(path),
+            "commit_sha": _get_commit_sha(path),
+            "has_uncommitted_changes": _has_uncommitted_changes(path),
+        }
+
+    return _git_info_cache

metaflow/metaflow_profile.py

@@ -2,6 +2,24 @@ import time
 
 from contextlib import contextmanager
 
+from .metaflow_config import PROFILE_FROM_START
+
+init_time = None
+
+
+if PROFILE_FROM_START:
+
+    def from_start(msg: str):
+        global init_time
+        if init_time is None:
+            init_time = time.time()
+        print("From start: %s took %dms" % (msg, int((time.time() - init_time) * 1000)))
+
+else:
+
+    def from_start(_msg: str):
+        pass
+
 
 @contextmanager
 def profile(label, stats_dict=None):

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.meta_files import read_info_file
+
+
+# True/False correspond to the value `public`` in get_version
+_version_cache = {True: None, False: None}
 
 __all__ = ("get_version",)
 
@@ -23,11 +27,11 @@ if name == "nt":
         """find the path to the git executable on Windows"""
         # first see if git is in the path
         try:
-            check_output(["where", "/Q", "git"])
+            subprocess.check_output(["where", "/Q", "git"])
             # if this command succeeded, git is in the path
             return "git"
         # catch the exception thrown if git was not found
-        except CalledProcessError:
+        except subprocess.CalledProcessError:
             pass
         # There are several locations where git.exe may be hiding
         possible_locations = []
@@ -57,87 +61,167 @@ 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):
-    """Tracks the version number.
+def make_public_version(version_string):
+    """
+    Takes a complex version string and returns a public, PEP 440-compliant version.
+    It removes local version identifiers (+...) and development markers (-...).
+    """
+    base_version = version_string.split("+", 1)[0]
+    public_version = base_version.split("-", 1)[0]
+    return public_version
 
-    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.
 
-    If the script is located within an active git repository,
-    git-describe is used to get the version information.
+def get_version(public=False):
+    """Tracks the version number.
 
-    Otherwise, the version logged by package installer is returned.
+    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 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.
+    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, 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:
+        # If we have a version from the INFO file, use it directly.
+        # However, if we are asked for a public version, we parse it to make sure
+        # that no local information is included.
+        if public:
+            version = make_public_version(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