ob-metaflow 2.9.10.1__py2.py3-none-any.whl → 2.10.2.6__py2.py3-none-any.whl

metaflow/plugins/pypi/micromamba.py

@@ -0,0 +1,294 @@
+import json
+import os
+import subprocess
+import tempfile
+
+from metaflow.exception import MetaflowException
+from metaflow.util import which
+
+from .utils import conda_platform
+
+
+class MicromambaException(MetaflowException):
+    headline = "Micromamba ran into an error while setting up environment"
+
+    def __init__(self, error):
+        if isinstance(error, (list,)):
+            error = "\n".join(error)
+        msg = "{error}".format(error=error)
+        super(MicromambaException, self).__init__(msg)
+
+
+class Micromamba(object):
+    def __init__(self):
+        # micromamba is a tiny version of the mamba package manager and comes with
+        # metaflow specific performance enhancements.
+
+        _path_to_hidden_micromamba = os.path.join(
+            os.path.expanduser(os.environ.get("METAFLOW_HOME", "~/.metaflowconfig")),
+            "micromamba",
+        )
+        self.bin = (
+            which(os.environ.get("METAFLOW_PATH_TO_MICROMAMBA") or "micromamba")
+            or which("./micromamba")  # to support remote execution
+            or which("./bin/micromamba")
+            or which(os.path.join(_path_to_hidden_micromamba, "bin/micromamba"))
+        )
+        if self.bin is None:
+            # Install Micromamba on the fly.
+            # TODO: Make this optional at some point.
+            _install_micromamba(_path_to_hidden_micromamba)
+            self.bin = which(os.path.join(_path_to_hidden_micromamba, "bin/micromamba"))
+
+        if self.bin is None:
+            msg = "No installation for *Micromamba* found.\n"
+            msg += "Visit https://mamba.readthedocs.io/en/latest/micromamba-installation.html for installation instructions."
+            raise MetaflowException(msg)
+
+    def solve(self, id_, packages, python, platform):
+        # Performance enhancements
+        # 1. Using zstd compressed repodata index files drops the index download time
+        #    by a factor of 10x - conda-forge/noarch/repodata.json has a
+        #    mean download time of 3.705s ± 2.283s vs 385.1ms ± 57.9ms for
+        #    conda-forge/noarch/repodata.json.zst. Thankfully, now micromamba pulls
+        #    zstd compressed files by default - https://github.com/conda-forge/conda-forge.github.io/issues/1835
+        # 2. Tweak repodata ttl to pull either only once a day or if a solve fails -
+        #    this ensures that repodata is not pulled everytime it becomes dirty by
+        #    default. This can result in an environment that has stale transitive
+        #    dependencies but still correct. (--repodata-ttl 86400 --retry-clean-cache)
+        # 3. Introduce pip as python dependency to resolve pip packages within conda
+        #    environment
+        # 4. Multiple solves can progress at the same time while relying on the same
+        #    index
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            env = {
+                "MAMBA_ADD_PIP_AS_PYTHON_DEPENDENCY": "true",
+                "CONDA_SUBDIR": platform,
+                # "CONDA_UNSATISFIABLE_HINTS_CHECK_DEPTH": "0" # https://github.com/conda/conda/issues/9862
+            }
+            cmd = [
+                "create",
+                "--yes",
+                "--quiet",
+                "--dry-run",
+                "--no-extra-safety-checks",
+                "--repodata-ttl=86400",
+                "--retry-clean-cache",
+                "--prefix=%s/prefix" % tmp_dir,
+            ]
+            # Introduce conda-forge as a default channel
+            for channel in self.info()["channels"] or ["conda-forge"]:
+                cmd.append("--channel=%s" % channel)
+
+            for package, version in packages.items():
+                cmd.append("%s==%s" % (package, version))
+            if python:
+                cmd.append("python==%s" % python)
+            # TODO: Ensure a human readable message is returned when the environment
+            #       can't be resolved for any and all reasons.
+            return [
+                {k: v for k, v in item.items() if k in ["url"]}
+                for item in self._call(cmd, env)["actions"]["LINK"]
+            ]
+
+    def download(self, id_, packages, python, platform):
+        # Unfortunately all the packages need to be catalogued in package cache
+        # because of which this function can't be parallelized
+
+        # Micromamba is painfully slow in determining if many packages are infact
+        # already cached. As a perf heuristic, we check if the environment already
+        # exists to short circuit package downloads.
+        if self.path_to_environment(id_, platform):
+            return
+
+        prefix = "{env_dirs}/{keyword}/{platform}/{id}".format(
+            env_dirs=self.info()["envs_dirs"][0],
+            platform=platform,
+            keyword="metaflow",  # indicates metaflow generated environment
+            id=id_,
+        )
+
+        # Another forced perf heuristic to skip cross-platform downloads.
+        if os.path.exists(f"{prefix}/fake.done"):
+            return
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            env = {
+                "CONDA_SUBDIR": platform,
+            }
+            cmd = [
+                "create",
+                "--yes",
+                "--no-deps",
+                "--download-only",
+                "--safety-checks=disabled",
+                "--no-extra-safety-checks",
+                "--repodata-ttl=86400",
+                "--prefix=%s/prefix" % tmp_dir,
+                "--quiet",
+            ]
+            for package in packages:
+                cmd.append("{url}".format(**package))
+
+            self._call(cmd, env)
+            # Perf optimization to skip cross-platform downloads.
+            if platform != self.platform():
+                os.makedirs(prefix, exist_ok=True) or open(
+                    f"{prefix}/fake.done", "w"
+                ).close()
+            return
+
+    def create(self, id_, packages, python, platform):
+        # create environment only if the platform matches system platform
+        if platform != self.platform() or self.path_to_environment(id_, platform):
+            return
+
+        prefix = "{env_dirs}/{keyword}/{platform}/{id}".format(
+            env_dirs=self.info()["envs_dirs"][0],
+            platform=platform,
+            keyword="metaflow",  # indicates metaflow generated environment
+            id=id_,
+        )
+
+        env = {
+            # use hardlinks when possible, otherwise copy files
+            # disabled for now since it adds to environment creation latencies
+            "CONDA_ALLOW_SOFTLINKS": "0",
+        }
+        cmd = [
+            "create",
+            "--yes",
+            "--no-extra-safety-checks",
+            # "--offline", # not needed since micromamba will first look at the cache
+            "--prefix",  # trick to ensure environments can be created in parallel
+            prefix,
+            "--quiet",
+            "--no-deps",  # important!
+        ]
+        for package in packages:
+            cmd.append("{url}".format(**package))
+        self._call(cmd, env)
+
+    def info(self):
+        return self._call(["config", "list", "-a"])
+
+    def path_to_environment(self, id_, platform=None):
+        if platform is None:
+            platform = self.platform()
+        suffix = "{keyword}/{platform}/{id}".format(
+            platform=platform,
+            keyword="metaflow",  # indicates metaflow generated environment
+            id=id_,
+        )
+        for env in self._call(["env", "list"])["envs"]:
+            # TODO: Check bin/python is available as a heuristic for well formed env
+            if env.endswith(suffix):
+                return env
+
+    def metadata(self, id_, packages, python, platform):
+        # this method unfortunately relies on the implementation detail for
+        # conda environments and has the potential to break all of a sudden.
+        packages_to_filenames = {
+            package["url"]: package["url"].split("/")[-1] for package in packages
+        }
+        directories = self.info()["pkgs_dirs"]
+        # search all package caches for packages
+        metadata = {
+            url: os.path.join(d, file)
+            for url, file in packages_to_filenames.items()
+            for d in directories
+            if os.path.isdir(d)
+            and file in os.listdir(d)
+            and os.path.isfile(os.path.join(d, file))
+        }
+        # set package tarball local paths to None if package tarballs are missing
+        for url in packages_to_filenames:
+            metadata.setdefault(url, None)
+        return metadata
+
+    def interpreter(self, id_):
+        return os.path.join(self.path_to_environment(id_), "bin/python")
+
+    def platform(self):
+        return self.info()["platform"]
+
+    def _call(self, args, env=None):
+        if env is None:
+            env = {}
+        try:
+            result = (
+                subprocess.check_output(
+                    [self.bin] + args,
+                    stderr=subprocess.PIPE,
+                    env={
+                        **os.environ,
+                        # prioritize metaflow-specific env vars
+                        **{k: v for k, v in env.items() if v is not None},
+                        **{
+                            "MAMBA_NO_BANNER": "1",
+                            "MAMBA_JSON": "true",
+                            # play with fire! needed for resolving cross-platform
+                            # environments
+                            "CONDA_SAFETY_CHECKS": "disabled",
+                            # "CONDA_UNSATISFIABLE_HINTS_CHECK_DEPTH": "0",
+                            # Support packages on S3
+                            # "CONDA_ALLOW_NON_CHANNEL_URLS": "1",
+                            "MAMBA_USE_LOCKFILES": "false",
+                        },
+                    },
+                )
+                .decode()
+                .strip()
+            )
+            if result:
+                return json.loads(result)
+            return {}
+        except subprocess.CalledProcessError as e:
+            msg = "command '{cmd}' returned error ({code})\n{stderr}"
+            try:
+                output = json.loads(e.output)
+                err = []
+                for error in output.get("solver_problems", []):
+                    err.append(error)
+                raise MicromambaException(
+                    msg.format(
+                        cmd=" ".join(e.cmd),
+                        code=e.returncode,
+                        output=e.output.decode(),
+                        stderr="\n".join(err),
+                    )
+                )
+            except (TypeError, ValueError) as ve:
+                pass
+            raise MicromambaException(
+                msg.format(
+                    cmd=" ".join(e.cmd),
+                    code=e.returncode,
+                    output=e.output.decode(),
+                    stderr=e.stderr.decode(),
+                )
+            )
+
+
+def _install_micromamba(installation_location):
+    # Unfortunately no 32bit binaries are available for micromamba, which ideally
+    # shouldn't be much of a problem in today's world.
+    platform = conda_platform()
+    try:
+        subprocess.Popen(f"mkdir -p {installation_location}", shell=True).wait()
+        # https://mamba.readthedocs.io/en/latest/micromamba-installation.html#manual-installation
+        result = subprocess.Popen(
+            f"curl -Ls https://micro.mamba.pm/api/micromamba/{platform}/latest | tar -xvj -C {installation_location} bin/micromamba",
+            shell=True,
+            stderr=subprocess.PIPE,
+        )
+        _, err = result.communicate()
+        if result.returncode != 0:
+            raise MicromambaException(
+                f"Micromamba installation '{result.args}' failed:\n{err.decode()}"
+            )
+
+    except subprocess.CalledProcessError as e:
+        raise MicromambaException(
+            "Micromamba installation failed:\n{}".format(e.stderr.decode())
+        )

metaflow/plugins/pypi/pip.py

@@ -0,0 +1,205 @@
+import json
+import os
+import subprocess
+import tempfile
+from itertools import chain, product
+
+from metaflow.exception import MetaflowException
+from metaflow.util import which
+
+from .micromamba import Micromamba
+from .utils import pip_tags
+
+
+class PipException(MetaflowException):
+    headline = "Pip ran into an error while setting up environment"
+
+    def __init__(self, error):
+        if isinstance(error, (list,)):
+            error = "\n".join(error)
+        msg = "{error}".format(error=error)
+        super(PipException, self).__init__(msg)
+
+
+METADATA_FILE = "{prefix}/.pip/metadata"
+INSTALLATION_MARKER = "{prefix}/.pip/id"
+
+# TODO:
+#     1. Support git repositories, local dirs, non-wheel like packages
+#     2. Support protected indices
+
+
+class Pip(object):
+    def __init__(self, micromamba=None):
+        # pip is assumed to be installed inside a conda environment managed by
+        # micromamba. pip commands are executed using `micromamba run --prefix`
+        self.micromamba = micromamba or Micromamba()
+
+    def solve(self, id_, packages, python, platform):
+        prefix = self.micromamba.path_to_environment(id_)
+        if prefix is None:
+            msg = "Unable to locate a Micromamba managed virtual environment\n"
+            msg += "for id {id}".format(id=id_)
+            raise PipException(msg)
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            report = "{tmp_dir}/report.json".format(tmp_dir=tmp_dir)
+            implementations, platforms, abis = zip(
+                *[
+                    (tag.interpreter, tag.platform, tag.abi)
+                    for tag in pip_tags(python, platform)
+                ]
+            )
+            extra_index_urls = self.extra_index_urls(prefix)
+            cmd = [
+                "install",
+                "--dry-run",
+                "--only-binary=:all:",  # only wheels
+                "--upgrade-strategy=only-if-needed",
+                "--target=%s" % tmp_dir,
+                "--report=%s" % report,
+                "--progress-bar=off",
+                "--quiet",
+                *(
+                    chain.from_iterable(
+                        product(["--extra-index-url"], set(extra_index_urls))
+                    )
+                ),
+                *(chain.from_iterable(product(["--abi"], set(abis)))),
+                *(chain.from_iterable(product(["--platform"], set(platforms)))),
+                # *(chain.from_iterable(product(["--implementations"], set(implementations)))),
+            ]
+            for package, version in packages.items():
+                if version.startswith(("<", ">", "!", "~")):
+                    cmd.append(f"{package}{version}")
+                else:
+                    cmd.append(f"{package}=={version}")
+            self._call(prefix, cmd)
+            with open(report, mode="r", encoding="utf-8") as f:
+                return [
+                    {k: v for k, v in item["download_info"].items() if k in ["url"]}
+                    for item in json.load(f)["install"]
+                ]
+
+    def download(self, id_, packages, python, platform):
+        prefix = self.micromamba.path_to_environment(id_)
+        metadata_file = METADATA_FILE.format(prefix=prefix)
+        # download packages only if they haven't ever been downloaded before
+        if os.path.isfile(metadata_file):
+            return
+        metadata = {}
+        implementations, platforms, abis = zip(
+            *[
+                (tag.interpreter, tag.platform, tag.abi)
+                for tag in pip_tags(python, platform)
+            ]
+        )
+        extra_index_urls = self.extra_index_urls(prefix)
+        cmd = [
+            "download",
+            "--no-deps",
+            "--no-index",
+            "--progress-bar=off",
+            #  if packages are present in Pip cache, this will be a local copy
+            "--dest=%s/.pip/wheels" % prefix,
+            "--quiet",
+            *(
+                chain.from_iterable(
+                    product(["--extra-index-url"], set(extra_index_urls))
+                )
+            ),
+            *(chain.from_iterable(product(["--abi"], set(abis)))),
+            *(chain.from_iterable(product(["--platform"], set(platforms)))),
+            # *(chain.from_iterable(product(["--implementations"], set(implementations)))),
+        ]
+        for package in packages:
+            cmd.append("{url}".format(**package))
+            metadata["{url}".format(**package)] = "{prefix}/.pip/wheels/{wheel}".format(
+                prefix=prefix, wheel=package["url"].split("/")[-1]
+            )
+        self._call(prefix, cmd)
+        # write the url to wheel mappings in a magic location
+        with open(metadata_file, "w") as file:
+            file.write(json.dumps(metadata))
+
+    def create(self, id_, packages, python, platform):
+        prefix = self.micromamba.path_to_environment(id_)
+        installation_marker = INSTALLATION_MARKER.format(prefix=prefix)
+        # install packages only if they haven't been installed before
+        if os.path.isfile(installation_marker):
+            return
+        # Pip can't install packages if the underlying virtual environment doesn't
+        # share the same platform
+        if self.micromamba.platform() == platform:
+            cmd = [
+                "install",
+                "--no-compile",
+                "--no-deps",
+                "--no-index",
+                "--progress-bar=off",
+                "--quiet",
+            ]
+            for package in packages:
+                cmd.append("{url}".format(**package))
+            self._call(prefix, cmd)
+        with open(installation_marker, "w") as file:
+            file.write(json.dumps({"id": id_}))
+
+    def metadata(self, id_, packages, python, platform):
+        # read the url to wheel mappings from a magic location
+        prefix = self.micromamba.path_to_environment(id_)
+        metadata_file = METADATA_FILE.format(prefix=prefix)
+        with open(metadata_file, "r") as file:
+            return json.loads(file.read())
+
+    def extra_index_urls(self, prefix):
+        # get extra index urls from Pip conf
+        extra_indices = []
+        for key in [":env:.extra-index-url", "global.extra-index-url"]:
+            try:
+                extras = self._call(prefix, args=["config", "get", key], isolated=False)
+                extra_indices.extend(extras.split(" "))
+            except Exception:
+                # Pip will throw an error when trying to get a config key that does
+                # not exist
+                pass
+        return extra_indices
+
+    def _call(self, prefix, args, env=None, isolated=True):
+        if env is None:
+            env = {}
+        try:
+            return (
+                subprocess.check_output(
+                    [
+                        self.micromamba.bin,
+                        "run",
+                        "--prefix",
+                        prefix,
+                        "pip3",
+                        "--disable-pip-version-check",
+                        "--no-input",
+                        "--no-color",
+                    ]
+                    + (["--isolated"] if isolated else [])
+                    + args,
+                    stderr=subprocess.PIPE,
+                    env={
+                        **os.environ,
+                        # prioritize metaflow-specific env vars
+                        **{"PYTHONNOUSERSITE": "1"},  # no user installation!
+                        **env,
+                    },
+                )
+                .decode()
+                .strip()
+            )
+        except subprocess.CalledProcessError as e:
+            raise PipException(
+                "command '{cmd}' returned error ({code}) {output}\n{stderr}".format(
+                    cmd=" ".join(e.cmd),
+                    code=e.returncode,
+                    output=e.output.decode(),
+                    stderr=e.stderr.decode(),
+                )
+            )

metaflow/plugins/pypi/pypi_decorator.py

@@ -0,0 +1,130 @@
+from metaflow.decorators import FlowDecorator, StepDecorator
+from metaflow.metaflow_environment import InvalidEnvironmentException
+
+
+class PyPIStepDecorator(StepDecorator):
+    """
+    Specifies the PyPI packages for the step.
+
+    Information in this decorator will augment any
+    attributes set in the `@pyi_base` flow-level decorator. Hence,
+    you can use `@pypi_base` to set packages required by all
+    steps and use `@pypi` to specify step-specific overrides.
+
+    Parameters
+    ----------
+    packages : Dict[str, str], default: {}
+        Packages to use for this step. The key is the name of the package
+        and the value is the version to use.
+    python : str, optional
+        Version of Python to use, e.g. '3.7.4'. A default value of None implies
+        that the version used will correspond to the version of the Python interpreter used to start the run.
+    """
+
+    name = "pypi"
+    defaults = {"packages": {}, "python": None, "disabled": None}  # wheels
+
+    def step_init(self, flow, graph, step, decos, environment, flow_datastore, logger):
+        # The init_environment hook for Environment creates the relevant virtual
+        # environments. The step_init hook sets up the relevant state for that hook to
+        # do it's magic.
+
+        self.flow = flow
+        self.step = step
+
+        # Support flow-level decorator
+        if "pypi_base" in self.flow._flow_decorators:
+            super_attributes = self.flow._flow_decorators["pypi_base"][0].attributes
+            self.attributes["packages"] = {
+                **super_attributes["packages"],
+                **self.attributes["packages"],
+            }
+            self.attributes["python"] = (
+                self.attributes["python"] or super_attributes["python"]
+            )
+            self.attributes["disabled"] = (
+                self.attributes["disabled"]
+                if self.attributes["disabled"] is not None
+                else super_attributes["disabled"]
+            )
+
+        # At the moment, @pypi uses a conda environment as a virtual environment. This
+        # is to ensure that we can have a dedicated Python interpreter within the
+        # virtual environment. The conda environment is currently created through
+        # micromamba. As a follow up, we can look into creating a virtualenv using
+        # venv.
+
+        # Currently, @pypi relies on pip for package resolution. We can introduce
+        # support for Poetry in the near future, if desired. Poetry is great for
+        # interactive use cases, but not so much for programmatic use cases like the
+        # one here. We can consider introducing a UX where @pypi is able to consume
+        # poetry.lock files in the future.
+
+        _supported_virtual_envs = ["conda"]  # , "venv"]
+
+        # To placate people who don't want to see a shred of conda in UX, we symlink
+        # --environment=pypi to --environment=conda
+        _supported_virtual_envs.extend(["pypi"])
+
+        # The --environment= requirement ensures that valid virtual environments are
+        # created for every step to execute it, greatly simplifying the @pypi
+        # implementation.
+        if environment.TYPE not in _supported_virtual_envs:
+            raise InvalidEnvironmentException(
+                "@%s decorator requires %s"
+                % (
+                    self.name,
+                    " or ".join(
+                        ["--environment=%s" % env for env in _supported_virtual_envs]
+                    ),
+                )
+            )
+
+
+class PyPIFlowDecorator(FlowDecorator):
+    """
+    Specifies the PyPI packages for all steps of the flow.
+
+    Use `@pypi_base` to set common packages required by all
+    steps and use `@pypi` to specify step-specific overrides.
+    Parameters
+    ----------
+    packages : Dict[str, str], default: {}
+        Packages to use for this flow. The key is the name of the package
+        and the value is the version to use.
+    python : str, optional
+        Version of Python to use, e.g. '3.7.4'. A default value of None implies
+        that the version used will correspond to the version of the Python interpreter used to start the run.
+    """
+
+    name = "pypi_base"
+    defaults = {"packages": {}, "python": None, "disabled": None}
+
+    def flow_init(
+        self, flow, graph, environment, flow_datastore, metadata, logger, echo, options
+    ):
+        from metaflow import decorators
+
+        decorators._attach_decorators(flow, ["pypi"])
+
+        # @pypi uses a conda environment to create a virtual environment.
+        # The conda environment can be created through micromamba.
+        _supported_virtual_envs = ["conda"]
+
+        # To placate people who don't want to see a shred of conda in UX, we symlink
+        # --environment=pypi to --environment=conda
+        _supported_virtual_envs.extend(["pypi"])
+
+        # The --environment= requirement ensures that valid virtual environments are
+        # created for every step to execute it, greatly simplifying the @conda
+        # implementation.
+        if environment.TYPE not in _supported_virtual_envs:
+            raise InvalidEnvironmentException(
+                "@%s decorator requires %s"
+                % (
+                    self.name,
+                    " or ".join(
+                        ["--environment=%s" % env for env in _supported_virtual_envs]
+                    ),
+                )
+            )

metaflow/plugins/pypi/pypi_environment.py

@@ -0,0 +1,7 @@
+from .conda_environment import CondaEnvironment
+
+
+# To placate people who don't want to see a shred of conda in UX, we symlink
+# --environment=pypi to --environment=conda
+class PyPIEnvironment(CondaEnvironment):
+    TYPE = "pypi"