fix: Add write permission to job output dirs for remote and step decorator running on non-root job user

Author: qidewenwhen

src/sagemaker/remote_function/runtime_environment/bootstrap_runtime_environment.py

@@ -14,6 +14,8 @@
 from __future__ import absolute_import
 
 import argparse
+import getpass
+import subprocess
 import sys
 import os
 import shutil
@@ -24,12 +26,14 @@
         RuntimeEnvironmentManager,
         _DependencySettings,
         get_logger,
+        RuntimeEnvironmentError,
     )
 else:
     from sagemaker.remote_function.runtime_environment.runtime_environment_manager import (
         RuntimeEnvironmentManager,
         _DependencySettings,
         get_logger,
+        RuntimeEnvironmentError,
     )
 
 SUCCESS_EXIT_CODE = 0
@@ -38,6 +42,7 @@
 REMOTE_FUNCTION_WORKSPACE = "sm_rf_user_ws"
 BASE_CHANNEL_PATH = "/opt/ml/input/data"
 FAILURE_REASON_PATH = "/opt/ml/output/failure"
+JOB_OUTPUT_DIRS_NEEDED_PERMISSION_CHANGE = ["/opt/ml/output", "/opt/ml/model", "/tmp"]
 PRE_EXECUTION_SCRIPT_NAME = "pre_exec.sh"
 JOB_REMOTE_FUNCTION_WORKSPACE = "sagemaker_remote_function_workspace"
 SCRIPT_AND_DEPENDENCIES_CHANNEL_NAME = "pre_exec_script_and_dependencies"
@@ -63,6 +68,17 @@ def main(sys_args=None):
 
         RuntimeEnvironmentManager()._validate_python_version(client_python_version, conda_env)
 
+        user = getpass.getuser()
+        if user != "root":
+            log_message = (
+                "The job is running on non-root user: %s. Adding write permissions to the "
+                "following job output directories: %s."
+            )
+            logger.info(log_message, user, JOB_OUTPUT_DIRS_NEEDED_PERMISSION_CHANGE)
+            _change_dir_permission(
+                dirs=JOB_OUTPUT_DIRS_NEEDED_PERMISSION_CHANGE, new_permission="777"
+            )
+
         if pipeline_execution_id:
             _bootstrap_runtime_env_for_pipeline_step(
                 client_python_version, func_step_workspace, conda_env, dependency_settings
@@ -81,6 +97,33 @@ def main(sys_args=None):
         sys.exit(exit_code)
 
 
+def _change_dir_permission(dirs: list, new_permission: str):
+    """Change the permission of given directories
+
+    Args:
+        dirs (list[str]): A list of directories for permission update.
+        new_permission (str): The new permission for the given directories.
+    """
+
+    _ERROR_MSG_PREFIX = "Failed to change directory permissions due to: "
+    command = ["sudo", "chmod", "-R", new_permission] + dirs
+    logger.info("Executing '%s'.", {" ".join(command)})
+
+    try:
+        subprocess.run(command, check=True, stderr=subprocess.PIPE)
+    except subprocess.CalledProcessError as called_process_err:
+        err_msg = called_process_err.stderr.decode("utf-8")
+        raise RuntimeEnvironmentError(f"{_ERROR_MSG_PREFIX} {err_msg}")
+    except FileNotFoundError as file_not_found_err:
+        if "[Errno 2] No such file or directory: 'sudo'" in str(file_not_found_err):
+            raise RuntimeEnvironmentError(
+                f"{_ERROR_MSG_PREFIX} {file_not_found_err}. "
+                "Please contact the image owner to install 'sudo' in the job container "
+                "and provide sudo privilege to the container user."
+            )
+        raise RuntimeEnvironmentError(file_not_found_err)
+
+
 def _bootstrap_runtime_env_for_remote_function(
     client_python_version: str,
     conda_env: str = None,

tests/integ/sagemaker/conftest.py

@@ -68,7 +68,12 @@
     "RUN curl 'https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip' -o 'awscliv2.zip' \
         && unzip awscliv2.zip \
         && ./aws/install\n\n"
+    "RUN apt install sudo\n"
     "RUN useradd -ms /bin/bash integ-test-user\n"
+    # Add the user to sudo group
+    "RUN usermod -aG sudo integ-test-user\n"
+    # Ensure passwords are not required for sudo group users
+    "RUN echo '%sudo ALL= (ALL) NOPASSWD:ALL' >> /etc/sudoers\n"
     "USER integ-test-user\n"
     "WORKDIR /home/integ-test-user\n"
     "COPY {source_archive} ./\n"

tests/integ/sagemaker/remote_function/test_decorator.py

@@ -747,6 +747,25 @@ def cuberoot(x):
     assert cuberoot(27) == 3
 
 
+def test_with_user_and_workdir_set_in_the_image_client_error_case(
+    sagemaker_session, dummy_container_with_user_and_workdir, cpu_instance_type
+):
+    client_error_message = "Testing client error in job."
+
+    @remote(
+        role=ROLE,
+        image_uri=dummy_container_with_user_and_workdir,
+        instance_type=cpu_instance_type,
+        sagemaker_session=sagemaker_session,
+    )
+    def my_func():
+        raise RuntimeError(client_error_message)
+
+    with pytest.raises(RuntimeError) as error:
+        my_func()
+    assert client_error_message in str(error)
+
+
 @pytest.mark.skip
 def test_decorator_with_spark_job(sagemaker_session, cpu_instance_type):
     @remote(

tests/integ/sagemaker/workflow/test_step_decorator.py

@@ -858,3 +858,49 @@ def cuberoot(x):
             pipeline.delete()
         except Exception:
             pass
+
+
+def test_with_user_and_workdir_set_in_the_image_client_error_case(
+    sagemaker_session, role, pipeline_name, region_name, dummy_container_with_user_and_workdir
+):
+    # This test aims to ensure client error in step decorated function
+    # can be successfully surfaced and the job can be failed.
+    os.environ["AWS_DEFAULT_REGION"] = region_name
+    client_error_message = "Testing client error in job."
+
+    @step(
+        role=role,
+        image_uri=dummy_container_with_user_and_workdir,
+        instance_type=INSTANCE_TYPE,
+    )
+    def my_func():
+        raise RuntimeError(client_error_message)
+
+    step_a = my_func()
+
+    pipeline = Pipeline(
+        name=pipeline_name,
+        steps=[step_a],
+        sagemaker_session=sagemaker_session,
+    )
+
+    try:
+        _, execution_steps = create_and_execute_pipeline(
+            pipeline=pipeline,
+            pipeline_name=pipeline_name,
+            region_name=region_name,
+            role=role,
+            no_of_steps=1,
+            last_step_name=get_step(step_a).name,
+            execution_parameters=dict(),
+            step_status="Failed",
+        )
+        assert (
+            f"ClientError: AlgorithmError: RuntimeError('{client_error_message}')"
+            in execution_steps[0]["FailureReason"]
+        )
+    finally:
+        try:
+            pipeline.delete()
+        except Exception:
+            pass