feature: micromamba download from mirror (#2178)

* add mirror url for downloading micromamba. add fallback logic for download

* support separate versions for client-side and server-side micromamba

* add retry to client-side micromamba download as well

* point micromamba mirror url to Outerbounds mirror

metaflow/plugins/pypi/bootstrap.py

@@ -12,6 +12,7 @@
 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
@@ -33,7 +34,6 @@ def wrapper(*args, **kwargs):
 
 
 if __name__ == "__main__":
-
     # TODO: Detect architecture on the fly when dealing with arm architectures.
     # ARCH=$(uname -m)
     # OS=$(uname)
@@ -80,43 +80,55 @@ def install_micromamba(architecture):
             return micromamba_path
 
         # Download and extract in one go
-        # TODO: Serve from cloudflare
-        url = f"https://micro.mamba.pm/api/micromamba/{architecture}/2.0.4"
+        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
-        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)
-
-                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)
+        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)
+
+                    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)
+
+        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)
 
         # Set executable permission
         os.chmod(micromamba_path, 0o755)
@@ -127,7 +139,6 @@ def install_micromamba(architecture):
 
     @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
@@ -156,7 +167,6 @@ def process_conda_package(args):
 
     @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))
@@ -195,7 +205,6 @@ def create_conda_environment(prefix, conda_pkgs_dir):
 
     @timer
     def install_pypi_packages(prefix, pypi_pkgs_dir):
-
         cmd = f"""set -e;
             export PATH=$PATH:$(pwd)/micromamba;
             export CONDA_PKGS_DIRS=$(pwd)/micromamba/pkgs;

metaflow/plugins/pypi/micromamba.py

@@ -3,11 +3,12 @@
 import os
 import subprocess
 import tempfile
+import time
 
 from metaflow.exception import MetaflowException
 from metaflow.util import which
 
-from .utils import conda_platform
+from .utils import MICROMAMBA_MIRROR_URL, MICROMAMBA_URL, conda_platform
 
 
 class MicromambaException(MetaflowException):
@@ -338,23 +339,37 @@ 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
-        # requires bzip2
-        result = subprocess.Popen(
-            f"curl -Ls https://micro.mamba.pm/api/micromamba/{platform}/1.5.7 | tar -xvj -C {installation_location} bin/micromamba",
-            shell=True,
-            stderr=subprocess.PIPE,
-            stdout=subprocess.PIPE,
-        )
-        _, err = result.communicate()
-        if result.returncode != 0:
-            raise MicromambaException(
-                f"Micromamba installation '{result.args}' failed:\n{err.decode()}"
-            )
+    url = MICROMAMBA_URL.format(platform=platform, version="1.5.7")
+    mirror_url = MICROMAMBA_MIRROR_URL.format(platform=platform, version="1.5.7")
+    os.makedirs(installation_location, exist_ok=True)
 
-    except subprocess.CalledProcessError as e:
-        raise MicromambaException(
-            "Micromamba installation failed:\n{}".format(e.stderr.decode())
-        )
+    def _download_and_extract(url):
+        max_retries = 3
+        for attempt in range(max_retries):
+            try:
+                # https://mamba.readthedocs.io/en/latest/micromamba-installation.html#manual-installation
+                # requires bzip2
+                result = subprocess.Popen(
+                    f"curl -Ls {url} | tar -xvj -C {installation_location} bin/micromamba",
+                    shell=True,
+                    stderr=subprocess.PIPE,
+                    stdout=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:
+                if attempt == max_retries - 1:
+                    raise MicromambaException(
+                        "Micromamba installation failed:\n{}".format(e.stderr.decode())
+                    )
+                time.sleep(2**attempt)
+
+    try:
+        # prioritize downloading from mirror
+        _download_and_extract(mirror_url)
+    except Exception:
+        # download from official source as a fallback
+        _download_and_extract(url)

metaflow/plugins/pypi/utils.py

@@ -20,6 +20,9 @@ class Tags:
 
 from metaflow.exception import MetaflowException
 
+MICROMAMBA_URL = "https://micro.mamba.pm/api/micromamba/{platform}/{version}"
+MICROMAMBA_MIRROR_URL = "https://micromamba.outerbounds.sh/{platform}/{version}.tar.bz2"
+
 
 def conda_platform():
     # Returns the conda platform for the Python interpreter