feature: support conda env creation for pathways

Author: Rohan Gujarathi

src/sagemaker/remote_function/core/runtime_environment.py

@@ -31,6 +31,14 @@ class RuntimeEnvironmentManager:
     def __init__(
         self, s3_base_uri: str = None, s3_kms_key: str = None, sagemaker_session: Session = None
     ):
+        """Initializes RuntimeEnvironemntManager instance
+
+        Args:
+            sagemaker_session (Session): Current sagemaker session
+            s3_base_uri (str): Base S3 URI where dependencies file is uploaded
+            s3_kms_key (str): KMS key to access the S3 bucket
+
+        """
         self.s3_base_uri = s3_base_uri
         self.s3_kms_key = s3_kms_key
         self.sagemaker_session = sagemaker_session
@@ -44,9 +52,7 @@ def snapshot(self, dependencies: str = None):
         user's active conda env and upload the yml file to S3
 
         Args:
-            sagemaker_session (Session): Current sagemaker session
-            s3_base_uri (str): Base S3 URI where dependencies file is uploaded
-            s3_kms_key (str): KMS key to access the S3 bucket
+            dependencies (str): Local path where dependencies file exists.
 
         Returns:
             S3 URI where the dependencies file is uploaded or None
@@ -86,14 +92,15 @@ def bootstrap(self, dependencies_s3_uri: str, job_conda_env: str = None):
         Returns: None
         """
 
-        if dependencies_s3_uri.endswith(".txt"):
-            local_path = os.getcwd()
-            S3Downloader.download(
-                dependencies_s3_uri, local_path, self.s3_kms_key, self.sagemaker_session
-            )
+        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])
+        conda_env = job_conda_env or os.getenv("DEFAULT_CONDA_ENV")
 
-            local_dependencies_file = os.path.join(local_path, dependencies_s3_uri.split("/")[-1])
-            conda_env = job_conda_env or os.getenv("DEFAULT_CONDA_ENV")
+        if dependencies_s3_uri.endswith(".txt"):
             if conda_env:
                 self._install_req_txt_in_conda_env(conda_env, local_dependencies_file)
                 self._write_conda_env_to_file(conda_env)
@@ -102,12 +109,10 @@ def bootstrap(self, dependencies_s3_uri: str, job_conda_env: str = None):
                 self._install_requirements_txt(local_dependencies_file, _python_executable())
 
         elif dependencies_s3_uri.endswith(".yml"):
-            # TODO: implement
-            #   1. verify is conda exists in the image
-            #   2. if job_conda_env: activate and update the conda env with yml
-            #   3. if not, create and activate conda env from conda yml file
-            return
-        return
+            # TODO: implement updating conda env if either job_conda_env or default_conda_env is provided.
+            conda_env_name = "sagemaker-runtime-env"
+            self._create_conda_env(conda_env_name, local_dependencies_file)
+            self._write_conda_env_to_file(conda_env_name)
 
     def _is_file_exists(self, dependencies):
         """Check whether the dependencies file exists at the given location.
@@ -127,10 +132,15 @@ def _install_requirements_txt(self, local_path, python_executable):
         _run_shell_cmd(cmd)
         logger.info("Command {} ran successfully".format(cmd))
 
-    def _create_conda_env(self):
+    def _create_conda_env(self, env_name, local_path):
         """Create conda env using conda yml file"""
-        # TODO: implement
-        pass  # pylint: disable=W0107
+
+        cmd = '/bin/bash -c "{} env create -n {} --file {}"'.format(
+            self._get_conda_exe(), env_name, local_path
+        )
+        logger.info("Creating conda environment {} using: {}.".format(env_name, cmd))
+        _run_shell_cmd(cmd)
+        logger.info("Conda environment {} created successfully.".format(env_name))
 
     def _install_req_txt_in_conda_env(self, env_name, local_path):
         """Install requirements.txt in the given conda environment"""

tests/integ/sagemaker/remote_function/conftest.py

@@ -48,12 +48,22 @@
     "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"
+        && mamba run -n integ_test_env pip install '{source_archive}[remote_function]'\n"
     "ENV SHELL=/bin/bash\n"
     "ENV DEFAULT_CONDA_ENV=base\n"
 )
 
+CONDA_YML_FILE_TEMPLATE = (
+    "name: test_conda_env\n"
+    "channels:\n"
+    "  - defaults\n"
+    "dependencies:\n"
+    "  - scipy=1.10.0\n"
+    "  - pip:\n"
+    "    - sagemaker-{sagemaker_version}.tar.gz[remote_function]\n"
+    "prefix: /opt/conda/bin/conda\n"
+)
+
 
 @pytest.fixture(scope="package")
 def dummy_container_without_error(sagemaker_session):
@@ -74,6 +84,23 @@ def dummy_container_with_conda(sagemaker_session):
     return ecr_uri
 
 
+@pytest.fixture(scope="package")
+def conda_env_yml():
+    """Write conda yml file needed for tests"""
+
+    conda_yml_file_name = "conda_env.yml"
+    with open(os.path.join(os.getcwd(), "VERSION"), "r") as version_file:
+        sagemaker_version = version_file.readline().strip()
+    conda_file_path = os.path.join(os.getcwd(), conda_yml_file_name)
+    with open(conda_file_path, "w") as yml_file:
+        yml_file.writelines(CONDA_YML_FILE_TEMPLATE.format(sagemaker_version=sagemaker_version))
+    yield conda_file_path
+
+    # cleanup
+    if os.path.isfile(conda_yml_file_name):
+        os.remove(conda_yml_file_name)
+
+
 def _build_container(sagemaker_session, py_version, docker_templete):
     """Build a dummy test container locally and push a container to an ecr repo"""
 

tests/integ/sagemaker/remote_function/test_decorator.py

@@ -203,6 +203,24 @@ def cuberoot(x):
         cuberoot(27) == 3
 
 
+def test_additional_dependencies_with_conda_yml_file(
+    sagemaker_session, dummy_container_with_conda, cpu_instance_type, conda_env_yml
+):
+    @remote(
+        role=ROLE,
+        image_uri=dummy_container_with_conda,
+        dependencies=conda_env_yml,
+        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_with_non_existent_dependencies(
     sagemaker_session, dummy_container_without_error, cpu_instance_type
 ):

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

@@ -203,5 +203,36 @@ def test_bootstrap_req_txt_with_conda_env(mock_s3_download, mock_cwd, mock_conda
         assert call_args == expected_cmd
 
 
-def test_bootstrap_conda_yml():
-    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(
+    "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_conda_yml_create_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)
+
+        dependencies_s3_uri = TEST_S3_BASE_URI + "/additional_dependencies/conda_env.yml"
+        runtime.bootstrap(dependencies_s3_uri)
+
+        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/conda_env.yml"
+
+        call_args = popen.call_args[0][0]
+        assert call_args is not None
+
+        expected_cmd = '/bin/bash -c "{} env create -n {} --file {}"'.format(
+            mock_conda_exe.return_value, "sagemaker-runtime-env", local_file_path
+        )
+        assert call_args == expected_cmd