Feature: support requirements.txt installation on conda env

Author: Rohan Gujarathi

src/sagemaker/remote_function/core/runtime_environment.py

@@ -16,7 +16,6 @@
 
 import sys
 import os
-import shlex
 import subprocess
 from sagemaker.s3 import s3_path_join, S3Uploader, S3Downloader
 from sagemaker.session import Session
@@ -88,20 +87,19 @@ def bootstrap(self, dependencies_s3_uri: str, job_conda_env: str = None):
         """
 
         if dependencies_s3_uri.endswith(".txt"):
-            if job_conda_env:
-                # TODO:
-                #   1. verify if conda exists in the image
-                #   2. activate the given conda env
-                #   3. update the conda env with req.txt file
-                return
-
             local_path = os.getcwd()
             S3Downloader.download(
                 dependencies_s3_uri, local_path, self.s3_kms_key, self.sagemaker_session
             )
 
             local_dependencies_file = os.path.join(local_path, dependencies_s3_uri.split("/")[-1])
-            self._install_requirements_txt(local_dependencies_file, _python_executable())
+            conda_env = job_conda_env or os.getenv("DEFAULT_CONDA_ENV")
+            if conda_env:
+                self._install_req_txt_in_conda_env(conda_env, local_dependencies_file)
+                self._write_conda_env_to_file(conda_env)
+
+            else:
+                self._install_requirements_txt(local_dependencies_file, _python_executable())
 
         elif dependencies_s3_uri.endswith(".yml"):
             # TODO: implement
@@ -134,31 +132,40 @@ def _create_conda_env(self):
         # TODO: implement
         pass  # pylint: disable=W0107
 
-    def _activate_conda_env(self):
-        """Activate conda environment"""
-        # TODO: implement
-        pass  # pylint: disable=W0107
+    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)
+        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))
 
     def _update_conda_env(self):
         """Update conda env using conda yml file"""
         # TODO: implement
         pass  # pylint: disable=W0107
 
+    def _write_conda_env_to_file(self, env_name):
+        file_name = "remote_function_conda_env.txt"
+        file_path = os.path.join(os.getcwd(), file_name)
+        with open(file_path, "w") as output_file:
+            output_file.write(env_name)
+
 
 def _run_shell_cmd(cmd: str):
     """This method runs a given shell command using subprocess
 
     Raises RuntimeEnvironmentError if the command fails
     """
 
-    process = subprocess.Popen(shlex.split(cmd), stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+    process = subprocess.Popen((cmd), stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
 
     _log_output(process)
     error_logs = _log_error(process)
     return_code = process.wait()
     if return_code:
-        error_message = "Encountered error while installing dependencies. Reason: {}".format(
-            error_logs
+        error_message = "Encountered error while running command '{}'. Reason: {}".format(
+            cmd, error_logs
         )
         raise RuntimeEnvironmentError(error_message)
 

src/sagemaker/remote_function/job_driver.sh

@@ -4,8 +4,16 @@
 
 set -eu
 
-printf "INFO: Bootstraping runtime environment\n"
+printf "INFO: Bootstraping runtime environment.\n"
 python -m sagemaker.remote_function.bootstrap "$@"
 
-printf "INFO: Invoking remote function\n"
-python -m sagemaker.remote_function.invoke_function "$@"
+if [ -f "remote_function_conda_env.txt" ]
+then
+    conda_env=$(cat remote_function_conda_env.txt)
+    printf "INFO: Invoking remote function inside conda environment: $conda_env.\n"
+    mamba 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 "$@"
+fi
+

tests/integ/sagemaker/remote_function/conftest.py

@@ -34,21 +34,47 @@
     "RUN rm {source_archive}\n"
 )
 
+DOCKERFILE_TEMPLATE_WITH_CONDA = (
+    "FROM public.ecr.aws/docker/library/python:{py_version}-slim\n\n"
+    "WORKDIR /opt/ml/remote_function/\n"
+    'SHELL ["/bin/bash", "-c"]\n'
+    "RUN apt-get update -y \
+        && apt-get install -y unzip curl\n\n"
+    "RUN curl -L -O 'https://github.com/conda-forge/miniforge/releases/latest/download/Mambaforge-Linux-x86_64.sh' \
+        && bash Mambaforge-Linux-x86_64.sh -b -p '/opt/conda' \
+        && /opt/conda/bin/conda init bash\n\n"
+    "ENV PATH $PATH:/opt/conda/bin\n"
+    "RUN mamba create -n integ_test_env python=3.10 -y\n"
+    "COPY {source_archive} ./\n"
+    "RUN pip install '{source_archive}[remote_function]' \
+        && mamba run -n base pip install '{source_archive}[remote_function]' \
+        && mamba run -n integ_test_env pip install '{source_archive}[remote_function]' \
+        && rm {source_archive}\n"
+    "ENV SHELL=/bin/bash\n"
+    "ENV DEFAULT_CONDA_ENV=base\n"
+)
+
 
 @pytest.fixture(scope="package")
 def dummy_container_without_error(sagemaker_session):
     # TODO: the python version should be dynamically specified instead of hardcoding
-    ecr_uri = _build_container(sagemaker_session, "3.10")
+    ecr_uri = _build_container(sagemaker_session, "3.10", DOCKERFILE_TEMPLATE)
     return ecr_uri
 
 
 @pytest.fixture(scope="package")
 def dummy_container_incompatible_python_runtime(sagemaker_session):
-    ecr_uri = _build_container(sagemaker_session, "3.7")
+    ecr_uri = _build_container(sagemaker_session, "3.7", DOCKERFILE_TEMPLATE)
+    return ecr_uri
+
+
+@pytest.fixture(scope="package")
+def dummy_container_with_conda(sagemaker_session):
+    ecr_uri = _build_container(sagemaker_session, "3.10", DOCKERFILE_TEMPLATE_WITH_CONDA)
     return ecr_uri
 
 
-def _build_container(sagemaker_session, py_version):
+def _build_container(sagemaker_session, py_version, docker_templete):
     """Build a dummy test container locally and push a container to an ecr repo"""
 
     region = sagemaker_session.boto_region_name
@@ -62,7 +88,7 @@ def _build_container(sagemaker_session, py_version):
         source_archive = _generate_and_move_sagemaker_sdk_tar(tmpdir)
         with open(os.path.join(tmpdir, "Dockerfile"), "w") as file:
             file.writelines(
-                DOCKERFILE_TEMPLATE.format(py_version=py_version, source_archive=source_archive)
+                docker_templete.format(py_version=py_version, source_archive=source_archive)
             )
 
         docker_client = docker.from_env()

tests/integ/sagemaker/remote_function/test_decorator.py

@@ -118,6 +118,69 @@ def cuberoot(x):
     assert cuberoot(27) == 3
 
 
+def test_additional_dependencies_with_job_conda_env(
+    sagemaker_session, dummy_container_with_conda, cpu_instance_type
+):
+    dependencies_path = os.path.join(DATA_DIR, "remote_function/requirements.txt")
+
+    @remote(
+        role=ROLE,
+        image_uri=dummy_container_with_conda,
+        dependencies=dependencies_path,
+        instance_type=cpu_instance_type,
+        sagemaker_session=sagemaker_session,
+        job_conda_env="integ_test_env",
+    )
+    def cuberoot(x):
+        from scipy.special import cbrt
+
+        return cbrt(x)
+
+    assert cuberoot(27) == 3
+
+
+def test_additional_dependencies_with_default_conda_env(
+    sagemaker_session, dummy_container_with_conda, cpu_instance_type
+):
+    dependencies_path = os.path.join(DATA_DIR, "remote_function/requirements.txt")
+
+    @remote(
+        role=ROLE,
+        image_uri=dummy_container_with_conda,
+        dependencies=dependencies_path,
+        instance_type=cpu_instance_type,
+        sagemaker_session=sagemaker_session,
+    )
+    def cuberoot(x):
+        from scipy.special import cbrt
+
+        return cbrt(x)
+
+    assert cuberoot(27) == 3
+
+
+def test_additional_dependencies_with_non_existent_conda_env(
+    sagemaker_session, dummy_container_with_conda, cpu_instance_type
+):
+    dependencies_path = os.path.join(DATA_DIR, "remote_function/requirements.txt")
+
+    @remote(
+        role=ROLE,
+        image_uri=dummy_container_with_conda,
+        dependencies=dependencies_path,
+        instance_type=cpu_instance_type,
+        sagemaker_session=sagemaker_session,
+        job_conda_env="non_existent_env",
+    )
+    def cuberoot(x):
+        from scipy.special import cbrt
+
+        return cbrt(x)
+
+    with pytest.raises(RuntimeEnvironmentError):
+        cuberoot(27) == 3
+
+
 def test_with_non_existent_dependencies(
     sagemaker_session, dummy_container_without_error, cpu_instance_type
 ):

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

@@ -135,9 +135,8 @@ def test_bootstrap_req_txt(mock_s3_download, mock_cwd, mock_session):
         call_args = popen.call_args[0][0]
         assert call_args is not None
 
-        expected = [python_exe, "-m", "pip", "install", "-r", local_file_path]
-        for i, v in enumerate(expected):
-            assert call_args[i] == v
+        expected_cmd = "{} -m pip install -r {}".format(python_exe, local_file_path)
+        assert call_args == expected_cmd
 
 
 @patch("sagemaker.remote_function.core.runtime_environment._log_error", Mock())
@@ -164,13 +163,40 @@ def test_bootstrap_req_txt_error(mock_s3_download, mock_cwd, mock_session):
         call_args = popen.call_args[0][0]
         assert call_args is not None
 
-        expected = [python_exe, "-m", "pip", "install", "-r", local_file_path]
-        for i, v in enumerate(expected):
-            assert call_args[i] == v
+        expected_cmd = "{} -m pip install -r {}".format(python_exe, local_file_path)
+        assert call_args == expected_cmd
 
 
-def test_bootstrap_req_txt_with_conda_env():
-    pass
+@patch("sagemaker.remote_function.core.runtime_environment._log_error", Mock())
+@patch("sagemaker.remote_function.core.runtime_environment._log_output", Mock())
+@patch(
+    "sagemaker.remote_function.core.runtime_environment.RuntimeEnvironmentManager._write_conda_env_to_file",
+    Mock(),
+)
+@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):
+    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)
+
+        dependencies_s3_uri = TEST_S3_BASE_URI + "/additional_dependencies/requirements.txt"
+        job_conda_env = "conda_env"
+        runtime.bootstrap(dependencies_s3_uri, job_conda_env)
+
+        mock_s3_download.assert_called_once_with(
+            dependencies_s3_uri, mock_cwd.return_value, TEST_S3_KMS_KEY, mock_session
+        )
+
+        local_file_path = "/usr/local/path/requirements.txt"
+
+        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
+        )
+        assert call_args == expected_cmd
 
 
 def test_bootstrap_conda_yml():