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

metaflow/plugins/pypi/conda_decorator.py

@@ -1,5 +1,3 @@
-import importlib
-import json
 import os
 import platform
 import re
@@ -7,12 +5,9 @@ import sys
 import tempfile
 
 from metaflow.decorators import FlowDecorator, StepDecorator
-from metaflow.extension_support import EXT_PKG
 from metaflow.metadata_provider import MetaDatum
 from metaflow.metaflow_environment import InvalidEnvironmentException
-from metaflow.util import get_metaflow_root
-
-from ...info_file import INFO_FILE
+from metaflow.packaging_sys import ContentType
 
 
 class CondaStepDecorator(StepDecorator):
@@ -45,26 +40,31 @@ class CondaStepDecorator(StepDecorator):
         "python": None,
         "disabled": None,
     }
+
+    _metaflow_home = None
+    _addl_env_vars = None
+
     # To define conda channels for the whole solve, users can specify
     # CONDA_CHANNELS in their environment. For pinning specific packages to specific
     # conda channels, users can specify channel::package as the package name.
 
-    def __init__(self, attributes=None, statically_defined=False):
+    def __init__(self, attributes=None, statically_defined=False, inserted_by=None):
         self._attributes_with_user_values = (
             set(attributes.keys()) if attributes is not None else set()
         )
 
-        super(CondaStepDecorator, self).__init__(attributes, statically_defined)
+        super(CondaStepDecorator, self).__init__(
+            attributes, statically_defined, inserted_by
+        )
 
     def init(self):
-        super(CondaStepDecorator, self).init()
-
         # Support legacy 'libraries=' attribute for the decorator.
         self.attributes["packages"] = {
             **self.attributes["libraries"],
             **self.attributes["packages"],
         }
-        del self.attributes["libraries"]
+        # Keep because otherwise make_decorator_spec will fail
+        self.attributes["libraries"] = {}
         if self.attributes["packages"]:
             self._attributes_with_user_values.add("packages")
 
@@ -152,67 +152,17 @@ class CondaStepDecorator(StepDecorator):
     def runtime_init(self, flow, graph, package, run_id):
         if self.disabled:
             return
-        # Create a symlink to metaflow installed outside the virtual environment.
-        self.metaflow_dir = tempfile.TemporaryDirectory(dir="/tmp")
-        os.symlink(
-            os.path.join(get_metaflow_root(), "metaflow"),
-            os.path.join(self.metaflow_dir.name, "metaflow"),
-        )
-
-        info = os.path.join(get_metaflow_root(), os.path.basename(INFO_FILE))
-        # Symlink the INFO file as well to properly propagate down the Metaflow version
-        if os.path.isfile(info):
-            os.symlink(
-                info, os.path.join(self.metaflow_dir.name, os.path.basename(INFO_FILE))
+        # We need to make all the code package available to the user code in
+        # a temporary directory which will be added to the PYTHONPATH.
+        if self.__class__._metaflow_home is None:
+            # Do this ONCE per flow
+            self.__class__._metaflow_home = tempfile.TemporaryDirectory(dir="/tmp")
+            package.extract_into(
+                self.__class__._metaflow_home.name, ContentType.ALL_CONTENT
+            )
+            self.__class__._addl_env_vars = package.get_post_extract_env_vars(
+                package.package_metadata, self.__class__._metaflow_home.name
             )
-        else:
-            # If there is no info file, we will actually create one in this new
-            # place because we won't be able to properly resolve the EXT_PKG extensions
-            # the same way as outside conda (looking at distributions, etc.). In a
-            # Conda environment, as shown below (where we set self.addl_paths), all
-            # EXT_PKG extensions are PYTHONPATH extensions. Instead of re-resolving,
-            # we use the resolved information that is written out to the INFO file.
-            with open(
-                os.path.join(self.metaflow_dir.name, os.path.basename(INFO_FILE)),
-                mode="wt",
-                encoding="utf-8",
-            ) as f:
-                f.write(
-                    json.dumps(
-                        self.environment.get_environment_info(include_ext_info=True)
-                    )
-                )
-
-        # Support metaflow extensions.
-        self.addl_paths = None
-        try:
-            m = importlib.import_module(EXT_PKG)
-        except ImportError:
-            # No additional check needed because if we are here, we already checked
-            # for other issues when loading at the toplevel.
-            pass
-        else:
-            custom_paths = list(set(m.__path__))
-            # For some reason, at times, unique paths appear multiple times. We
-            # simplify to avoid un-necessary links.
-
-            if len(custom_paths) == 1:
-                # Regular package; we take a quick shortcut here.
-                os.symlink(
-                    custom_paths[0],
-                    os.path.join(self.metaflow_dir.name, EXT_PKG),
-                )
-            else:
-                # This is a namespace package, we therefore create a bunch of
-                # directories so that we can symlink in those separately, and we will
-                # add those paths to the PYTHONPATH for the interpreter. Note that we
-                # don't symlink to the parent of the package because that could end up
-                # including more stuff we don't want
-                self.addl_paths = []
-                for p in custom_paths:
-                    temp_dir = tempfile.mkdtemp(dir=self.metaflow_dir.name)
-                    os.symlink(p, os.path.join(temp_dir, EXT_PKG))
-                    self.addl_paths.append(temp_dir)
 
         # # Also install any environment escape overrides directly here to enable
         # # the escape to work even in non metaflow-created subprocesses
@@ -291,11 +241,17 @@ class CondaStepDecorator(StepDecorator):
         if self.disabled:
             return
         # Ensure local installation of Metaflow is visible to user code
-        python_path = self.metaflow_dir.name
-        if self.addl_paths is not None:
-            addl_paths = os.pathsep.join(self.addl_paths)
-            python_path = os.pathsep.join([addl_paths, python_path])
-        cli_args.env["PYTHONPATH"] = python_path
+        python_path = self.__class__._metaflow_home.name
+        addl_env_vars = {}
+        if self.__class__._addl_env_vars:
+            for key, value in self.__class__._addl_env_vars.items():
+                if key.endswith(":"):
+                    addl_env_vars[key[:-1]] = value
+                elif key == "PYTHONPATH":
+                    addl_env_vars[key] = os.pathsep.join([value, python_path])
+                else:
+                    addl_env_vars[key] = value
+        cli_args.env.update(addl_env_vars)
         if self.interpreter:
             # https://github.com/conda/conda/issues/7707
             # Also ref - https://github.com/Netflix/metaflow/pull/178
@@ -306,7 +262,9 @@ class CondaStepDecorator(StepDecorator):
     def runtime_finished(self, exception):
         if self.disabled:
             return
-        self.metaflow_dir.cleanup()
+        if self.__class__._metaflow_home is not None:
+            self.__class__._metaflow_home.cleanup()
+            self.__class__._metaflow_home = None
 
 
 class CondaFlowDecorator(FlowDecorator):
@@ -339,22 +297,23 @@ class CondaFlowDecorator(FlowDecorator):
         "disabled": None,
     }
 
-    def __init__(self, attributes=None, statically_defined=False):
+    def __init__(self, attributes=None, statically_defined=False, inserted_by=None):
         self._attributes_with_user_values = (
             set(attributes.keys()) if attributes is not None else set()
         )
 
-        super(CondaFlowDecorator, self).__init__(attributes, statically_defined)
+        super(CondaFlowDecorator, self).__init__(
+            attributes, statically_defined, inserted_by
+        )
 
     def init(self):
-        super(CondaFlowDecorator, self).init()
-
         # Support legacy 'libraries=' attribute for the decorator.
         self.attributes["packages"] = {
             **self.attributes["libraries"],
             **self.attributes["packages"],
         }
-        del self.attributes["libraries"]
+        # Keep because otherwise make_decorator_spec will fail
+        self.attributes["libraries"] = {}
         if self.attributes["python"]:
             self.attributes["python"] = str(self.attributes["python"])
 

metaflow/plugins/pypi/conda_environment.py

@@ -17,6 +17,7 @@ from metaflow.debug import debug
 from metaflow.exception import MetaflowException
 from metaflow.metaflow_config import get_pinned_conda_libs
 from metaflow.metaflow_environment import MetaflowEnvironment
+from metaflow.packaging_sys import ContentType
 
 from . import MAGIC_FILE, _datastore_packageroot
 from .utils import conda_platform
@@ -32,8 +33,10 @@ class CondaEnvironmentException(MetaflowException):
 class CondaEnvironment(MetaflowEnvironment):
     TYPE = "conda"
     _filecache = None
+    _force_rebuild = False
 
     def __init__(self, flow):
+        super().__init__(flow)
         self.flow = flow
 
     def set_local_root(self, local_root):
@@ -71,7 +74,7 @@ class CondaEnvironment(MetaflowEnvironment):
         self.logger = make_thread_safe(logger)
 
         # TODO: Wire up logging
-        micromamba = Micromamba(self.logger)
+        micromamba = Micromamba(self.logger, self._force_rebuild)
         self.solvers = {"conda": micromamba, "pypi": Pip(micromamba, self.logger)}
 
     def init_environment(self, echo, only_steps=None):
@@ -107,7 +110,10 @@ class CondaEnvironment(MetaflowEnvironment):
             return (
                 id_,
                 (
-                    self.read_from_environment_manifest([id_, platform, type_])
+                    (
+                        not self._force_rebuild
+                        and self.read_from_environment_manifest([id_, platform, type_])
+                    )
                     or self.write_to_environment_manifest(
                         [id_, platform, type_],
                         self.solvers[type_].solve(id_, **environment),
@@ -153,7 +159,7 @@ class CondaEnvironment(MetaflowEnvironment):
             _meta = copy.deepcopy(local_packages)
             for id_, packages, _, _ in results:
                 for package in packages:
-                    if package.get("path"):
+                    if package.get("path") and not self._force_rebuild:
                         # Cache only those packages that manifest is unaware of
                         local_packages.pop(package["url"], None)
                     else:
@@ -186,7 +192,7 @@ class CondaEnvironment(MetaflowEnvironment):
             storage.save_bytes(
                 list_of_path_and_filehandle,
                 len_hint=len(list_of_path_and_filehandle),
-                # overwrite=True,
+                overwrite=self._force_rebuild,
             )
             for id_, packages, _, platform in results:
                 if id_ in dirty:
@@ -331,7 +337,7 @@ class CondaEnvironment(MetaflowEnvironment):
                     environment[decorator.name] = {
                         k: copy.deepcopy(decorator.attributes[k])
                         for k in decorator.attributes
-                        if k != "disabled"
+                        if k not in ("disabled", "libraries")
                     }
                 else:
                     return {}
@@ -473,7 +479,9 @@ class CondaEnvironment(MetaflowEnvironment):
         files = []
         manifest = self.get_environment_manifest_path()
         if os.path.exists(manifest):
-            files.append((manifest, os.path.basename(manifest)))
+            files.append(
+                (manifest, os.path.basename(manifest), ContentType.OTHER_CONTENT)
+            )
         return files
 
     def bootstrap_commands(self, step_name, datastore_type):

metaflow/plugins/pypi/micromamba.py

@@ -2,6 +2,7 @@ import functools
 import json
 import os
 import re
+import shutil
 import subprocess
 import tempfile
 import time
@@ -30,7 +31,7 @@ _double_equal_match = re.compile("==(?=[<=>!~])")
 
 
 class Micromamba(object):
-    def __init__(self, logger=None):
+    def __init__(self, logger=None, force_rebuild=False):
         # micromamba is a tiny version of the mamba package manager and comes with
         # metaflow specific performance enhancements.
 
@@ -60,6 +61,8 @@ class Micromamba(object):
         # which causes a race condition in case micromamba needs to be installed first.
         self.install_mutex = Lock()
 
+        self.force_rebuild = force_rebuild
+
     @property
     def bin(self) -> str:
         "Defer installing Micromamba until when the binary path is actually requested"
@@ -152,6 +155,11 @@ class Micromamba(object):
             keyword="metaflow",  # indicates metaflow generated environment
             id=id_,
         )
+        # If we are forcing a rebuild of the environment, we make sure to remove existing files beforehand.
+        # This is to ensure that no irrelevant packages get bundled relative to the resolved environment.
+        # NOTE: download always happens before create, so we want to do the cleanup here instead.
+        if self.force_rebuild:
+            shutil.rmtree(self.path_to_environment(id_, platform), ignore_errors=True)
 
         # cheap check
         if os.path.exists(f"{prefix}/fake.done"):

metaflow/plugins/pypi/pip.py

@@ -12,7 +12,7 @@ from metaflow.debug import debug
 from metaflow.exception import MetaflowException
 
 from .micromamba import Micromamba
-from .utils import pip_tags, wheel_tags
+from .utils import pip_tags, wheel_tags, markers_from_platform, conda_platform
 
 
 class PipException(MetaflowException):
@@ -60,6 +60,15 @@ class Pip(object):
         else:
             self.logger = lambda *args, **kwargs: None  # No-op logger if not provided
 
+    def _get_resolved_python_version(self, prefix):
+        try:
+            result = self.micromamba._call(["list", "--prefix", prefix, "--json"])
+            for package in result:
+                if package.get("name") == "python":
+                    return package["version"]
+        except Exception:
+            return None
+
     def solve(self, id_, packages, python, platform):
         prefix = self.micromamba.path_to_environment(id_)
         if prefix is None:
@@ -67,13 +76,19 @@ class Pip(object):
             msg += "for id {id}".format(id=id_)
             raise PipException(msg)
 
+        resolved_python = self._get_resolved_python_version(prefix)
+        if not resolved_python:
+            raise PipException(
+                "Could not determine Python version from conda environment"
+            )
+
         debug.conda_exec("Solving packages for PyPI environment %s" % id_)
         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)
+                    for tag in pip_tags(resolved_python, platform)
                 ]
             )
             custom_index_url, extra_index_urls = self.indices(prefix)
@@ -104,7 +119,17 @@ class Pip(object):
                 else:
                     cmd.append(f"{package}=={version}")
             try:
-                self._call(prefix, cmd)
+                env = {}
+                if conda_platform() != platform:
+                    # cross-platform resolving requires patching the machine and system info for pip to pick up all the relevant packages.
+                    marker_overrides = markers_from_platform(platform)
+                    env = {
+                        "PIP_CUSTOMIZE_OVERRIDES": json.dumps(marker_overrides),
+                        "PYTHONPATH": os.path.join(
+                            os.path.dirname(__file__), "pip_patcher"
+                        ),
+                    }
+                self._call(prefix, cmd, env)
             except PipPackageNotFound as ex:
                 # pretty print package errors
                 raise PipException(
@@ -142,6 +167,13 @@ class Pip(object):
 
     def download(self, id_, packages, python, platform):
         prefix = self.micromamba.path_to_environment(id_)
+
+        resolved_python = self._get_resolved_python_version(prefix)
+        if not resolved_python:
+            raise PipException(
+                "Could not determine Python version from conda environment"
+            )
+
         metadata_file = METADATA_FILE.format(prefix=prefix)
         # download packages only if they haven't ever been downloaded before
         if os.path.isfile(metadata_file):
@@ -192,7 +224,11 @@ class Pip(object):
                     if os.path.isfile(os.path.join(path, f)) and f.endswith(".whl")
                 ]
                 if (
-                    len(set(pip_tags(python, platform)).intersection(wheel_tags(wheel)))
+                    len(
+                        set(pip_tags(resolved_python, platform)).intersection(
+                            wheel_tags(wheel)
+                        )
+                    )
                     == 0
                 ):
                     raise PipException(
@@ -211,7 +247,7 @@ class Pip(object):
         implementations, platforms, abis = zip(
             *[
                 (tag.interpreter, tag.platform, tag.abi)
-                for tag in pip_tags(python, platform)
+                for tag in pip_tags(resolved_python, platform)
             ]
         )
 

metaflow/plugins/pypi/pypi_decorator.py

@@ -24,12 +24,12 @@ class PyPIStepDecorator(StepDecorator):
     name = "pypi"
     defaults = {"packages": {}, "python": None, "disabled": None}  # wheels
 
-    def __init__(self, attributes=None, statically_defined=False):
+    def __init__(self, attributes=None, statically_defined=False, inserted_by=None):
         self._attributes_with_user_values = (
             set(attributes.keys()) if attributes is not None else set()
         )
 
-        super().__init__(attributes, statically_defined)
+        super().__init__(attributes, statically_defined, inserted_by)
 
     def step_init(self, flow, graph, step, decos, environment, flow_datastore, logger):
         # The init_environment hook for Environment creates the relevant virtual
@@ -128,12 +128,12 @@ class PyPIFlowDecorator(FlowDecorator):
     name = "pypi_base"
     defaults = {"packages": {}, "python": None, "disabled": None}
 
-    def __init__(self, attributes=None, statically_defined=False):
+    def __init__(self, attributes=None, statically_defined=False, inserted_by=None):
         self._attributes_with_user_values = (
             set(attributes.keys()) if attributes is not None else set()
         )
 
-        super().__init__(attributes, statically_defined)
+        super().__init__(attributes, statically_defined, inserted_by)
 
     def flow_init(
         self, flow, graph, environment, flow_datastore, metadata, logger, echo, options

metaflow/plugins/pypi/utils.py

@@ -41,6 +41,28 @@ def conda_platform():
             return "osx-64"
 
 
+def markers_from_platform(platform):
+    plat, mach = platform.split("-")
+
+    platform_system = {"osx": "Darwin", "linux": "Linux"}.get(plat)
+    platform_machine = {
+        "32": "x86",
+        "64": "x86_64",
+        "arm64": "aarch64",
+        "aarch64": "aarch64",
+    }.get(mach)
+
+    markers = {
+        k: v
+        for k, v in {
+            "platform_system": platform_system,
+            "platform_machine": platform_machine,
+        }.items()
+        if v is not None
+    }
+    return markers
+
+
 def wheel_tags(wheel):
     _, _, _, tags = parse_wheel_filename(wheel)
     return list(tags)

metaflow/plugins/secrets/__init__.py

@@ -9,3 +9,6 @@ class SecretsProvider(abc.ABC):
     def get_secret_as_dict(self, secret_id, options={}, role=None) -> Dict[str, str]:
         """Retrieve the secret from secrets backend, and return a dictionary of
         environment variables."""
+
+
+from .secrets_func import get_secret