ob-metaflow 2.11.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 import MetaDatum
+from metaflow.metadata_provider import MetaDatum
 from metaflow.metaflow_environment import InvalidEnvironmentException
-from metaflow.util import get_metaflow_root
-
-from ... import INFO_FILE
+from metaflow.packaging_sys import ContentType
 
 
 class CondaStepDecorator(StepDecorator):
@@ -45,19 +40,36 @@ 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):
-        super(CondaStepDecorator, self).__init__(attributes, statically_defined)
+    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, inserted_by
+        )
 
+    def init(self):
         # 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")
+
+    def is_attribute_user_defined(self, name):
+        return name in self._attributes_with_user_values
 
     def step_init(self, flow, graph, step, decos, environment, flow_datastore, logger):
         # The init_environment hook for Environment creates the relevant virtual
@@ -71,11 +83,16 @@ class CondaStepDecorator(StepDecorator):
 
         # Support flow-level decorator.
         if "conda_base" in self.flow._flow_decorators:
-            super_attributes = self.flow._flow_decorators["conda_base"][0].attributes
+            conda_base = self.flow._flow_decorators["conda_base"][0]
+            super_attributes = conda_base.attributes
             self.attributes["packages"] = {
                 **super_attributes["packages"],
                 **self.attributes["packages"],
             }
+            self._attributes_with_user_values.update(
+                conda_base._attributes_with_user_values
+            )
+
             self.attributes["python"] = (
                 self.attributes["python"] or super_attributes["python"]
             )
@@ -100,6 +117,10 @@ class CondaStepDecorator(StepDecorator):
         # --environment=pypi to --environment=conda
         _supported_virtual_envs.extend(["pypi"])
 
+        # TODO: Hardcoded for now to support the fast bakery environment.
+        # We should introduce a more robust mechanism for appending supported environments, for example from within extensions.
+        _supported_virtual_envs.extend(["fast-bakery"])
+
         # The --environment= requirement ensures that valid virtual environments are
         # created for every step to execute it, greatly simplifying the @conda
         # implementation.
@@ -131,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
@@ -206,7 +177,8 @@ class CondaStepDecorator(StepDecorator):
         self.interpreter = (
             self.environment.interpreter(self.step)
             if not any(
-                decorator.name in ["batch", "kubernetes"]
+                decorator.name
+                in ["batch", "kubernetes", "nvidia", "snowpark", "slurm", "nvct"]
                 for decorator in next(
                     step for step in self.flow if step.name == self.step
                 ).decorators
@@ -269,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
@@ -284,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):
@@ -317,21 +297,38 @@ class CondaFlowDecorator(FlowDecorator):
         "disabled": None,
     }
 
-    def __init__(self, attributes=None, statically_defined=False):
-        super(CondaFlowDecorator, self).__init__(attributes, statically_defined)
+    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, inserted_by
+        )
 
+    def init(self):
         # 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"])
 
+    def is_attribute_user_defined(self, name):
+        return name in self._attributes_with_user_values
+
     def flow_init(
         self, flow, graph, environment, flow_datastore, metadata, logger, echo, options
     ):
+        # NOTE: Important for extensions implementing custom virtual environments.
+        # Without this steps will not have an implicit conda step decorator on them unless the environment adds one in its decospecs.
+        from metaflow import decorators
+
+        decorators._attach_decorators(flow, ["conda"])
+
         # @conda uses a conda environment to create a virtual environment.
         # The conda environment can be created through micromamba.
         _supported_virtual_envs = ["conda"]
@@ -340,6 +337,10 @@ class CondaFlowDecorator(FlowDecorator):
         # --environment=pypi to --environment=conda
         _supported_virtual_envs.extend(["pypi"])
 
+        # TODO: Hardcoded for now to support the fast bakery environment.
+        # We should introduce a more robust mechanism for appending supported environments, for example from within extensions.
+        _supported_virtual_envs.extend(["fast-bakery"])
+
         # The --environment= requirement ensures that valid virtual environments are
         # created for every step to execute it, greatly simplifying the @conda
         # implementation.

metaflow/plugins/pypi/conda_environment.py

@@ -5,21 +5,19 @@ import functools
 import io
 import json
 import os
-import sys
 import tarfile
-import time
-from concurrent.futures import ThreadPoolExecutor
+import threading
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from functools import wraps
 from hashlib import sha256
 from io import BufferedIOBase, BytesIO
-from itertools import chain
 from urllib.parse import unquote, urlparse
 
-import requests
-
+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.metaflow_profile import profile
+from metaflow.packaging_sys import ContentType
 
 from . import MAGIC_FILE, _datastore_packageroot
 from .utils import conda_platform
@@ -35,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):
@@ -48,9 +48,8 @@ class CondaEnvironment(MetaflowEnvironment):
         # Apply conda decorator to manage the task execution lifecycle.
         return ("conda",) + super().decospecs()
 
-    def validate_environment(self, echo, datastore_type):
+    def validate_environment(self, logger, datastore_type):
         self.datastore_type = datastore_type
-        self.echo = echo
 
         # Avoiding circular imports.
         from metaflow.plugins import DATASTORES
@@ -62,10 +61,23 @@ class CondaEnvironment(MetaflowEnvironment):
         from .micromamba import Micromamba
         from .pip import Pip
 
-        micromamba = Micromamba()
-        self.solvers = {"conda": micromamba, "pypi": Pip(micromamba)}
+        print_lock = threading.Lock()
+
+        def make_thread_safe(func):
+            @wraps(func)
+            def wrapper(*args, **kwargs):
+                with print_lock:
+                    return func(*args, **kwargs)
+
+            return wrapper
+
+        self.logger = make_thread_safe(logger)
 
-    def init_environment(self, echo):
+        # TODO: Wire up logging
+        micromamba = Micromamba(self.logger, self._force_rebuild)
+        self.solvers = {"conda": micromamba, "pypi": Pip(micromamba, self.logger)}
+
+    def init_environment(self, echo, only_steps=None):
         # The implementation optimizes for latency to ensure as many operations can
         # be turned into cheap no-ops as feasible. Otherwise, we focus on maintaining
         # a balance between latency and maintainability of code without re-implementing
@@ -77,6 +89,8 @@ class CondaEnvironment(MetaflowEnvironment):
         def environments(type_):
             seen = set()
             for step in self.flow:
+                if only_steps and step.name not in only_steps:
+                    continue
                 environment = self.get_environment(step)
                 if type_ in environment and environment["id_"] not in seen:
                     seen.add(environment["id_"])
@@ -96,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),
@@ -107,6 +124,11 @@ class CondaEnvironment(MetaflowEnvironment):
             )
 
         def cache(storage, results, type_):
+            debug.conda_exec(
+                "Caching packages for %s environments %s"
+                % (type_, [result[0] for result in results])
+            )
+
             def _path(url, local_path):
                 # Special handling for VCS packages
                 if url.startswith("git+"):
@@ -137,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:
@@ -148,6 +170,9 @@ class CondaEnvironment(MetaflowEnvironment):
                 (
                     package["path"],
                     # Lazily fetch package from the interweb if needed.
+                    # TODO: Depending on the len_hint, the package might be downloaded from
+                    #       the interweb prematurely. save_bytes needs to be adjusted to handle
+                    #       this scenario.
                     LazyOpen(
                         package["local_path"],
                         "rb",
@@ -156,41 +181,131 @@ class CondaEnvironment(MetaflowEnvironment):
                 )
                 for url, package in local_packages.items()
             ]
+            debug.conda_exec(
+                "Caching %s new packages to the datastore for %s environment %s"
+                % (
+                    len(list_of_path_and_filehandle),
+                    type_,
+                    [result[0] for result in results],
+                )
+            )
             storage.save_bytes(
                 list_of_path_and_filehandle,
                 len_hint=len(list_of_path_and_filehandle),
+                overwrite=self._force_rebuild,
             )
             for id_, packages, _, platform in results:
                 if id_ in dirty:
                     self.write_to_environment_manifest([id_, platform, type_], packages)
 
-        # First resolve environments through Conda, before PyPI.
-        echo("Bootstrapping virtual environment(s) ...")
-        for solver in ["conda", "pypi"]:
-            with ThreadPoolExecutor() as executor:
-                results = list(
-                    executor.map(lambda x: solve(*x, solver), environments(solver))
+            debug.conda_exec("Finished caching packages.")
+
+        storage = None
+        if self.datastore_type not in ["local"]:
+            # Initialize storage for caching if using a remote datastore
+            storage = self.datastore(_datastore_packageroot(self.datastore, echo))
+
+        self.logger("Bootstrapping virtual environment(s) ...")
+        # Sequence of operations:
+        #  1. Start all conda solves in parallel
+        #  2. Download conda packages sequentially
+        #  3. Create and cache conda environments in parallel
+        #  4. Start PyPI solves in parallel after each conda environment is created
+        #  5. Download PyPI packages sequentially
+        #  6. Create and cache PyPI environments in parallel
+        with ThreadPoolExecutor() as executor:
+            # Start all conda solves in parallel
+            debug.conda_exec("Solving packages for Conda environments..")
+            conda_solve_futures = [
+                executor.submit(lambda x: solve(*x, "conda"), env)
+                for env in environments("conda")
+            ]
+            conda_create_futures = []
+
+            pypi_envs = {env[0]: env for env in environments("pypi")}
+            pypi_solve_futures = []
+            pypi_create_futures = []
+
+            cache_futures = []
+            # Process conda results sequentially for downloads
+            for future in as_completed(conda_solve_futures):
+                result = future.result()
+                # Sequential conda download
+                debug.conda_exec(
+                    "Downloading packages for Conda environment %s" % result[0]
                 )
-            _ = list(map(lambda x: self.solvers[solver].download(*x), results))
-            with ThreadPoolExecutor() as executor:
-                _ = list(
-                    executor.map(lambda x: self.solvers[solver].create(*x), results)
+                self.solvers["conda"].download(*result)
+                # Parallel conda create and cache
+                conda_create_future = executor.submit(
+                    self.solvers["conda"].create, *result
                 )
-            if self.datastore_type not in ["local"]:
-                # Cache packages only when a remote datastore is in play.
-                storage = self.datastore(
-                    _datastore_packageroot(self.datastore, self.echo)
+                if storage:
+                    cache_futures.append(
+                        executor.submit(cache, storage, [result], "conda")
+                    )
+
+                # Queue PyPI solve to start after conda create
+                if result[0] in pypi_envs:
+                    debug.conda_exec(
+                        "Solving packages for PyPI environment %s" % result[0]
+                    )
+                    # solve pypi envs uniquely
+                    pypi_env = pypi_envs.pop(result[0])
+
+                    def pypi_solve(env):
+                        conda_create_future.result()  # Wait for conda create
+                        return solve(*env, "pypi")
+
+                    pypi_solve_futures.append(executor.submit(pypi_solve, pypi_env))
+                else:
+                    # add conda create future to the generic list
+                    conda_create_futures.append(conda_create_future)
+
+            # Process PyPI results sequentially for downloads
+            for solve_future in as_completed(pypi_solve_futures):
+                result = solve_future.result()
+                # Sequential PyPI download
+                debug.conda_exec(
+                    "Downloading packages for PyPI environment %s" % result[0]
+                )
+                self.solvers["pypi"].download(*result)
+                # Parallel PyPI create and cache
+                pypi_create_futures.append(
+                    executor.submit(self.solvers["pypi"].create, *result)
                 )
-                cache(storage, results, solver)
-        echo("Virtual environment(s) bootstrapped!")
+                if storage:
+                    cache_futures.append(
+                        executor.submit(cache, storage, [result], "pypi")
+                    )
+
+            # Raise exceptions for conda create
+            debug.conda_exec("Checking results for Conda create..")
+            for future in as_completed(conda_create_futures):
+                future.result()
+
+            # Raise exceptions for pypi create
+            debug.conda_exec("Checking results for PyPI create..")
+            for future in as_completed(pypi_create_futures):
+                future.result()
+
+            # Raise exceptions for caching
+            debug.conda_exec("Checking results for caching..")
+            for future in as_completed(cache_futures):
+                # check for result in order to raise any exceptions.
+                future.result()
+
+        self.logger("Virtual environment(s) bootstrapped!")
 
     def executable(self, step_name, default=None):
-        step = next(step for step in self.flow if step.name == step_name)
+        step = next((step for step in self.flow if step.name == step_name), None)
+        if step is None:
+            # requesting internal steps e.g. _parameters
+            return super().executable(step_name, default)
         id_ = self.get_environment(step).get("id_")
         if id_:
             # bootstrap.py is responsible for ensuring the validity of this executable.
             # -s is important! Can otherwise leak packages to other environments.
-            return os.path.join("linux-64", id_, "bin/python -s")
+            return os.path.join("$MF_ARCH", id_, "bin/python -s")
         else:
             # for @conda/@pypi(disabled=True).
             return super().executable(step_name, default)
@@ -220,9 +335,9 @@ class CondaEnvironment(MetaflowEnvironment):
                 disabled = decorator.attributes["disabled"]
                 if not disabled or str(disabled).lower() == "false":
                     environment[decorator.name] = {
-                        k: decorator.attributes[k]
+                        k: copy.deepcopy(decorator.attributes[k])
                         for k in decorator.attributes
-                        if k != "disabled"
+                        if k not in ("disabled", "libraries")
                     }
                 else:
                     return {}
@@ -263,14 +378,21 @@ class CondaEnvironment(MetaflowEnvironment):
         # 5. All resolved packages (Conda or PyPI) are cached
         # 6. PyPI packages are only installed for local platform
 
-        # Resolve `linux-64` Conda environments if @batch or @kubernetes are in play
         target_platform = conda_platform()
         for decorator in step.decorators:
-            # TODO: rather than relying on decorator names, rely on attributes
-            #       to make them extensible.
-            if decorator.name in ["batch", "kubernetes", "nvidia"]:
-                # TODO: Support arm architectures
-                target_platform = "linux-64"
+            # NOTE: Keep the list of supported decorator names for backward compatibility purposes.
+            # Older versions did not implement the 'support_conda_environment' attribute.
+            if getattr(
+                decorator, "supports_conda_environment", False
+            ) or decorator.name in [
+                "batch",
+                "kubernetes",
+                "nvidia",
+                "snowpark",
+                "slurm",
+                "nvct",
+            ]:
+                target_platform = getattr(decorator, "target_platform", "linux-64")
                 break
 
         environment["conda"]["platforms"] = [target_platform]
@@ -288,19 +410,20 @@ class CondaEnvironment(MetaflowEnvironment):
             # PyPI registries, the usage of environment variable `GOOGLE_APPLICATION_CREDENTIALS`
             # demands that `keyrings.google-artifactregistry-auth` has to be installed
             # and available in the underlying python environment.
-            if os.getenv("GOOGLE_APPLICATION_CREDENTIALS"):
-                environment["conda"]["packages"][
-                    "keyrings.google-artifactregistry-auth"
-                ] = ">=1.1.1"
+            
+            # commenting this out per https://outerboundsco.slack.com/archives/C040K733FND/p1719262399355449
+            # this should be a temporary workaround. Need to find a better fix
+            # if os.getenv("GOOGLE_APPLICATION_CREDENTIALS"):
+            #     environment["conda"]["packages"][
+            #         "keyrings.google-artifactregistry-auth"
+            #     ] = ">=1.1.1"
 
         # Z combinator for a recursive lambda
         deep_sort = (lambda f: f(f))(
             lambda f: lambda obj: (
                 {k: f(f)(v) for k, v in sorted(obj.items())}
                 if isinstance(obj, dict)
-                else sorted([f(f)(e) for e in obj])
-                if isinstance(obj, list)
-                else obj
+                else sorted([f(f)(e) for e in obj]) if isinstance(obj, list) else obj
             )
         )
 
@@ -316,7 +439,7 @@ class CondaEnvironment(MetaflowEnvironment):
                             **environment,
                             **{
                                 "package_root": _datastore_packageroot(
-                                    self.datastore, self.echo
+                                    self.datastore, self.logger
                                 )
                             },
                         }
@@ -356,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):
@@ -366,14 +491,18 @@ class CondaEnvironment(MetaflowEnvironment):
         if id_:
             return [
                 "echo 'Bootstrapping virtual environment...'",
+                "flush_mflogs",
                 # We have to prevent the tracing module from loading,
                 # as the bootstrapping process uses the internal S3 client which would fail to import tracing
                 # due to the required dependencies being bundled into the conda environment,
                 # which is yet to be initialized at this point.
-                'DISABLE_TRACING=True python -m metaflow.plugins.pypi.bootstrap "%s" %s "%s" linux-64'
+                'DISABLE_TRACING=True python -m metaflow.plugins.pypi.bootstrap "%s" %s "%s"'
                 % (self.flow.name, id_, self.datastore_type),
                 "echo 'Environment bootstrapped.'",
-                "export PATH=$PATH:$(pwd)/micromamba",
+                "flush_mflogs",
+                # To avoid having to install micromamba in the PATH in micromamba.py, we add it to the PATH here.
+                "export PATH=$PATH:$(pwd)/micromamba/bin",
+                "export MF_ARCH=$(case $(uname)/$(uname -m) in Darwin/arm64)echo osx-arm64;;Darwin/*)echo osx-64;;Linux/aarch64)echo linux-aarch64;;*)echo linux-64;;esac)",
             ]
         else:
             # for @conda/@pypi(disabled=True).
@@ -434,6 +563,7 @@ class LazyOpen(BufferedIOBase):
         self._file = None
         self._buffer = None
         self._position = 0
+        self.requests = None
 
     def _ensure_file(self):
         if not self._file:
@@ -450,8 +580,13 @@ class LazyOpen(BufferedIOBase):
                 raise ValueError("Both filename and url are missing")
 
     def _download_to_buffer(self):
+        if self.requests is None:
+            # TODO: Remove dependency on requests
+            import requests
+
+            self.requests = requests
         # TODO: Stream it in chunks?
-        response = requests.get(self.url, stream=True)
+        response = self.requests.get(self.url, stream=True)
         response.raise_for_status()
         return response.content