fix: use copy instead of move in bootstrap script (#1339)

Co-authored-by: Rohan Gujarathi <[email protected]>

Author: 2 people

src/sagemaker/remote_function/runtime_environment/bootstrap_runtime_environment.py

@@ -140,7 +140,7 @@ def _bootstrap_runtime_env_for_pipeline_step(
     for file in os.listdir(pre_exec_script_and_dependencies_dir):
         src_path = os.path.join(pre_exec_script_and_dependencies_dir, file)
         dest_path = os.path.join(workspace_dir, file)
-        shutil.move(src_path, dest_path)
+        shutil.copy(src_path, dest_path)
 
     _handle_pre_exec_scripts(workspace_dir)
 

tests/integ/sagemaker/conftest.py

@@ -61,6 +61,21 @@
     "ENV SAGEMAKER_JOB_CONDA_ENV=default_env\n"
 )
 
+DOCKERFILE_TEMPLATE_WITH_USER_AND_WORKDIR = (
+    "FROM public.ecr.aws/docker/library/python:{py_version}-slim\n\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' \
+        && unzip awscliv2.zip \
+        && ./aws/install\n\n"
+    "RUN useradd -ms /bin/bash integ-test-user\n"
+    "USER integ-test-user\n"
+    "WORKDIR /home/integ-test-user\n"
+    "COPY {source_archive} ./\n"
+    "RUN pip install '{source_archive}'\n"
+    "RUN rm {source_archive}\n"
+)
+
 AUTO_CAPTURE_CLIENT_DOCKER_TEMPLATE = (
     "FROM public.ecr.aws/docker/library/python:{py_version}-slim\n\n"
     'SHELL ["/bin/bash", "-c"]\n'
@@ -105,6 +120,14 @@ def dummy_container_without_error(sagemaker_session, compatible_python_version):
     return ecr_uri
 
 
+@pytest.fixture(scope="session")
+def dummy_container_with_user_and_workdir(sagemaker_session, compatible_python_version):
+    ecr_uri = _build_container(
+        sagemaker_session, compatible_python_version, DOCKERFILE_TEMPLATE_WITH_USER_AND_WORKDIR
+    )
+    return ecr_uri
+
+
 @pytest.fixture(scope="session")
 def dummy_container_incompatible_python_runtime(sagemaker_session, incompatible_python_version):
     ecr_uri = _build_container(sagemaker_session, incompatible_python_version, DOCKERFILE_TEMPLATE)

tests/integ/sagemaker/remote_function/test_decorator.py

@@ -727,6 +727,26 @@ def load_checkpoints(checkpoint_path: Union[str, os.PathLike]):
     )
 
 
+def test_with_user_and_workdir_set_in_the_image(
+    sagemaker_session, dummy_container_with_user_and_workdir, cpu_instance_type
+):
+    dependencies_path = os.path.join(DATA_DIR, "remote_function", "requirements.txt")
+
+    @remote(
+        role=ROLE,
+        image_uri=dummy_container_with_user_and_workdir,
+        dependencies=dependencies_path,
+        instance_type=cpu_instance_type,
+        sagemaker_session=sagemaker_session,
+    )
+    def cuberoot(x):
+        from scipy.special import cbrt
+
+        return cbrt(27)
+
+    assert cuberoot(27) == 3
+
+
 @pytest.mark.skip
 def test_decorator_with_spark_job(sagemaker_session, cpu_instance_type):
     @remote(

tests/integ/sagemaker/workflow/test_step_decorator.py

@@ -813,3 +813,48 @@ def updated_func(x):
             pipeline.delete()
         except Exception:
             pass
+
+
+def test_with_user_and_workdir_set_in_the_image(
+    sagemaker_session, role, pipeline_name, region_name, dummy_container_with_user_and_workdir
+):
+    os.environ["AWS_DEFAULT_REGION"] = region_name
+    dependencies_path = os.path.join(DATA_DIR, "workflow", "requirements.txt")
+
+    @step(
+        role=role,
+        image_uri=dummy_container_with_user_and_workdir,
+        dependencies=dependencies_path,
+        instance_type=INSTANCE_TYPE,
+    )
+    def cuberoot(x):
+        from scipy.special import cbrt
+
+        return cbrt(x)
+
+    step_a = cuberoot(8)
+
+    pipeline = Pipeline(
+        name=pipeline_name,
+        steps=[step_a],
+        sagemaker_session=sagemaker_session,
+    )
+
+    try:
+        create_and_execute_pipeline(
+            pipeline=pipeline,
+            pipeline_name=pipeline_name,
+            region_name=region_name,
+            role=role,
+            no_of_steps=1,
+            last_step_name="cuberoot",
+            execution_parameters=dict(),
+            step_status="Succeeded",
+            step_result_type=numpy.float64,
+            step_result_value=2.0,
+        )
+    finally:
+        try:
+            pipeline.delete()
+        except Exception:
+            pass

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

@@ -370,7 +370,7 @@ def test_bootstrap_runtime_env_for_remote_no_workspace(
     install_dependencies.assert_not_called()
 
 
-@patch("shutil.move")
+@patch("shutil.copy")
 @patch("os.listdir", return_value=[PRE_EXECUTION_SCRIPT_NAME, TEST_DEPENDENCY_FILE_NAME])
 @patch("os.path.exists", return_value=True)
 @patch(
@@ -384,7 +384,7 @@ def test_bootstrap_runtime_env_for_remote_no_workspace(
     return_value=pathlib.Path(TEST_DEPENDENCIES_PATH),
 )
 def test_bootstrap_runtime_env_for_step(
-    unpack_workspace, install_scripts, install_dependencies, path_exists, list_dir, move_files
+    unpack_workspace, install_scripts, install_dependencies, path_exists, list_dir, copy_files
 ):
     bootstrap._bootstrap_runtime_env_for_pipeline_step(
         TEST_PYTHON_VERSION,
@@ -397,7 +397,7 @@ def test_bootstrap_runtime_env_for_step(
     dependency_dir = TEST_BASE_CHANNEL_PATH + PIPELINE_STEP_CHANNEL
     path_exists.assert_called_once_with(dependency_dir)
     list_dir.assert_called_once_with(dependency_dir)
-    assert move_files.call_count == 2
+    assert copy_files.call_count == 2
     install_scripts.assert_called_once_with(unpack_workspace.return_value)
     install_dependencies.assert_called_once_with(
         unpack_workspace.return_value,
@@ -410,7 +410,7 @@ def test_bootstrap_runtime_env_for_step(
 
 @patch("os.getcwd", return_value=CURR_WORKING_DIR)
 @patch("os.mkdir")
-@patch("shutil.move")
+@patch("shutil.copy")
 @patch("os.listdir", return_value=[PRE_EXECUTION_SCRIPT_NAME, TEST_DEPENDENCY_FILE_NAME])
 @patch("os.path.exists", return_value=True)
 @patch(
@@ -429,7 +429,7 @@ def test_bootstrap_runtime_env_for_step_no_workspace(
     install_dependencies,
     path_exists,
     list_dir,
-    move_files,
+    copy_files,
     mkdir,
     get_cwd,
 ):
@@ -447,7 +447,7 @@ def test_bootstrap_runtime_env_for_step_no_workspace(
     dependency_dir = TEST_BASE_CHANNEL_PATH + PIPELINE_STEP_CHANNEL
     path_exists.assert_called_once_with(dependency_dir)
     list_dir.assert_called_once_with(dependency_dir)
-    assert move_files.call_count == 2
+    assert copy_files.call_count == 2
     install_scripts.assert_called_once_with(pathlib.Path(TEST_DEPENDENCIES_PATH))
     install_dependencies.assert_called_once_with(
         pathlib.Path(TEST_DEPENDENCIES_PATH),