Fix:Use prefix for env auto-capture (aws#909)

Co-authored-by: Mourya Baddam <[email protected]>

Author: 2 people

src/sagemaker/remote_function/job.py

@@ -61,6 +61,7 @@
 # training channel names
 RUNTIME_SCRIPTS_CHANNEL_NAME = "sagemaker_remote_function_bootstrap"
 REMOTE_FUNCTION_WORKSPACE = "sm_rf_user_ws"
+JOB_REMOTE_FUNCTION_WORKSPACE = "sagemaker_remote_function_workspace"
 
 # run context dictionary keys
 KEY_EXPERIMENT_NAME = "experiment_name"
@@ -88,6 +89,16 @@
 printf "INFO: Bootstraping runtime environment.\\n"
 python /opt/ml/input/data/{RUNTIME_SCRIPTS_CHANNEL_NAME}/{BOOTSTRAP_SCRIPT_NAME} "$@"
 
+if [ -d {JOB_REMOTE_FUNCTION_WORKSPACE} ]
+then
+    if [ -f "remote_function_conda_env.txt" ]
+    then
+        cp remote_function_conda_env.txt {JOB_REMOTE_FUNCTION_WORKSPACE}/remote_function_conda_env.txt
+    fi
+    printf "INFO: Changing workspace to {JOB_REMOTE_FUNCTION_WORKSPACE}.\\n"
+    cd {JOB_REMOTE_FUNCTION_WORKSPACE}
+fi
+
 if [ -f "remote_function_conda_env.txt" ]
 then
     conda_env=$(cat remote_function_conda_env.txt)
@@ -553,9 +564,11 @@ def _prepare_and_upload_dependencies(
     ):
         return None
 
-    with _tmpdir() as tmp_workspace:
+    with _tmpdir() as tmp_dir:
+        tmp_workspace_dir = os.path.join(tmp_dir, "temp_workspace/")
+        os.mkdir(tmp_workspace_dir)
         # TODO Remove the following hack to avoid dir_exists error in the copy_tree call below.
-        tmp_workspace = os.path.join(tmp_workspace, "remote_function/")
+        tmp_workspace = os.path.join(tmp_workspace_dir, JOB_REMOTE_FUNCTION_WORKSPACE)
 
         if include_local_workdir:
             shutil.copytree(
@@ -593,8 +606,10 @@ def _prepare_and_upload_dependencies(
                     pre_execution_script,
                 )
 
-        workspace_archive_path = os.path.join(tmp_workspace, "workspace")
-        workspace_archive_path = shutil.make_archive(workspace_archive_path, "zip", tmp_workspace)
+        workspace_archive_path = os.path.join(tmp_dir, "workspace")
+        workspace_archive_path = shutil.make_archive(
+            workspace_archive_path, "zip", tmp_workspace_dir
+        )
         logger.info("Successfully created workdir archive at '%s'", workspace_archive_path)
 
         upload_path = S3Uploader.upload(

src/sagemaker/remote_function/runtime_environment/bootstrap_runtime_environment.py

@@ -34,6 +34,7 @@
 BASE_CHANNEL_PATH = "/opt/ml/input/data"
 FAILURE_REASON_PATH = "/opt/ml/output/failure"
 PRE_EXECUTION_SCRIPT_NAME = "pre_exec.sh"
+JOB_REMOTE_FUNCTION_WORKSPACE = "sagemaker_remote_function_workspace"
 
 
 logger = get_logger()
@@ -94,6 +95,7 @@ def _bootstrap_runtime_environment(
     workspace_unpack_dir = pathlib.Path(os.getcwd()).absolute()
     shutil.unpack_archive(filename=workspace_archive_path, extract_dir=workspace_unpack_dir)
     logger.info("Successfully unpacked workspace archive at '%s'.", workspace_unpack_dir)
+    workspace_unpack_dir = pathlib.Path(workspace_unpack_dir, JOB_REMOTE_FUNCTION_WORKSPACE)
 
     # Handle pre-execution commands
     path_to_pre_exec_script = os.path.join(workspace_unpack_dir, PRE_EXECUTION_SCRIPT_NAME)

src/sagemaker/remote_function/runtime_environment/runtime_environment_manager.py

@@ -95,28 +95,12 @@ def _capture_from_local_runtime(self) -> str:
 
         # Try to capture dependencies from the conda environment, if any.
         conda_env_name = self._get_active_conda_env_name()
-        logger.info("Found conda_env_name: '%s'", conda_env_name)
-        conda_env_prefix = None
-
-        if conda_env_name is None:
-            conda_env_prefix = self._get_active_conda_env_prefix()
+        conda_env_prefix = self._get_active_conda_env_prefix()
+        if conda_env_name:
+            logger.info("Found conda_env_name: '%s'", conda_env_name)
+        elif conda_env_prefix:
             logger.info("Found conda_env_prefix: '%s'", conda_env_prefix)
-            if conda_env_prefix is None and self._get_studio_image_uri() is not None:
-                logger.info(
-                    "Neither conda env name or prefix is set. Running Studio fallback logic"
-                )
-                # Fallback for Studio Notebooks since conda env is not activated to use as a
-                # Jupyter kernel from images.
-                # TODO: Remove after fixing the behavior for Studio Notebooks.
-                which_python = self._get_which_python()
-                prefix_candidate = which_python.replace("/bin/python", "")
-                conda_env_list = self._get_conda_envs_list()
-                if (
-                    conda_env_list.find(prefix_candidate + "\n") > 0
-                ):  # need "\n" to match exact prefix; -1 for not found.
-                    conda_env_prefix = prefix_candidate
-
-        if conda_env_name is None and conda_env_prefix is None:
+        else:
             raise ValueError("No conda environment seems to be active.")
 
         if conda_env_name == "base":
@@ -126,25 +110,10 @@ def _capture_from_local_runtime(self) -> str:
             )
 
         local_dependencies_path = os.path.join(os.getcwd(), "env_snapshot.yml")
-        if conda_env_name is not None:
-            self._export_conda_env_from_env_name(conda_env_name, local_dependencies_path)
-        else:
-            self._export_conda_env_from_prefix(conda_env_prefix, local_dependencies_path)
+        self._export_conda_env_from_prefix(conda_env_prefix, local_dependencies_path)
 
         return local_dependencies_path
 
-    def _get_conda_envs_list(self) -> str:
-        """Returns the registered list of conda environments on the system."""
-        return _run_and_get_output_shell_cmd(f"{self._get_conda_exe()} env list")
-
-    def _get_which_python(self) -> str:
-        """Return the location of the current Python interpreter."""
-        return _python_executable()
-
-    def _get_studio_image_uri(self) -> str:
-        """Returns the Sagemaker Image URI from the set environment variable. None otherwise."""
-        return os.getenv("SAGEMAKER_INTERNAL_IMAGE_URI")
-
     def _get_active_conda_env_prefix(self) -> str:
         """Returns the conda prefix from the set environment variable. None otherwise."""
         return os.getenv("CONDA_PREFIX")
@@ -242,14 +211,6 @@ def _update_conda_env(self, env_name, local_path):
         _run_shell_cmd(cmd)
         logger.info("Conda env %s updated succesfully", env_name)
 
-    def _export_conda_env_from_env_name(self, env_name, local_path):
-        """Export the conda env to a conda yml file"""
-
-        cmd = f"{self._get_conda_exe()} env export -n {env_name} --no-builds > {local_path}"
-        logger.info("Exporting conda environment: %s", cmd)
-        _run_shell_cmd(cmd)
-        logger.info("Conda environment %s exported successfully", env_name)
-
     def _export_conda_env_from_prefix(self, prefix, local_path):
         """Export the conda env to a conda yml file"""
 
@@ -324,9 +285,13 @@ def _run_pre_execution_command_script(script_path: str):
 
     Raises RuntimeEnvironmentError if the shell script fails
     """
+    current_dir = os.path.dirname(script_path)
 
     process = subprocess.Popen(
-        ["/bin/bash", "-eu", script_path], stdout=subprocess.PIPE, stderr=subprocess.PIPE
+        ["/bin/bash", "-eu", script_path],
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        cwd=current_dir,
     )
 
     _log_output(process)

tests/integ/sagemaker/remote_function/conftest.py

@@ -27,7 +27,6 @@
 
 DOCKERFILE_TEMPLATE = (
     "FROM public.ecr.aws/docker/library/python:{py_version}-slim\n\n"
-    "WORKDIR /opt/ml/remote_function/\n"
     "RUN apt-get update -y \
         && apt-get install -y unzip curl\n\n"
     "RUN curl 'https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip' -o 'awscliv2.zip' \
@@ -40,7 +39,6 @@
 
 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"
@@ -66,7 +64,7 @@
     "dependencies:\n"
     "  - scipy=1.7.3\n"
     "  - pip:\n"
-    "    - /opt/ml/remote_function/sagemaker-{sagemaker_version}.tar.gz\n"
+    "    - /sagemaker-{sagemaker_version}.tar.gz\n"
     "prefix: /opt/conda/bin/conda\n"
 )
 

tests/unit/sagemaker/remote_function/runtime_environment/test_bootstrap_runtime_environment.py

@@ -21,7 +21,8 @@
 import pathlib
 
 TEST_JOB_CONDA_ENV = "conda_env"
-TEST_DEPENDENCIES_PATH = "/user/set/workdir"
+CURR_WORKING_DIR = "/user/set/workdir"
+TEST_DEPENDENCIES_PATH = "/user/set/workdir/sagemaker_remote_function_workspace"
 TEST_PYTHON_VERSION = "3.10"
 TEST_WORKSPACE_ARCHIVE_DIR_PATH = "/opt/ml/input/data/sm_rf_user_ws"
 TEST_WORKSPACE_ARCHIVE_PATH = "/opt/ml/input/data/sm_rf_user_ws/workspace.zip"
@@ -44,7 +45,7 @@ def mock_args():
 )
 @patch("sys.exit")
 @patch("shutil.unpack_archive", Mock())
-@patch("os.getcwd", return_value=TEST_DEPENDENCIES_PATH)
+@patch("os.getcwd", return_value=CURR_WORKING_DIR)
 @patch("os.path.exists", return_value=True)
 @patch("os.path.isfile", return_value=True)
 @patch("os.listdir", return_value=["fileA.py", "fileB.sh", "requirements.txt"])
@@ -175,7 +176,7 @@ def test_main_no_workspace_archive(
 @patch("shutil.unpack_archive", Mock())
 @patch("os.path.exists", return_value=True)
 @patch("os.path.isfile", return_value=True)
-@patch("os.getcwd", return_value=TEST_DEPENDENCIES_PATH)
+@patch("os.getcwd", return_value=CURR_WORKING_DIR)
 @patch("os.listdir", return_value=["fileA.py", "fileB.sh"])
 @patch(
     "sagemaker.remote_function.runtime_environment.runtime_environment_manager."

tests/unit/sagemaker/remote_function/runtime_environment/test_runtime_environment_manager.py

@@ -90,17 +90,6 @@ def test__get_active_conda_env_prefix():
         assert call_arg == "CONDA_PREFIX"
 
 
-def test__get_studio_image_uri():
-    with patch("os.getenv") as getenv_patch:
-        getenv_patch.return_value = "some-sagemaker-studio-image-ecr-uri"
-
-        result = RuntimeEnvironmentManager()._get_studio_image_uri()
-
-        assert result == "some-sagemaker-studio-image-ecr-uri"
-        call_arg = getenv_patch.call_args[0][0]
-        assert call_arg == "SAGEMAKER_INTERNAL_IMAGE_URI"
-
-
 @patch(
     "sagemaker.remote_function.runtime_environment.runtime_environment_manager._log_error", Mock()
 )
@@ -117,7 +106,14 @@ def test__get_studio_image_uri():
     ".RuntimeEnvironmentManager._get_active_conda_env_name",
     return_value="test_env",
 )
-def test_snapshot_from_active_conda_env_when_name_available(conda_default_env, stub_conda_exe):
+@patch(
+    "sagemaker.remote_function.runtime_environment.runtime_environment_manager"
+    ".RuntimeEnvironmentManager._get_active_conda_env_prefix",
+    return_value="/some/conda/env/prefix",
+)
+def test_snapshot_from_active_conda_env_when_name_available(
+    conda_env_prefix, conda_default_env, stub_conda_exe
+):
     expected_result = os.path.join(os.getcwd(), "env_snapshot.yml")
     with patch("subprocess.Popen") as popen:
         popen.return_value.wait.return_value = 0
@@ -128,7 +124,7 @@ def test_snapshot_from_active_conda_env_when_name_available(conda_default_env, s
         call_args = popen.call_args[0][0]
         assert call_args is not None
         expected_cmd = (
-            f"{stub_conda_exe.return_value} env export -n {conda_default_env.return_value} "
+            f"{stub_conda_exe.return_value} env export -p {conda_env_prefix.return_value} "
             f"--no-builds > {expected_result}"
         )
         assert call_args == expected_cmd
@@ -189,81 +185,7 @@ def test_snapshot_from_active_conda_env_when_prefix_available(
     ".RuntimeEnvironmentManager._get_active_conda_env_prefix",
     return_value=None,
 )
-@patch(
-    "sagemaker.remote_function.runtime_environment.runtime_environment_manager"
-    ".RuntimeEnvironmentManager._get_studio_image_uri",
-    return_value=None,
-)
-def test_snapshot_auto_capture_non_studio_no_active_conda_env(
-    no_studio_image_uri, no_conda_env_prefix, no_conda_env_name
-):
-    with pytest.raises(ValueError):
-        RuntimeEnvironmentManager().snapshot("auto_capture")
-
-
-@patch(
-    "sagemaker.remote_function.runtime_environment.runtime_environment_manager._log_error", Mock()
-)
-@patch(
-    "sagemaker.remote_function.runtime_environment.runtime_environment_manager._log_output", Mock()
-)
-@patch(
-    "sagemaker.remote_function.runtime_environment.runtime_environment_manager"
-    ".RuntimeEnvironmentManager._get_conda_exe",
-    return_value="some-exe",
-)
-@patch(
-    "sagemaker.remote_function.runtime_environment.runtime_environment_manager"
-    ".RuntimeEnvironmentManager._get_active_conda_env_name",
-    return_value=None,
-)
-@patch(
-    "sagemaker.remote_function.runtime_environment.runtime_environment_manager"
-    ".RuntimeEnvironmentManager._get_active_conda_env_prefix",
-    return_value=None,
-)
-@patch(
-    "sagemaker.remote_function.runtime_environment.runtime_environment_manager"
-    ".RuntimeEnvironmentManager._get_studio_image_uri",
-    return_value="some-image-uri",
-)
-@patch(
-    "sagemaker.remote_function.runtime_environment.runtime_environment_manager"
-    ".RuntimeEnvironmentManager._get_which_python",
-    return_value="/some/prefix/bin/python",
-)
-@patch(
-    "sagemaker.remote_function.runtime_environment.runtime_environment_manager"
-    ".RuntimeEnvironmentManager._get_conda_envs_list",
-    return_value="something something /some/prefix\n something /someother/prefix\n\n",
-)
-def test_snapshot_auto_capture_in_studio_no_active_conda_env(
-    ouptut_conda_env_list,
-    output_which_python,
-    some_studio_image_uri,
-    no_conda_env_prefix,
-    no_conda_env_name,
-    stub_conda_exe,
-):
-    expected_result = os.path.join(os.getcwd(), "env_snapshot.yml")
-    with patch("subprocess.Popen") as popen:
-        popen.return_value.wait.return_value = 0
-
-        result = RuntimeEnvironmentManager().snapshot("auto_capture")
-        assert result == expected_result
-
-        call_args = popen.call_args[0][0]
-        assert call_args is not None
-        expected_cmd = f"{stub_conda_exe.return_value} env export -p /some/prefix --no-builds > {expected_result}"
-        assert call_args == expected_cmd
-
-
-@patch(
-    "sagemaker.remote_function.runtime_environment.runtime_environment_manager"
-    ".RuntimeEnvironmentManager._get_active_conda_env_name",
-    return_value=None,
-)
-def test_snapshot_from_active_conda_env_error(conda_default_env):
+def test_snapshot_auto_capture_no_active_conda_env(no_conda_env_prefix, no_conda_env_name):
     with pytest.raises(ValueError):
         RuntimeEnvironmentManager().snapshot("auto_capture")