ob-metaflow 2.12.30.2__py2.py3-none-any.whl → 2.13.6.1__py2.py3-none-any.whl

metaflow/plugins/pypi/bootstrap.py

@@ -1,4 +1,5 @@
 import bz2
+import concurrent.futures
 import io
 import json
 import os
@@ -6,21 +7,33 @@ import shutil
 import subprocess
 import sys
 import tarfile
-
+import time
+from urllib.error import URLError
+from urllib.request import urlopen
 from metaflow.metaflow_config import DATASTORE_LOCAL_DIR
 from metaflow.plugins import DATASTORES
+from metaflow.plugins.pypi.utils import MICROMAMBA_MIRROR_URL, MICROMAMBA_URL
 from metaflow.util import which
+from urllib.request import Request
+import warnings
 
 from . import MAGIC_FILE, _datastore_packageroot
 
 # Bootstraps a valid conda virtual environment composed of conda and pypi packages
 
-if __name__ == "__main__":
-    if len(sys.argv) != 5:
-        print("Usage: bootstrap.py <flow_name> <id> <datastore_type> <architecture>")
-        sys.exit(1)
-    _, flow_name, id_, datastore_type, architecture = sys.argv
 
+def timer(func):
+    def wrapper(*args, **kwargs):
+        start_time = time.time()
+        result = func(*args, **kwargs)
+        duration = time.time() - start_time
+        # print(f"Time taken for {func.__name__}: {duration:.2f} seconds")
+        return result
+
+    return wrapper
+
+
+if __name__ == "__main__":
     # TODO: Detect architecture on the fly when dealing with arm architectures.
     # ARCH=$(uname -m)
     # OS=$(uname)
@@ -45,96 +58,251 @@ if __name__ == "__main__":
     #     fi
     # fi
 
-    prefix = os.path.join(os.getcwd(), architecture, id_)
-    pkgs_dir = os.path.join(os.getcwd(), ".pkgs")
-    manifest_dir = os.path.join(os.getcwd(), DATASTORE_LOCAL_DIR, flow_name)
+    def run_cmd(cmd):
+        result = subprocess.run(
+            cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True
+        )
+        if result.returncode != 0:
+            print(f"Bootstrap failed while executing: {cmd}")
+            print("Stdout:", result.stdout)
+            print("Stderr:", result.stderr)
+            sys.exit(1)
 
-    datastores = [d for d in DATASTORES if d.TYPE == datastore_type]
-    if not datastores:
-        print(f"No datastore found for type: {datastore_type}")
-        sys.exit(1)
+    @timer
+    def install_micromamba(architecture):
+        micromamba_dir = os.path.join(os.getcwd(), "micromamba")
+        micromamba_path = os.path.join(micromamba_dir, "bin", "micromamba")
+
+        if which("micromamba"):
+            return which("micromamba")
+        if os.path.exists(micromamba_path):
+            os.environ["PATH"] += os.pathsep + os.path.dirname(micromamba_path)
+            return micromamba_path
+
+        # Download and extract in one go
+        url = MICROMAMBA_URL.format(platform=architecture, version="2.0.4")
+        mirror_url = MICROMAMBA_MIRROR_URL.format(
+            platform=architecture, version="2.0.4"
+        )
+
+        # Prepare directory once
+        os.makedirs(os.path.dirname(micromamba_path), exist_ok=True)
+
+        # Download and decompress in one go
+        def _download_and_extract(url):
+            headers = {
+                "Accept-Encoding": "gzip, deflate, br",
+                "Connection": "keep-alive",
+                "User-Agent": "python-urllib",
+            }
+
+            max_retries = 3
+            for attempt in range(max_retries):
+                try:
+                    req = Request(url, headers=headers)
 
-    storage = datastores[0](
-        _datastore_packageroot(datastores[0], lambda *args, **kwargs: None)
-    )
+                    with urlopen(req) as response:
+                        decompressor = bz2.BZ2Decompressor()
+                        with warnings.catch_warnings():
+                            warnings.filterwarnings(
+                                "ignore", category=DeprecationWarning
+                            )
+                            with tarfile.open(
+                                fileobj=io.BytesIO(
+                                    decompressor.decompress(response.read())
+                                ),
+                                mode="r:",
+                            ) as tar:
+                                member = tar.getmember("bin/micromamba")
+                                tar.extract(member, micromamba_dir)
+                    break
+                except (URLError, IOError) as e:
+                    if attempt == max_retries - 1:
+                        raise Exception(
+                            f"Failed to download micromamba after {max_retries} attempts: {e}"
+                        )
+                    time.sleep(2**attempt)
 
-    # Move MAGIC_FILE inside local datastore.
-    os.makedirs(manifest_dir, exist_ok=True)
-    shutil.move(
-        os.path.join(os.getcwd(), MAGIC_FILE),
-        os.path.join(manifest_dir, MAGIC_FILE),
-    )
+        try:
+            # first try from mirror
+            _download_and_extract(mirror_url)
+        except Exception:
+            # download from mirror failed, try official source before failing.
+            _download_and_extract(url)
 
-    with open(os.path.join(manifest_dir, MAGIC_FILE)) as f:
-        env = json.load(f)[id_][architecture]
+        # Set executable permission
+        os.chmod(micromamba_path, 0o755)
 
-    # Download Conda packages.
-    conda_pkgs_dir = os.path.join(pkgs_dir, "conda")
-    with storage.load_bytes([package["path"] for package in env["conda"]]) as results:
-        for key, tmpfile, _ in results:
+        # Update PATH only once at the end
+        os.environ["PATH"] += os.pathsep + os.path.dirname(micromamba_path)
+        return micromamba_path
+
+    @timer
+    def download_conda_packages(storage, packages, dest_dir):
+        def process_conda_package(args):
             # Ensure that conda packages go into architecture specific folders.
             # The path looks like REPO/CHANNEL/CONDA_SUBDIR/PACKAGE. We trick
             # Micromamba into believing that all packages are coming from a local
             # channel - the only hurdle is ensuring that packages are organised
             # properly.
-
-            # TODO: consider RAM disk
-            dest = os.path.join(conda_pkgs_dir, "/".join(key.split("/")[-2:]))
+            key, tmpfile, dest_dir = args
+            dest = os.path.join(dest_dir, "/".join(key.split("/")[-2:]))
             os.makedirs(os.path.dirname(dest), exist_ok=True)
             shutil.move(tmpfile, dest)
 
-    # Create Conda environment.
-    cmds = [
-        # TODO: check if mamba or conda are already available on the image
-        # TODO: micromamba installation can be pawned off to micromamba.py
-        f"""set -e;
-        if ! command -v micromamba >/dev/null 2>&1; then
-            mkdir -p micromamba;
-            python -c "import requests, bz2, sys; data = requests.get('https://micro.mamba.pm/api/micromamba/{architecture}/1.5.7').content; sys.stdout.buffer.write(bz2.decompress(data))" | tar -xv -C $(pwd)/micromamba bin/micromamba --strip-components 1;
+        os.makedirs(dest_dir, exist_ok=True)
+        with storage.load_bytes([package["path"] for package in packages]) as results:
+            with concurrent.futures.ThreadPoolExecutor() as executor:
+                executor.map(
+                    process_conda_package,
+                    [(key, tmpfile, dest_dir) for key, tmpfile, _ in results],
+                )
+            # for key, tmpfile, _ in results:
+
+            #     # TODO: consider RAM disk
+            #     dest = os.path.join(dest_dir, "/".join(key.split("/")[-2:]))
+            #     os.makedirs(os.path.dirname(dest), exist_ok=True)
+            #     shutil.move(tmpfile, dest)
+        return dest_dir
+
+    @timer
+    def download_pypi_packages(storage, packages, dest_dir):
+        def process_pypi_package(args):
+            key, tmpfile, dest_dir = args
+            dest = os.path.join(dest_dir, os.path.basename(key))
+            shutil.move(tmpfile, dest)
+
+        os.makedirs(dest_dir, exist_ok=True)
+        with storage.load_bytes([package["path"] for package in packages]) as results:
+            with concurrent.futures.ThreadPoolExecutor() as executor:
+                executor.map(
+                    process_pypi_package,
+                    [(key, tmpfile, dest_dir) for key, tmpfile, _ in results],
+                )
+            # for key, tmpfile, _ in results:
+            #     dest = os.path.join(dest_dir, os.path.basename(key))
+            #     shutil.move(tmpfile, dest)
+        return dest_dir
+
+    @timer
+    def create_conda_environment(prefix, conda_pkgs_dir):
+        cmd = f'''set -e;
+            tmpfile=$(mktemp);
+            echo "@EXPLICIT" > "$tmpfile";
+            ls -d {conda_pkgs_dir}/*/* >> "$tmpfile";
+            export PATH=$PATH:$(pwd)/micromamba;
+            export CONDA_PKGS_DIRS=$(pwd)/micromamba/pkgs;
+            export MAMBA_NO_LOW_SPEED_LIMIT=1;
+            export MAMBA_USE_INDEX_CACHE=1;
+            export MAMBA_NO_PROGRESS_BARS=1;
+            export CONDA_FETCH_THREADS=1;
+            micromamba create --yes --offline --no-deps \
+                --safety-checks=disabled --no-extra-safety-checks \
+                --prefix {prefix} --file "$tmpfile" \
+                --no-pyc --no-rc --always-copy;
+            rm "$tmpfile"'''
+        run_cmd(cmd)
+
+    @timer
+    def install_pypi_packages(prefix, pypi_pkgs_dir):
+        cmd = f"""set -e;
             export PATH=$PATH:$(pwd)/micromamba;
-            if ! command -v micromamba >/dev/null 2>&1; then
-                echo "Failed to install Micromamba!";
-                exit 1;
-            fi;
-        fi""",
-        # Create a conda environment through Micromamba.
-        f'''set -e;
-        tmpfile=$(mktemp);
-        echo "@EXPLICIT" > "$tmpfile";
-        ls -d {conda_pkgs_dir}/*/* >> "$tmpfile";
-        export PATH=$PATH:$(pwd)/micromamba;
-        export CONDA_PKGS_DIRS=$(pwd)/micromamba/pkgs;
-        micromamba create --yes --offline --no-deps --safety-checks=disabled --no-extra-safety-checks --prefix {prefix} --file "$tmpfile";
-        rm "$tmpfile"''',
-    ]
-
-    # Download PyPI packages.
-    if "pypi" in env:
+            export CONDA_PKGS_DIRS=$(pwd)/micromamba/pkgs;
+            micromamba run --prefix {prefix} python -m pip --disable-pip-version-check \
+                install --root-user-action=ignore --no-compile --no-index \
+                --no-cache-dir --no-deps --prefer-binary \
+                --find-links={pypi_pkgs_dir}  --no-user \
+                --no-warn-script-location --no-input \
+                {pypi_pkgs_dir}/*.whl
+            """
+        run_cmd(cmd)
+
+    @timer
+    def setup_environment(
+        architecture, storage, env, prefix, conda_pkgs_dir, pypi_pkgs_dir
+    ):
+        with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
+            # install micromamba, download conda and pypi packages in parallel
+            futures = {
+                "micromamba": executor.submit(install_micromamba, architecture),
+                "conda_pkgs": executor.submit(
+                    download_conda_packages, storage, env["conda"], conda_pkgs_dir
+                ),
+            }
+            if "pypi" in env:
+                futures["pypi_pkgs"] = executor.submit(
+                    download_pypi_packages, storage, env["pypi"], pypi_pkgs_dir
+                )
+
+            # create conda environment after micromamba is installed and conda packages are downloaded
+            done, _ = concurrent.futures.wait(
+                [futures["micromamba"], futures["conda_pkgs"]],
+                return_when=concurrent.futures.ALL_COMPLETED,
+            )
+
+            for future in done:
+                future.result()
+
+            # start conda environment creation
+            futures["conda_env"] = executor.submit(
+                create_conda_environment, prefix, conda_pkgs_dir
+            )
+
+            if "pypi" in env:
+                # install pypi packages after conda environment is created and pypi packages are downloaded
+                done, _ = concurrent.futures.wait(
+                    [futures["conda_env"], futures["pypi_pkgs"]],
+                    return_when=concurrent.futures.ALL_COMPLETED,
+                )
+
+                for future in done:
+                    future.result()
+
+                # install pypi packages
+                futures["pypi_install"] = executor.submit(
+                    install_pypi_packages, prefix, pypi_pkgs_dir
+                )
+                # wait for pypi packages to be installed
+                futures["pypi_install"].result()
+            else:
+                # wait for conda environment to be created
+                futures["conda_env"].result()
+
+    if len(sys.argv) != 5:
+        print("Usage: bootstrap.py <flow_name> <id> <datastore_type> <architecture>")
+        sys.exit(1)
+
+    try:
+        _, flow_name, id_, datastore_type, architecture = sys.argv
+
+        prefix = os.path.join(os.getcwd(), architecture, id_)
+        pkgs_dir = os.path.join(os.getcwd(), ".pkgs")
+        conda_pkgs_dir = os.path.join(pkgs_dir, "conda")
         pypi_pkgs_dir = os.path.join(pkgs_dir, "pypi")
-        with storage.load_bytes(
-            [package["path"] for package in env["pypi"]]
-        ) as results:
-            for key, tmpfile, _ in results:
-                dest = os.path.join(pypi_pkgs_dir, os.path.basename(key))
-                os.makedirs(os.path.dirname(dest), exist_ok=True)
-                shutil.move(tmpfile, dest)
-
-        # Install PyPI packages.
-        cmds.extend(
-            [
-                f"""set -e;
-                export PATH=$PATH:$(pwd)/micromamba;
-                export CONDA_PKGS_DIRS=$(pwd)/micromamba/pkgs;
-                micromamba run --prefix {prefix} python -m pip --disable-pip-version-check install --root-user-action=ignore --no-compile {pypi_pkgs_dir}/*.whl --no-user"""
-            ]
+        manifest_dir = os.path.join(os.getcwd(), DATASTORE_LOCAL_DIR, flow_name)
+
+        datastores = [d for d in DATASTORES if d.TYPE == datastore_type]
+        if not datastores:
+            print(f"No datastore found for type: {datastore_type}")
+            sys.exit(1)
+
+        storage = datastores[0](
+            _datastore_packageroot(datastores[0], lambda *args, **kwargs: None)
         )
 
-    for cmd in cmds:
-        result = subprocess.run(
-            cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE
+        # Move MAGIC_FILE inside local datastore.
+        os.makedirs(manifest_dir, exist_ok=True)
+        shutil.move(
+            os.path.join(os.getcwd(), MAGIC_FILE),
+            os.path.join(manifest_dir, MAGIC_FILE),
         )
-        if result.returncode != 0:
-            print(f"Bootstrap failed while executing: {cmd}")
-            print("Stdout:", result.stdout.decode())
-            print("Stderr:", result.stderr.decode())
-            sys.exit(1)
+        with open(os.path.join(manifest_dir, MAGIC_FILE)) as f:
+            env = json.load(f)[id_][architecture]
+
+        setup_environment(
+            architecture, storage, env, prefix, conda_pkgs_dir, pypi_pkgs_dir
+        )
+
+    except Exception as e:
+        print(f"Error: {str(e)}", file=sys.stderr)
+        sys.exit(1)

metaflow/plugins/pypi/conda_decorator.py

@@ -50,20 +50,26 @@ class CondaStepDecorator(StepDecorator):
     # conda channels, users can specify channel::package as the package name.
 
     def __init__(self, attributes=None, statically_defined=False):
-        self._user_defined_attributes = (
-            attributes.copy() if attributes is not None else {}
+        self._attributes_with_user_values = (
+            set(attributes.keys()) if attributes is not None else set()
         )
+
         super(CondaStepDecorator, self).__init__(attributes, statically_defined)
 
+    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"]
+        if self.attributes["packages"]:
+            self._attributes_with_user_values.add("packages")
 
     def is_attribute_user_defined(self, name):
-        return name in self._user_defined_attributes
+        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
@@ -83,10 +89,10 @@ class CondaStepDecorator(StepDecorator):
                 **super_attributes["packages"],
                 **self.attributes["packages"],
             }
-            self._user_defined_attributes = {
-                **self._user_defined_attributes,
-                **conda_base._user_defined_attributes,
-            }
+            self._attributes_with_user_values.update(
+                conda_base._attributes_with_user_values
+            )
+
             self.attributes["python"] = (
                 self.attributes["python"] or super_attributes["python"]
             )
@@ -333,11 +339,15 @@ class CondaFlowDecorator(FlowDecorator):
     }
 
     def __init__(self, attributes=None, statically_defined=False):
-        self._user_defined_attributes = (
-            attributes.copy() if attributes is not None else {}
+        self._attributes_with_user_values = (
+            set(attributes.keys()) if attributes is not None else set()
         )
+
         super(CondaFlowDecorator, self).__init__(attributes, statically_defined)
 
+    def init(self):
+        super(CondaFlowDecorator, self).init()
+
         # Support legacy 'libraries=' attribute for the decorator.
         self.attributes["packages"] = {
             **self.attributes["libraries"],
@@ -348,7 +358,7 @@ class CondaFlowDecorator(FlowDecorator):
             self.attributes["python"] = str(self.attributes["python"])
 
     def is_attribute_user_defined(self, name):
-        return name in self._user_defined_attributes
+        return name in self._attributes_with_user_values
 
     def flow_init(
         self, flow, graph, environment, flow_datastore, metadata, logger, echo, options

metaflow/plugins/pypi/conda_environment.py

@@ -5,21 +5,17 @@ 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.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 . import MAGIC_FILE, _datastore_packageroot
 from .utils import conda_platform
@@ -50,7 +46,6 @@ class CondaEnvironment(MetaflowEnvironment):
 
     def validate_environment(self, logger, datastore_type):
         self.datastore_type = datastore_type
-        self.logger = logger
 
         # Avoiding circular imports.
         from metaflow.plugins import DATASTORES
@@ -62,8 +57,21 @@ 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)
+
+        # TODO: Wire up logging
+        micromamba = Micromamba(self.logger)
+        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
@@ -150,6 +158,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",
@@ -166,22 +177,60 @@ class CondaEnvironment(MetaflowEnvironment):
                 if id_ in dirty:
                     self.write_to_environment_manifest([id_, platform, type_], packages)
 
-        # First resolve environments through Conda, before PyPI.
+        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) ...")
-        for solver in ["conda", "pypi"]:
-            with ThreadPoolExecutor() as executor:
-                results = list(
-                    executor.map(lambda x: solve(*x, solver), environments(solver))
-                )
-            _ = 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)
-                )
-            if self.datastore_type not in ["local"]:
-                # Cache packages only when a remote datastore is in play.
-                storage = self.datastore(_datastore_packageroot(self.datastore, echo))
-                cache(storage, results, solver)
+        # 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
+            conda_futures = [
+                executor.submit(lambda x: solve(*x, "conda"), env)
+                for env in environments("conda")
+            ]
+
+            pypi_envs = {env[0]: env for env in environments("pypi")}
+            pypi_futures = []
+
+            # Process conda results sequentially for downloads
+            for future in as_completed(conda_futures):
+                result = future.result()
+                # Sequential conda download
+                self.solvers["conda"].download(*result)
+                # Parallel conda create and cache
+                create_future = executor.submit(self.solvers["conda"].create, *result)
+                if storage:
+                    executor.submit(cache, storage, [result], "conda")
+
+                # Queue PyPI solve to start after conda create
+                if result[0] in pypi_envs:
+
+                    def pypi_solve(env):
+                        create_future.result()  # Wait for conda create
+                        return solve(*env, "pypi")
+
+                    pypi_futures.append(
+                        executor.submit(pypi_solve, pypi_envs[result[0]])
+                    )
+
+            # Process PyPI results sequentially for downloads
+            for solve_future in pypi_futures:
+                result = solve_future.result()
+                # Sequential PyPI download
+                self.solvers["pypi"].download(*result)
+                # Parallel PyPI create and cache
+                executor.submit(self.solvers["pypi"].create, *result)
+                if storage:
+                    executor.submit(cache, storage, [result], "pypi")
         self.logger("Virtual environment(s) bootstrapped!")
 
     def executable(self, step_name, default=None):
@@ -385,7 +434,8 @@ class CondaEnvironment(MetaflowEnvironment):
                 'DISABLE_TRACING=True python -m metaflow.plugins.pypi.bootstrap "%s" %s "%s" linux-64'
                 % (self.flow.name, id_, self.datastore_type),
                 "echo 'Environment bootstrapped.'",
-                "export PATH=$PATH:$(pwd)/micromamba",
+                # 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",
             ]
         else:
             # for @conda/@pypi(disabled=True).
@@ -446,6 +496,7 @@ class LazyOpen(BufferedIOBase):
         self._file = None
         self._buffer = None
         self._position = 0
+        self.requests = None
 
     def _ensure_file(self):
         if not self._file:
@@ -462,8 +513,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