pre-execution command support (aws#905)

Co-authored-by: Ao Guo <[email protected]>

Author: 2 people

src/sagemaker/config/config_schema.py

@@ -46,6 +46,8 @@
 MODULES = "Modules"
 REMOTE_FUNCTION = "RemoteFunction"
 DEPENDENCIES = "Dependencies"
+PRE_EXECUTION_SCRIPT = "PreExecutionScript"
+PRE_EXECUTION_COMMANDS = "PreExecutionCommands"
 ENVIRONMENT_VARIABLES = "EnvironmentVariables"
 IMAGE_URI = "ImageUri"
 INCLUDE_LOCAL_WORKDIR = "IncludeLocalWorkDir"
@@ -233,6 +235,12 @@ def _simple_path(*args: str):
 REMOTE_FUNCTION_DEPENDENCIES = _simple_path(
     SAGEMAKER, PYTHON_SDK, MODULES, REMOTE_FUNCTION, DEPENDENCIES
 )
+REMOTE_FUNCTION_PRE_EXECUTION_COMMANDS = _simple_path(
+    SAGEMAKER, PYTHON_SDK, MODULES, REMOTE_FUNCTION, PRE_EXECUTION_COMMANDS
+)
+REMOTE_FUNCTION_PRE_EXECUTION_SCRIPT = _simple_path(
+    SAGEMAKER, PYTHON_SDK, MODULES, REMOTE_FUNCTION, PRE_EXECUTION_SCRIPT
+)
 REMOTE_FUNCTION_ENVIRONMENT_VARIABLES = _simple_path(
     SAGEMAKER, PYTHON_SDK, MODULES, REMOTE_FUNCTION, ENVIRONMENT_VARIABLES
 )
@@ -266,9 +274,6 @@ def _simple_path(*args: str):
 REMOTE_FUNCTION_ENABLE_INTER_CONTAINER_TRAFFIC_ENCRYPTION = _simple_path(
     SAGEMAKER, PYTHON_SDK, MODULES, REMOTE_FUNCTION, ENABLE_INTER_CONTAINER_TRAFFIC_ENCRYPTION
 )
-REMOTE_FUNCTION_ENABLE_NETWORK_ISOLATION = _simple_path(
-    SAGEMAKER, PYTHON_SDK, MODULES, REMOTE_FUNCTION, ENABLE_NETWORK_ISOLATION
-)
 
 # Paths for reference elsewhere in the SDK.
 # Names include the schema version since the paths could change with other schema versions
@@ -440,6 +445,8 @@ def _simple_path(*args: str):
         },
         # Regex is taken from https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_S3DataSource.html#sagemaker-Type-S3DataSource-S3Uri
         "s3Uri": {TYPE: "string", "pattern": "^(https|s3)://([^/]+)/?(.*)$", "maxLength": 1024},
+        # Regex is taken from https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_AlgorithmSpecification.html#sagemaker-Type-AlgorithmSpecification-ContainerEntrypoint
+        "preExecutionCommand": {TYPE: "string", "pattern": r".*"},
     },
     PROPERTIES: {
         SCHEMA_VERSION: {
@@ -475,10 +482,14 @@ def _simple_path(*args: str):
                                     ADDITIONAL_PROPERTIES: False,
                                     PROPERTIES: {
                                         DEPENDENCIES: {TYPE: "string"},
+                                        PRE_EXECUTION_COMMANDS: {
+                                            TYPE: "array",
+                                            "items": {"$ref": "#/definitions/preExecutionCommand"},
+                                        },
+                                        PRE_EXECUTION_SCRIPT: {TYPE: "string"},
                                         ENABLE_INTER_CONTAINER_TRAFFIC_ENCRYPTION: {
                                             TYPE: "boolean"
                                         },
-                                        ENABLE_NETWORK_ISOLATION: {TYPE: "boolean"},
                                         ENVIRONMENT_VARIABLES: {
                                             "$ref": "#/definitions/environmentVariables"
                                         },

src/sagemaker/remote_function/client.py

@@ -59,6 +59,8 @@ def remote(
     _func=None,
     *,
     dependencies: str = None,
+    pre_execution_commands: List[str] = None,
+    pre_execution_script: str = None,
     environment_variables: Dict[str, str] = None,
     image_uri: str = None,
     include_local_workdir: bool = False,
@@ -79,14 +81,19 @@ def remote(
     volume_kms_key: str = None,
     volume_size: int = 30,
     encrypt_inter_container_traffic: bool = None,
-    enable_network_isolation: bool = None,
 ):
     """Function that starts a new SageMaker job synchronously with overridden runtime settings.
 
     Args:
         _func (Optional): Python function to be executed on the SageMaker job runtime environment.
         dependencies (str): Path to dependencies file or a reserved keyword
             ``auto_capture``. Defaults to None.
+        pre_execution_commands (List[str]): List of commands to be executed prior to executing
+            remote function. Only one of ``pre_execution_commands`` or ``pre_execution_script``
+            can be specified at the same time. Defaults to None.
+        pre_execution_script (str): Path to script file to be executed prior to executing
+            remote function. Only one of ``pre_execution_commands`` or ``pre_execution_script``
+            can be specified at the same time. Defaults to None.
         environment_variables (Dict): environment variables
         image_uri (str): Docker image URI on ECR.
         include_local_workdir (bool): Set to ``True`` if the remote function code imports local
@@ -118,11 +125,6 @@ def remote(
             data. Default is 30.
         encrypt_inter_container_traffic (bool): Specifies whether traffic between training
             containers is encrypted for the training job. (default: ``False``).
-        enable_network_isolation (bool): Specifies whether container will
-            run in network isolation mode (default: ``False``). Network
-            isolation mode restricts the container access to outside networks
-            (such as the Internet). The container does not make any inbound or
-            outbound network calls. Also known as Internet-free mode.
     """
 
     def _remote(func):
@@ -133,6 +135,8 @@ def wrapper(*args, **kwargs):
 
             job_settings = _JobSettings(
                 dependencies=dependencies,
+                pre_execution_commands=pre_execution_commands,
+                pre_execution_script=pre_execution_script,
                 environment_variables=environment_variables,
                 image_uri=image_uri,
                 include_local_workdir=include_local_workdir,
@@ -153,7 +157,6 @@ def wrapper(*args, **kwargs):
                 volume_kms_key=volume_kms_key,
                 volume_size=volume_size,
                 encrypt_inter_container_traffic=encrypt_inter_container_traffic,
-                enable_network_isolation=enable_network_isolation,
             )
             job = _Job.start(job_settings, func, args, kwargs)
 
@@ -322,6 +325,8 @@ def __init__(
         self,
         *,
         dependencies: str = None,
+        pre_execution_commands: List[str] = None,
+        pre_execution_script: str = None,
         environment_variables: Dict[str, str] = None,
         image_uri: str = None,
         include_local_workdir: bool = False,
@@ -343,13 +348,18 @@ def __init__(
         volume_kms_key: str = None,
         volume_size: int = 30,
         encrypt_inter_container_traffic: bool = None,
-        enable_network_isolation: bool = None,
     ):
         """Initiates a ``RemoteExecutor`` instance.
 
         Args:
             dependencies (str): Path to dependencies file or a reserved keyword
                 ``auto_capture``. Defaults to None.
+            pre_execution_commands (List[str]): List of commands to be executed prior to executing
+                remote function. Only one of ``pre_execution_commands`` or ``pre_execution_script``
+                can be specified at the same time. Defaults to None.
+            pre_execution_script (str): Path to script file to be executed prior to executing
+                remote function. Only one of ``pre_execution_commands`` or ``pre_execution_script``
+                can be specified at the same time. Defaults to None.
             environment_variables (Dict): Environment variables passed to the underlying sagemaker
                 job. Defaults to None
             image_uri (str): Docker image URI on ECR. Defaults to base Python image.
@@ -388,11 +398,6 @@ def __init__(
                 data. Defaults to 30.
             encrypt_inter_container_traffic (bool): Specifies whether traffic between training
                 containers is encrypted for the training job. (default: ``False``).
-            enable_network_isolation (bool): Specifies whether container will
-                run in network isolation mode (default: ``False``). Network
-                isolation mode restricts the container access to outside networks
-                (such as the Internet). The container does not make any inbound or
-                outbound network calls. Also known as Internet-free mode.
         """
         self.max_parallel_jobs = max_parallel_jobs
 
@@ -401,6 +406,8 @@ def __init__(
 
         self.job_settings = _JobSettings(
             dependencies=dependencies,
+            pre_execution_commands=pre_execution_commands,
+            pre_execution_script=pre_execution_script,
             environment_variables=environment_variables,
             image_uri=image_uri,
             include_local_workdir=include_local_workdir,
@@ -421,7 +428,6 @@ def __init__(
             volume_kms_key=volume_kms_key,
             volume_size=volume_size,
             encrypt_inter_container_traffic=encrypt_inter_container_traffic,
-            enable_network_isolation=enable_network_isolation,
         )
 
         self._state_condition = threading.Condition()

src/sagemaker/remote_function/core/serialization.py

@@ -106,9 +106,10 @@ def serialize(obj: Any, sagemaker_session, s3_uri: str, s3_kms_key: str = None):
                 e, NotImplementedError
             ) and "Instance of Run type is not allowed to be pickled." in str(e):
                 raise SerializationError(
-                    """You are trying to reference to a sagemaker.experiments.run.Run instance from within the function
-                       or passing it as a function argument.
-                       Instantiate a Run in the function or use load_run instead."""
+                    """You are trying to pass a sagemaker.experiments.run.Run object to a remote function
+                       or are trying to access a global sagemaker.experiments.run.Run object from within the function.
+                       This is not supported. You must use `load_run` to load an existing Run in the remote function
+                       or instantiate a new Run in the function."""
                 ) from e
 
             raise SerializationError(

src/sagemaker/remote_function/job.py

@@ -25,6 +25,8 @@
     REMOTE_FUNCTION_ENVIRONMENT_VARIABLES,
     REMOTE_FUNCTION_IMAGE_URI,
     REMOTE_FUNCTION_DEPENDENCIES,
+    REMOTE_FUNCTION_PRE_EXECUTION_COMMANDS,
+    REMOTE_FUNCTION_PRE_EXECUTION_SCRIPT,
     REMOTE_FUNCTION_INCLUDE_LOCAL_WORKDIR,
     REMOTE_FUNCTION_INSTANCE_TYPE,
     REMOTE_FUNCTION_JOB_CONDA_ENV,
@@ -36,7 +38,6 @@
     REMOTE_FUNCTION_VPC_CONFIG_SUBNETS,
     REMOTE_FUNCTION_VPC_CONFIG_SECURITY_GROUP_IDS,
     REMOTE_FUNCTION_ENABLE_INTER_CONTAINER_TRAFFIC_ENCRYPTION,
-    REMOTE_FUNCTION_ENABLE_NETWORK_ISOLATION,
 )
 from sagemaker.experiments._run_context import _RunContext
 from sagemaker.experiments.run import Run
@@ -53,6 +54,7 @@
 # runtime script names
 BOOTSTRAP_SCRIPT_NAME = "bootstrap_runtime_environment.py"
 ENTRYPOINT_SCRIPT_NAME = "job_driver.sh"
+PRE_EXECUTION_SCRIPT_NAME = "pre_exec.sh"
 RUNTIME_MANAGER_SCRIPT_NAME = "runtime_environment_manager.py"
 
 # training channel names
@@ -121,6 +123,8 @@ def __init__(
         self,
         *,
         dependencies: str = None,
+        pre_execution_commands: List[str] = None,
+        pre_execution_script: str = None,
         environment_variables: Dict[str, str] = None,
         image_uri: str = None,
         include_local_workdir: bool = None,
@@ -141,7 +145,6 @@ def __init__(
         volume_kms_key: str = None,
         volume_size: int = 30,
         encrypt_inter_container_traffic: bool = None,
-        enable_network_isolation: bool = None,
     ):
 
         self.sagemaker_session = sagemaker_session or Session()
@@ -171,6 +174,24 @@ def __init__(
             config_path=REMOTE_FUNCTION_DEPENDENCIES,
             sagemaker_session=self.sagemaker_session,
         )
+
+        self.pre_execution_commands = resolve_value_from_config(
+            direct_input=pre_execution_commands,
+            config_path=REMOTE_FUNCTION_PRE_EXECUTION_COMMANDS,
+            sagemaker_session=self.sagemaker_session,
+        )
+
+        self.pre_execution_script = resolve_value_from_config(
+            direct_input=pre_execution_script,
+            config_path=REMOTE_FUNCTION_PRE_EXECUTION_SCRIPT,
+            sagemaker_session=self.sagemaker_session,
+        )
+
+        if self.pre_execution_commands is not None and self.pre_execution_script is not None:
+            raise ValueError(
+                "Only one of pre_execution_commands or pre_execution_script can be specified!"
+            )
+
         self.include_local_workdir = resolve_value_from_config(
             direct_input=include_local_workdir,
             config_path=REMOTE_FUNCTION_INCLUDE_LOCAL_WORKDIR,
@@ -202,12 +223,7 @@ def __init__(
             default_value=False,
             sagemaker_session=self.sagemaker_session,
         )
-        self.enable_network_isolation = resolve_value_from_config(
-            direct_input=enable_network_isolation,
-            config_path=REMOTE_FUNCTION_ENABLE_NETWORK_ISOLATION,
-            default_value=False,
-            sagemaker_session=self.sagemaker_session,
-        )
+        self.enable_network_isolation = False
 
         _role = resolve_value_from_config(
             direct_input=role,
@@ -326,6 +342,8 @@ def start(job_settings: _JobSettings, func, func_args, func_kwargs, run_info=Non
         user_dependencies_s3uri = _prepare_and_upload_dependencies(
             local_dependencies_path=dependencies_list_path,
             include_local_workdir=job_settings.include_local_workdir,
+            pre_execution_commands=job_settings.pre_execution_commands,
+            pre_execution_script_local_path=job_settings.pre_execution_script,
             s3_base_uri=s3_base_uri,
             s3_kms_key=job_settings.s3_kms_key,
             sagemaker_session=job_settings.sagemaker_session,
@@ -540,13 +558,20 @@ def _prepare_and_upload_runtime_scripts(
 def _prepare_and_upload_dependencies(
     local_dependencies_path: str,
     include_local_workdir: bool,
+    pre_execution_commands: List[str],
+    pre_execution_script_local_path: str,
     s3_base_uri: str,
     s3_kms_key: str,
     sagemaker_session: Session,
 ) -> str:
     """Upload the job dependencies to S3 if present"""
 
-    if not local_dependencies_path and not include_local_workdir:
+    if not (
+        local_dependencies_path
+        or include_local_workdir
+        or pre_execution_commands
+        or pre_execution_script_local_path
+    ):
         return None
 
     with _tmpdir() as tmp_workspace:
@@ -570,6 +595,25 @@ def _prepare_and_upload_dependencies(
                 "Copied dependencies file at '%s' to '%s'", local_dependencies_path, dst_path
             )
 
+        if pre_execution_commands or pre_execution_script_local_path:
+            if not os.path.isdir(tmp_workspace):
+                os.mkdir(tmp_workspace)
+            pre_execution_script = os.path.join(tmp_workspace, PRE_EXECUTION_SCRIPT_NAME)
+            if pre_execution_commands:
+                with open(pre_execution_script, "w") as target_script:
+                    commands = [cmd + "\n" for cmd in pre_execution_commands]
+                    target_script.writelines(commands)
+                    logger.info(
+                        "Generated pre-execution script from commands to '%s'", pre_execution_script
+                    )
+            else:
+                shutil.copy(pre_execution_script_local_path, pre_execution_script)
+                logger.info(
+                    "Copied pre-execution commands from script at '%s' to '%s'",
+                    pre_execution_script_local_path,
+                    pre_execution_script,
+                )
+
         workspace_archive_path = os.path.join(tmp_workspace, "workspace")
         workspace_archive_path = shutil.make_archive(workspace_archive_path, "zip", tmp_workspace)
         logger.info("Successfully created workdir archive at '%s'", workspace_archive_path)

src/sagemaker/remote_function/runtime_environment/bootstrap_runtime_environment.py

@@ -35,6 +35,7 @@
 FAILURE_REASON_PATH = "/opt/ml/output/failure"
 SAGEMAKER_WHL_FILE_NAME = "sagemaker-2.132.1.dev0-py2.py3-none-any.whl"
 SAGEMAKER_WHL_CHANNEL = "sagemaker_whl_file"
+PRE_EXECUTION_SCRIPT_NAME = "pre_exec.sh"
 
 
 logger = get_logger()
@@ -54,8 +55,6 @@ def main():
 
         RuntimeEnvironmentManager()._validate_python_version(client_python_version, conda_env)
 
-        _execute_pre_exec_cmds()
-
         _bootstrap_runtime_environment(client_python_version, conda_env)
 
         exit_code = SUCCESS_EXIT_CODE
@@ -67,12 +66,6 @@ def main():
         sys.exit(exit_code)
 
 
-def _execute_pre_exec_cmds():
-    """Execute pre-flight commands before invkoing remote function"""
-    # TODO: complete me
-    pass  # pylint: disable=W0107
-
-
 def _bootstrap_runtime_environment(
     client_python_version: str,
     conda_env: str = None,
@@ -104,6 +97,10 @@ def _bootstrap_runtime_environment(
     shutil.unpack_archive(filename=workspace_archive_path, extract_dir=workspace_unpack_dir)
     logger.info("Successfully unpacked workspace archive at '%s'.", workspace_unpack_dir)
 
+    # Handle pre-execution commands
+    path_to_pre_exec_script = os.path.join(workspace_unpack_dir, PRE_EXECUTION_SCRIPT_NAME)
+    RuntimeEnvironmentManager().run_pre_exec_script(pre_exec_script_path=path_to_pre_exec_script)
+
     # Handle dependencies file.
     dependencies_file = None
     for file in os.listdir(workspace_unpack_dir):