pathways: verify if conda or mamba is installed

Author: Rohan Gujarathi

src/sagemaker/remote_function/core/runtime_environment.py

@@ -135,7 +135,9 @@ def _create_conda_env(self):
     def _install_req_txt_in_conda_env(self, env_name, local_path):
         """Install requirements.txt in the given conda environment"""
 
-        cmd = '/bin/bash -c "mamba run -n {} pip install -r {}"'.format(env_name, local_path)
+        cmd = '/bin/bash -c "{} run -n {} pip install -r {}"'.format(
+            self._get_conda_exe(), env_name, local_path
+        )
         logger.info("Activating conda env and installing requirements: {}".format(cmd))
         _run_shell_cmd(cmd)
         logger.info("Requirements installed successfully in conda env {}".format(env_name))
@@ -151,6 +153,16 @@ def _write_conda_env_to_file(self, env_name):
         with open(file_path, "w") as output_file:
             output_file.write(env_name)
 
+    def _get_conda_exe(self):
+        """Checks whether conda or mamba is available to use"""
+
+        if not subprocess.Popen(["which", "mamba"]).wait():
+            return "mamba"
+        elif not subprocess.Popen(["which", "conda"]).wait():
+            return "conda"
+        else:
+            raise ValueError("Neither conda nor mamba is installed on the image")
+
 
 def _run_shell_cmd(cmd: str):
     """This method runs a given shell command using subprocess

src/sagemaker/remote_function/job_driver.sh

@@ -10,8 +10,15 @@ python -m sagemaker.remote_function.bootstrap "$@"
 if [ -f "remote_function_conda_env.txt" ]
 then
     conda_env=$(cat remote_function_conda_env.txt)
+
+    if which mamba >/dev/null; then
+        conda_exe="mamba"
+    else
+        conda_exe="conda"
+    fi
+
     printf "INFO: Invoking remote function inside conda environment: $conda_env.\n"
-    mamba run -n $conda_env python -m sagemaker.remote_function.invoke_function "$@"
+    $conda_exe run -n $conda_env python -m sagemaker.remote_function.invoke_function "$@"
 else
     printf "INFO: No conda env provided. Invoking remote function\n"
     python -m sagemaker.remote_function.invoke_function "$@"

tests/unit/sagemaker/remote_function/core/test_runtime_environment.py

@@ -173,9 +173,13 @@ def test_bootstrap_req_txt_error(mock_s3_download, mock_cwd, mock_session):
     "sagemaker.remote_function.core.runtime_environment.RuntimeEnvironmentManager._write_conda_env_to_file",
     Mock(),
 )
+@patch(
+    "sagemaker.remote_function.core.runtime_environment.RuntimeEnvironmentManager._get_conda_exe",
+    return_value="conda",
+)
 @patch("os.getcwd", return_value="/usr/local/path")
 @patch("sagemaker.s3.S3Downloader.download")
-def test_bootstrap_req_txt_with_conda_env(mock_s3_download, mock_cwd, mock_session):
+def test_bootstrap_req_txt_with_conda_env(mock_s3_download, mock_cwd, mock_conda_exe, mock_session):
     with patch("subprocess.Popen") as popen:
         popen.return_value.wait.return_value = 0
         runtime = RuntimeEnvironmentManager(TEST_S3_BASE_URI, TEST_S3_KMS_KEY, mock_session)
@@ -193,8 +197,8 @@ def test_bootstrap_req_txt_with_conda_env(mock_s3_download, mock_cwd, mock_sessi
         call_args = popen.call_args[0][0]
         assert call_args is not None
 
-        expected_cmd = '/bin/bash -c "mamba run -n {} pip install -r {}"'.format(
-            job_conda_env, local_file_path
+        expected_cmd = '/bin/bash -c "{} run -n {} pip install -r {}"'.format(
+            mock_conda_exe.return_value, job_conda_env, local_file_path
         )
         assert call_args == expected_cmd