search for micromamba in more locations (#1581)

* search for micromamba in more locations * update todo comment * remove spurious code
Netflix · Oct 10, 2023 · 11da905 · 11da905
1 parent 754ef8a
commit 11da905
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 7 deletions.
diff --git a/metaflow/plugins/pypi/bootstrap.py b/metaflow/plugins/pypi/bootstrap.py
@@ -86,19 +86,24 @@
     # Create Conda environment.
     cmds = [
         # TODO: check if mamba or conda are already available on the image
-        # TODO: search for micromamba everywhere
-        f"""if ! command -v ./micromamba >/dev/null 2>&1; then
-            wget -qO- https://micro.mamba.pm/api/micromamba/{architecture}/latest | python -c "import sys, bz2; sys.stdout.buffer.write(bz2.decompress(sys.stdin.buffer.read()))" | tar -xv bin/micromamba --strip-components=1 ;
-            if ! command -v ./micromamba >/dev/null 2>&1; then
+        # TODO: micromamba installation can be pawned off to micromamba.py
+        f"""set -e;
+        if ! command -v micromamba >/dev/null 2>&1; then
+            mkdir micromamba;
+            wget -qO- https://micro.mamba.pm/api/micromamba/{architecture}/latest | python -c "import sys, bz2; sys.stdout.buffer.write(bz2.decompress(sys.stdin.buffer.read()))" | tar -xv -C $(pwd)/micromamba bin/micromamba --strip-components 1;
+            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'''tmpfile=$(mktemp);
+        f'''set -e;
+        tmpfile=$(mktemp);
         echo "@EXPLICIT" > "$tmpfile";
         ls -d {conda_pkgs_dir}/*/* >> "$tmpfile";
-        ./micromamba create --yes --offline --no-deps --safety-checks=disabled --no-extra-safety-checks --prefix {prefix} --file "$tmpfile";
+        export PATH=$PATH:$(pwd)/micromamba;
+        micromamba create --yes --offline --no-deps --safety-checks=disabled --no-extra-safety-checks --prefix {prefix} --file "$tmpfile";
         rm "$tmpfile"''',
     ]
 
@@ -116,7 +121,9 @@
         # Install PyPI packages.
         cmds.extend(
             [
-                f"""./micromamba run --prefix {prefix} pip --disable-pip-version-check install --root-user-action=ignore --no-compile {pypi_pkgs_dir}/*.whl"""
+                f"""set -e;
+                export PATH=$PATH:$(pwd)/micromamba;
+                micromamba run --prefix {prefix} pip --disable-pip-version-check install --root-user-action=ignore --no-compile {pypi_pkgs_dir}/*.whl"""
             ]
         )
 

diff --git a/metaflow/plugins/pypi/conda_environment.py b/metaflow/plugins/pypi/conda_environment.py
@@ -314,6 +314,7 @@ def bootstrap_commands(self, step_name, datastore_type):
                 '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",
             ]
         else:
             # for @conda/@pypi(disabled=True).

diff --git a/metaflow/plugins/pypi/micromamba.py b/metaflow/plugins/pypi/micromamba.py
@@ -277,10 +277,12 @@ def _install_micromamba(installation_location):
     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}/latest | tar -xvj -C {installation_location} bin/micromamba",
             shell=True,
             stderr=subprocess.PIPE,
+            stdout=subprocess.PIPE,
         )
         _, err = result.communicate()
         if result.returncode != 0: