sagemaker-core 1.0.47__py3-none-any.whl → 2.1.1__py3-none-any.whl

sagemaker/core/remote_function/job.py

@@ -0,0 +1,2140 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""Helper classes that interact with SageMaker Training service."""
+from __future__ import absolute_import
+
+import dataclasses
+import json
+import os
+import re
+import secrets
+import shutil
+import sys
+import time
+from io import BytesIO
+from typing import Callable, Dict, List, Optional, Tuple, Union, TYPE_CHECKING
+from urllib.parse import urlparse
+
+import botocore
+from botocore.exceptions import ClientError
+
+from sagemaker.core.config.config_schema import (
+    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,
+    REMOTE_FUNCTION_ROLE_ARN,
+    REMOTE_FUNCTION_S3_ROOT_URI,
+    REMOTE_FUNCTION_S3_KMS_KEY_ID,
+    REMOTE_FUNCTION_VOLUME_KMS_KEY_ID,
+    REMOTE_FUNCTION_TAGS,
+    REMOTE_FUNCTION_VPC_CONFIG_SUBNETS,
+    REMOTE_FUNCTION_VPC_CONFIG_SECURITY_GROUP_IDS,
+    REMOTE_FUNCTION_ENABLE_INTER_CONTAINER_TRAFFIC_ENCRYPTION,
+)
+from sagemaker.core.experiments._run_context import _RunContext
+from sagemaker.core.experiments.run import Run
+from sagemaker.core.image_uris import get_base_python_image_uri
+from sagemaker.core import image_uris
+from sagemaker.core.remote_function.checkpoint_location import CheckpointLocation
+from sagemaker.core.helper.session_helper import get_execution_role, expand_role, Session
+from sagemaker.core.common_utils import (
+    name_from_base,
+    _tmpdir,
+    resolve_value_from_config,
+    format_tags,
+    Tags,
+)
+from sagemaker.core.s3 import s3_path_join, S3Uploader
+
+from sagemaker.core.remote_function.core.stored_function import StoredFunction, _SerializedData
+from sagemaker.core.remote_function.core.pipeline_variables import Context
+
+from sagemaker.core.remote_function.runtime_environment.runtime_environment_manager import (
+    RuntimeEnvironmentManager,
+    _DependencySettings,
+)
+from sagemaker.core.remote_function import logging_config
+from sagemaker.core.remote_function.spark_config import SparkConfig
+from sagemaker.core.remote_function.custom_file_filter import (
+    CustomFileFilter,
+    copy_workdir,
+    resolve_custom_file_filter_from_config_file,
+)
+
+# Lazy import to avoid circular dependency - DelayedReturn is in MLOps which depends on Core
+# from sagemaker.mlops.workflow.function_step import DelayedReturn
+from sagemaker.core.workflow.step_outputs import get_step
+from sagemaker.core import exceptions
+from sagemaker.core import network as vpc_utils
+
+from sagemaker.core import logs as sagemaker_logs
+
+from sagemaker.core.common_utils import (
+    _wait_until,
+    secondary_training_status_changed,
+    secondary_training_status_message,
+)
+from sagemaker.core.config.config_utils import _append_sagemaker_config_tags
+
+if TYPE_CHECKING:
+    from sagemaker.core.helper.pipeline_variable import PipelineVariable
+
+# runtime script names
+BOOTSTRAP_SCRIPT_NAME = "bootstrap_runtime_environment.py"
+MPI_UTILS_SCRIPT_NAME = "mpi_utils_remote.py"
+ENTRYPOINT_SCRIPT_NAME = "job_driver.sh"
+PRE_EXECUTION_SCRIPT_NAME = "pre_exec.sh"
+RUNTIME_MANAGER_SCRIPT_NAME = "runtime_environment_manager.py"
+SPARK_APP_SCRIPT_NAME = "spark_app.py"
+
+# 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"
+SCRIPT_AND_DEPENDENCIES_CHANNEL_NAME = "pre_exec_script_and_dependencies"
+
+# Spark config channel and file name
+SPARK_CONF_CHANNEL_NAME = "conf"
+SPARK_CONF_FILE_NAME = "configuration.json"
+
+# Spark submitted files workspace names on S3
+SPARK_SUBMIT_JARS_WORKSPACE = "sm_rf_spark_jars"
+SPARK_SUBMIT_PY_FILES_WORKSPACE = "sm_rf_spark_py_files"
+SPARK_SUBMIT_FILES_WORKSPACE = "sm_rf_spark_data_files"
+SPARK_CONF_WORKSPACE = "sm_rf_spark_conf"
+
+# default spark version
+DEFAULT_SPARK_VERSION = "3.3"
+DEFAULT_SPARK_CONTAINER_VERSION = "v1"
+
+SPARK_NAME = "spark"
+
+# run context dictionary keys
+KEY_EXPERIMENT_NAME = "experiment_name"
+KEY_RUN_NAME = "run_name"
+
+JOBS_CONTAINER_ENTRYPOINT = [
+    "/bin/bash",
+    f"/opt/ml/input/data/{RUNTIME_SCRIPTS_CHANNEL_NAME}/{ENTRYPOINT_SCRIPT_NAME}",
+]
+
+SPARK_APP_SCRIPT_PATH = f"/opt/ml/input/data/{RUNTIME_SCRIPTS_CHANNEL_NAME}/{SPARK_APP_SCRIPT_NAME}"
+
+ENTRYPOINT_SCRIPT = f"""
+#!/bin/bash
+
+# Entry point for bootstrapping runtime environment and invoking remote function
+
+set -eu
+
+PERSISTENT_CACHE_DIR=${{SAGEMAKER_MANAGED_WARMPOOL_CACHE_DIRECTORY:-/opt/ml/cache}}
+export CONDA_PKGS_DIRS=${{PERSISTENT_CACHE_DIR}}/sm_remotefunction_user_dependencies_cache/conda/pkgs
+printf "INFO: CONDA_PKGS_DIRS is set to '$CONDA_PKGS_DIRS'\\n"
+export PIP_CACHE_DIR=${{PERSISTENT_CACHE_DIR}}/sm_remotefunction_user_dependencies_cache/pip
+printf "INFO: PIP_CACHE_DIR is set to '$PIP_CACHE_DIR'\\n"
+
+printf "INFO: /opt/ml/input/config/resourceconfig.json:\\n"
+cat /opt/ml/input/config/resourceconfig.json
+
+printf "INFO: Bootstraping runtime environment.\\n"
+python /opt/ml/input/data/{RUNTIME_SCRIPTS_CHANNEL_NAME}/{BOOTSTRAP_SCRIPT_NAME} "$@"
+source /opt/ml/input/sm_training.env
+
+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)
+
+    if which mamba >/dev/null; then
+        conda_exe="mamba"
+    else
+        conda_exe="conda"
+    fi
+
+    printf "INFO: Invoking remote function inside conda environment: $conda_env.\\n"
+    printf "INFO: $conda_exe run -n $conda_env python -m sagemaker.train.remote_function.invoke_function \\n"
+    $conda_exe run -n $conda_env python -m sagemaker.train.remote_function.invoke_function "$@"
+else
+    printf "INFO: No conda env provided. Invoking remote function\\n"
+    printf "INFO: python -m sagemaker.train.remote_function.invoke_function \\n"
+    python -m sagemaker.train.remote_function.invoke_function "$@"
+fi
+"""
+
+ENTRYPOINT_MPIRUN_SCRIPT = f"""
+#!/bin/bash
+
+# Entry point for bootstrapping runtime environment and invoking remote function with mpirun
+
+set -eu
+
+PERSISTENT_CACHE_DIR=${{SAGEMAKER_MANAGED_WARMPOOL_CACHE_DIRECTORY:-/opt/ml/cache}}
+export CONDA_PKGS_DIRS=${{PERSISTENT_CACHE_DIR}}/sm_remotefunction_user_dependencies_cache/conda/pkgs
+printf "INFO: CONDA_PKGS_DIRS is set to '$CONDA_PKGS_DIRS'\\n"
+export PIP_CACHE_DIR=${{PERSISTENT_CACHE_DIR}}/sm_remotefunction_user_dependencies_cache/pip
+printf "INFO: PIP_CACHE_DIR is set to '$PIP_CACHE_DIR'\\n"
+
+printf "INFO: /opt/ml/input/config/resourceconfig.json:\\n"
+cat /opt/ml/input/config/resourceconfig.json
+
+printf "INFO: Bootstraping runtime environment.\\n"
+python /opt/ml/input/data/{RUNTIME_SCRIPTS_CHANNEL_NAME}/{BOOTSTRAP_SCRIPT_NAME} "$@"
+source /opt/ml/input/sm_training.env
+
+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)
+
+    if which mamba >/dev/null; then
+        conda_exe="mamba"
+    else
+        conda_exe="conda"
+    fi
+
+    if [ "$SM_CURRENT_HOST" = "$SM_MASTER_ADDR" ]; then
+        python /opt/ml/input/data/{RUNTIME_SCRIPTS_CHANNEL_NAME}/{MPI_UTILS_SCRIPT_NAME}
+
+        printf "INFO: Invoking remote function with mpirun inside conda environment: $conda_env.\\n"
+        printf "INFO: $conda_exe run -n $conda_env mpirun --host $SM_HOSTS_LIST -np $SM_NPROC_PER_NODE \
+        --allow-run-as-root --display-map --tag-output -mca btl_tcp_if_include $SM_NETWORK_INTERFACE_NAME \
+        -mca plm_rsh_no_tree_spawn 1 -mca pml ob1 -mca btl ^openib -mca orte_abort_on_non_zero_status 1 \
+        -mca btl_vader_single_copy_mechanism none -mca plm_rsh_num_concurrent $SM_HOST_COUNT \
+        -x NCCL_SOCKET_IFNAME=$SM_NETWORK_INTERFACE_NAME -x LD_LIBRARY_PATH -x PATH \
+
+        python -m mpi4py -m sagemaker.train.remote_function.invoke_function \\n"
+        $conda_exe run -n $conda_env mpirun --host $SM_HOSTS_LIST -np $SM_NPROC_PER_NODE \
+        --allow-run-as-root --display-map --tag-output -mca btl_tcp_if_include $SM_NETWORK_INTERFACE_NAME \
+        -mca plm_rsh_no_tree_spawn 1 -mca pml ob1 -mca btl ^openib -mca orte_abort_on_non_zero_status 1 \
+        -mca btl_vader_single_copy_mechanism none -mca plm_rsh_num_concurrent $SM_HOST_COUNT \
+        -x NCCL_SOCKET_IFNAME=$SM_NETWORK_INTERFACE_NAME -x LD_LIBRARY_PATH -x PATH \
+        $SM_FI_PROVIDER $SM_NCCL_PROTO $SM_FI_EFA_USE_DEVICE_RDMA \
+        python -m mpi4py -m sagemaker.train.remote_function.invoke_function "$@"
+
+        python /opt/ml/input/data/{RUNTIME_SCRIPTS_CHANNEL_NAME}/{MPI_UTILS_SCRIPT_NAME} --job_ended 1
+    else
+        printf "INFO: This is the instance $SM_CURRENT_HOST. mpirun command terminated\\n"
+        python /opt/ml/input/data/{RUNTIME_SCRIPTS_CHANNEL_NAME}/{MPI_UTILS_SCRIPT_NAME}
+    fi
+else
+    if [ "$SM_CURRENT_HOST" = "$SM_MASTER_ADDR" ]; then
+        python /opt/ml/input/data/{RUNTIME_SCRIPTS_CHANNEL_NAME}/{MPI_UTILS_SCRIPT_NAME}
+
+        printf "INFO: No conda env provided. Invoking remote function with mpirun\\n"
+        printf "INFO: mpirun --host $SM_HOSTS_LIST -np $SM_NPROC_PER_NODE \
+        --allow-run-as-root --display-map --tag-output -mca btl_tcp_if_include $SM_NETWORK_INTERFACE_NAME \
+        -mca plm_rsh_no_tree_spawn 1 -mca pml ob1 -mca btl ^openib -mca orte_abort_on_non_zero_status 1 \
+        -mca btl_vader_single_copy_mechanism none -mca plm_rsh_num_concurrent $SM_HOST_COUNT \
+        -x NCCL_SOCKET_IFNAME=$SM_NETWORK_INTERFACE_NAME -x LD_LIBRARY_PATH -x PATH \
+        $SM_FI_PROVIDER $SM_NCCL_PROTO $SM_FI_EFA_USE_DEVICE_RDMA \
+        python -m mpi4py -m sagemaker.train.remote_function.invoke_function \\n"
+
+        mpirun --host $SM_HOSTS_LIST -np $SM_NPROC_PER_NODE \
+        --allow-run-as-root --display-map --tag-output -mca btl_tcp_if_include $SM_NETWORK_INTERFACE_NAME \
+        -mca plm_rsh_no_tree_spawn 1 -mca pml ob1 -mca btl ^openib -mca orte_abort_on_non_zero_status 1 \
+        -mca btl_vader_single_copy_mechanism none -mca plm_rsh_num_concurrent $SM_HOST_COUNT \
+        -x NCCL_SOCKET_IFNAME=$SM_NETWORK_INTERFACE_NAME -x LD_LIBRARY_PATH -x PATH \
+        $SM_FI_PROVIDER $SM_NCCL_PROTO $SM_FI_EFA_USE_DEVICE_RDMA \
+        python -m mpi4py -m sagemaker.train.remote_function.invoke_function "$@"
+
+        python /opt/ml/input/data/{RUNTIME_SCRIPTS_CHANNEL_NAME}/{MPI_UTILS_SCRIPT_NAME} --job_ended 1
+    else
+        printf "INFO: This is the instance $SM_CURRENT_HOST.\\n"
+        python /opt/ml/input/data/{RUNTIME_SCRIPTS_CHANNEL_NAME}/{MPI_UTILS_SCRIPT_NAME}
+    fi
+fi
+"""
+
+ENTRYPOINT_TORCHRUN_SCRIPT = f"""
+#!/bin/bash
+
+# Entry point for bootstrapping runtime environment and invoking remote function with torchrun
+
+set -eu
+
+PERSISTENT_CACHE_DIR=${{SAGEMAKER_MANAGED_WARMPOOL_CACHE_DIRECTORY:-/opt/ml/cache}}
+export CONDA_PKGS_DIRS=${{PERSISTENT_CACHE_DIR}}/sm_remotefunction_user_dependencies_cache/conda/pkgs
+printf "INFO: CONDA_PKGS_DIRS is set to '$CONDA_PKGS_DIRS'\\n"
+export PIP_CACHE_DIR=${{PERSISTENT_CACHE_DIR}}/sm_remotefunction_user_dependencies_cache/pip
+printf "INFO: PIP_CACHE_DIR is set to '$PIP_CACHE_DIR'\\n"
+
+printf "INFO: /opt/ml/input/config/resourceconfig.json:\\n"
+cat /opt/ml/input/config/resourceconfig.json
+
+printf "INFO: Bootstraping runtime environment.\\n"
+python /opt/ml/input/data/{RUNTIME_SCRIPTS_CHANNEL_NAME}/{BOOTSTRAP_SCRIPT_NAME} "$@"
+source /opt/ml/input/sm_training.env
+
+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)
+
+    if which mamba >/dev/null; then
+        conda_exe="mamba"
+    else
+        conda_exe="conda"
+    fi
+
+    printf "INFO: Invoking remote function with torchrun inside conda environment: $conda_env.\\n"
+    printf "INFO: $conda_exe run -n $conda_env torchrun --nnodes $SM_HOST_COUNT --nproc_per_node $SM_NPROC_PER_NODE \
+    --master_addr $SM_MASTER_ADDR --master_port $SM_MASTER_PORT --node_rank $SM_CURRENT_HOST_RANK \
+    -m sagemaker.train.remote_function.invoke_function \\n"
+
+    $conda_exe run -n $conda_env torchrun --nnodes $SM_HOST_COUNT --nproc_per_node $SM_NPROC_PER_NODE \
+    --master_addr $SM_MASTER_ADDR --master_port $SM_MASTER_PORT --node_rank $SM_CURRENT_HOST_RANK \
+    -m sagemaker.train.remote_function.invoke_function "$@"
+else
+    printf "INFO: No conda env provided. Invoking remote function with torchrun\\n"
+    printf "INFO: torchrun --nnodes $SM_HOST_COUNT --nproc_per_node $SM_NPROC_PER_NODE --master_addr $SM_MASTER_ADDR \
+    --master_port $SM_MASTER_PORT --node_rank $SM_CURRENT_HOST_RANK -m sagemaker.train.remote_function.invoke_function \\n"
+
+    torchrun --nnodes $SM_HOST_COUNT --nproc_per_node $SM_NPROC_PER_NODE --master_addr $SM_MASTER_ADDR \
+    --master_port $SM_MASTER_PORT --node_rank $SM_CURRENT_HOST_RANK -m sagemaker.train.remote_function.invoke_function "$@"
+fi
+"""
+
+SPARK_ENTRYPOINT_SCRIPT = f"""
+#!/bin/bash
+
+# Entry point for bootstrapping runtime environment and invoking remote function for Spark
+
+set -eu
+
+printf "INFO: Bootstraping Spark runtime environment.\\n"
+
+python3 /opt/ml/input/data/{RUNTIME_SCRIPTS_CHANNEL_NAME}/{BOOTSTRAP_SCRIPT_NAME} "$@"
+
+# Spark Container entry point script to initiate the spark application
+smspark-submit "$@"
+"""
+
+_STATUS_CODE_TABLE = {
+    "COMPLETED": "Completed",
+    "INPROGRESS": "InProgress",
+    "IN_PROGRESS": "InProgress",
+    "FAILED": "Failed",
+    "STOPPED": "Stopped",
+    "STOPPING": "Stopping",
+    "STARTING": "Starting",
+    "PENDING": "Pending",
+}
+
+logger = logging_config.get_logger()
+
+
+class LogState(object):
+    """Placeholder docstring"""
+
+    STARTING = 1
+    WAIT_IN_PROGRESS = 2
+    TAILING = 3
+    JOB_COMPLETE = 4
+    COMPLETE = 5
+
+
+class _JobSettings:
+    """Helper class that processes the job settings.
+
+    It validates the job settings and provides default values if necessary.
+    """
+
+    def __init__(
+        self,
+        *,
+        dependencies: str = None,
+        pre_execution_commands: List[str] = None,
+        pre_execution_script: str = None,
+        environment_variables: Dict[str, Union[str, "PipelineVariable"]] = None,
+        image_uri: Union[str, "PipelineVariable"] = None,
+        include_local_workdir: bool = None,
+        custom_file_filter: Optional[Union[Callable[[str, List], List], CustomFileFilter]] = None,
+        instance_count: Union[int, "PipelineVariable"] = 1,
+        instance_type: Union[str, "PipelineVariable"] = None,
+        job_conda_env: Union[str, "PipelineVariable"] = None,
+        job_name_prefix: str = None,
+        keep_alive_period_in_seconds: Union[int, "PipelineVariable"] = 0,
+        max_retry_attempts: Union[int, "PipelineVariable"] = 1,
+        max_runtime_in_seconds: Union[int, "PipelineVariable"] = 24 * 60 * 60,
+        role: str = None,
+        s3_kms_key: Union[str, "PipelineVariable"] = None,
+        s3_root_uri: str = None,
+        sagemaker_session: Session = None,
+        security_group_ids: List[Union[str, "PipelineVariable"]] = None,
+        subnets: List[Union[str, "PipelineVariable"]] = None,
+        tags: Optional[Tags] = None,
+        volume_kms_key: Union[str, "PipelineVariable"] = None,
+        volume_size: Union[int, "PipelineVariable"] = 30,
+        encrypt_inter_container_traffic: Union[bool, "PipelineVariable"] = None,
+        spark_config: SparkConfig = None,
+        use_spot_instances=False,
+        max_wait_time_in_seconds=None,
+        disable_output_compression: bool = False,
+        use_torchrun: bool = False,
+        use_mpirun: bool = False,
+        nproc_per_node: Optional[int] = None,
+    ):
+        """Initialize a _JobSettings instance which configures the remote job.
+
+        Args:
+            dependencies (str): Either the path to a dependencies file or the reserved keyword
+              ``auto_capture``. Defaults to ``None``.
+              If ``dependencies`` is provided, the value must be one of the following:
+
+              * A path to a conda environment.yml file. The following conditions apply.
+
+                * If job_conda_env is set, then the conda environment is updated by installing
+                  dependencies from the yaml file and the function is invoked within that
+                  conda environment. For this to succeed, the specified conda environment must
+                  already exist in the image.
+                * If the environment variable ``SAGEMAKER_JOB_CONDA_ENV`` is set in the image,
+                  then the conda environment is updated by installing dependencies from the
+                  yaml file and the function is invoked within that conda environment. For
+                  this to succeed, the conda environment name must already be set in
+                  ``SAGEMAKER_JOB_CONDA_ENV``, and ``SAGEMAKER_JOB_CONDA_ENV`` must already
+                  exist in the image.
+                * If none of the previous conditions are met, a new conda environment named
+                  ``sagemaker-runtime-env`` is created and the function annotated with the remote
+                  decorator is invoked in that conda environment.
+
+              * A path to a requirements.txt file. The following conditions apply.
+
+                * If ``job_conda_env`` is set in the remote decorator, dependencies are installed
+                  within that conda environment and the function annotated with the remote decorator
+                  is invoked in the same conda environment. For this to succeed, the specified
+                  conda environment must already exist in the image.
+                * If an environment variable ``SAGEMAKER_JOB_CONDA_ENV`` is set in the image,
+                  dependencies are installed within that conda environment and the function
+                  annotated with the remote decorator is invoked in the same. For this to succeed,
+                  the conda environment name must already be set in ``SAGEMAKER_JOB_CONDA_ENV``, and
+                  ``SAGEMAKER_JOB_CONDA_ENV`` must already exist in the image.
+                * If none of the above conditions are met, conda is not used. Dependencies are
+                  installed at the system level, without any virtual environment, and the function
+                  annotated with the remote decorator is invoked using the Python runtime available
+                  in the system path.
+
+              * The parameter dependencies is set to ``auto_capture``. SageMaker will automatically
+                generate an env_snapshot.yml corresponding to the current active conda environment’s
+                snapshot. You do not need to provide a dependencies file. The following conditions
+                apply:
+
+                * You must run the remote function within an active conda environment.
+                * When installing the dependencies on the training job, the same conditions
+                  as when dependencies is set to a path to a conda environment file apply.
+                  These conditions are as follows:
+
+                  * If job_conda_env is set, then the conda environment is updated by installing
+                    dependencies from the yaml file and the function is invoked within that
+                    conda environment. For this to succeed, the specified conda environment must
+                    already exist in the image.
+                  * If the environment variable ``SAGEMAKER_JOB_CONDA_ENV`` is set in the image,
+                    then the conda environment is updated by installing dependencies from the yaml
+                    file and the function is invoked within that conda environment. For this to
+                    succeed, the conda environment name must already be set in
+                    ``SAGEMAKER_JOB_CONDA_ENV``, and ``SAGEMAKER_JOB_CONDA_ENV`` must already exist
+                    in the image.
+                  * If none of the previous conditions are met, a new conda environment with name
+                    ``sagemaker-runtime-env`` is created and the function annotated with the
+                    remote decorator is invoked in that conda environment.
+
+              * ``None``. SageMaker will assume that there are no dependencies to install while
+                executing the remote annotated function in the training job.
+
+            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[str, str] or dict[str, PipelineVariable]): The environment
+              variables used inside the decorator function. Defaults to ``None``.
+
+            image_uri (str, PipelineVariable): The universal resource identifier (URI) location of
+              a Docker image on Amazon Elastic Container Registry (ECR). Defaults to the following
+              based on where the SDK is running:
+
+                * For users who specify ``spark_config`` and want to run the function in a Spark
+                  application, the ``image_uri`` should be ``None``. A SageMaker Spark image will
+                  be used for training, otherwise a ``ValueError`` is thrown.
+                * For users on SageMaker Studio notebooks, the image used as the kernel image for
+                  the notebook is used.
+                * For other users, it is resolved to base python image with the same python version
+                  as the environment running the local code.
+
+              If no compatible image is found, a ValueError is thrown.
+
+            include_local_workdir (bool): A flag to indicate that the remote function should include
+              local directories. Set to ``True`` if the remote function code imports local modules
+              and methods that are not available via PyPI or conda. Default value is ``False``.
+
+            custom_file_filter (Callable[[str, List], List], CustomFileFilter): Either a function
+              that filters job dependencies to be uploaded to S3 or a ``CustomFileFilter`` object
+              that specifies the local directories and files to be included in the remote function.
+              If a callable is passed in, that function is passed to the ``ignore``  argument of
+              ``shutil.copytree``. Defaults to ``None``, which means only python
+              files are accepted and uploaded to S3.
+
+            instance_count (int, PipelineVariable): The number of instances to use. Defaults to 1.
+
+            instance_type (str, PipelineVariable): The Amazon Elastic Compute Cloud (EC2) instance
+              type to use to run the SageMaker job. e.g. ml.c4.xlarge. If not provided,
+              a ValueError is thrown.
+
+            job_conda_env (str, PipelineVariable): The name of the conda environment to activate
+              during job's runtime. Defaults to ``None``.
+
+            job_name_prefix (str, PipelineVariable): The prefix used to create the underlying
+              SageMaker job.
+
+            keep_alive_period_in_seconds (int, PipelineVariable): The duration in seconds to retain
+              and reuse provisioned infrastructure after the completion of a training job, also
+              known as SageMaker managed warm pools. The use of warm pools reduces the latency time
+              spent to provision new resources. The default value for
+              ``keep_alive_period_in_seconds`` is 0.
+              NOTE: Additional charges associated with warm pools may apply. Using this parameter
+              also activates a new persistent cache feature, which will further reduce job start up
+              latency than over using SageMaker managed warm pools alone by caching the package
+              source downloaded in the previous runs.
+
+            max_retry_attempts (int, PipelineVariable): The max number of times the job is retried
+              on ``InternalServerFailure`` Error from SageMaker service. Defaults to 1.
+
+            max_runtime_in_seconds (int, PipelineVariable): The upper limit in seconds to be used
+              for training. After this specified amount of time, SageMaker terminates the job
+              regardless of its current status. Defaults to 1 day or (86400 seconds).
+
+            role (str): The IAM role (either name or full ARN) used to run your SageMaker training
+              job. Defaults to:
+
+              * the SageMaker default IAM role if the SDK is running in SageMaker Notebooks or
+                SageMaker Studio Notebooks.
+              * if not above, a ValueError is thrown.
+
+            s3_kms_key (str): The key used to encrypt the input and output data.
+              Default to ``None``.
+
+            s3_root_uri (str): The root S3 folder to which the code archives and data are
+              uploaded to. Defaults to ``s3://<sagemaker-default-bucket>``.
+
+            sagemaker_session (sagemaker.core.helper.session.Session): The underlying SageMaker session to
+              which SageMaker service calls are delegated to (default: None). If not provided,
+              one is created using a default configuration chain.
+
+            security_group_ids (List[str, PipelineVariable]): A list of security group IDs.
+              Defaults to ``None`` and the training job is created without VPC config.
+
+            subnets (List[str, PipelineVariable]): A list of subnet IDs. Defaults to ``None``
+              and the job is created without VPC config.
+
+            tags (Optional[Tags]): Tags attached to the job. Defaults to ``None``
+                and the training job is created without tags.
+
+            volume_kms_key (str, PipelineVariable): An Amazon Key Management Service (KMS) key
+              used to encrypt an Amazon Elastic Block Storage (EBS) volume attached to the
+              training instance. Defaults to ``None``.
+
+            volume_size (int, PipelineVariable): The size in GB of the storage volume for storing
+              input and output data during training. Defaults to ``30``.
+
+            encrypt_inter_container_traffic (bool, PipelineVariable): A flag that specifies
+              whether traffic between training containers is encrypted for the training job.
+              Defaults to ``False``.
+
+            spark_config (SparkConfig): Configurations to the Spark application that runs on
+              Spark image. If ``spark_config`` is specified, a SageMaker Spark image uri
+              will be used for training. Note that ``image_uri`` can not be specified at the
+              same time otherwise a ``ValueError`` is thrown. Defaults to ``None``.
+
+            use_spot_instances (bool, PipelineVariable): Specifies whether to use SageMaker
+              Managed Spot instances for training. If enabled then the ``max_wait`` arg should
+              also be set. Defaults to ``False``.
+
+            max_wait_time_in_seconds (int): Timeout in seconds waiting for spot training job.
+              After this amount of time Amazon SageMaker will stop waiting for managed spot
+              training job to complete. Defaults to ``None``.
+
+            disable_output_compression (bool): Optional. When set to true, Model is uploaded to
+              Amazon S3 without compression after training finishes.
+
+            use_torchrun (bool): Specifies whether to use torchrun for distributed training.
+              Defaults to ``False``.
+
+            use_mpirun (bool): Specifies whether to use mpirun for distributed training.
+              Defaults to ``False``.
+
+            nproc_per_node (int): Optional. Specifies the number of processes per node for
+              distributed training. Defaults to ``None``.
+              This is defined automatically configured on the instance type.
+        """
+        self.sagemaker_session = sagemaker_session or Session()
+        self.environment_variables = resolve_value_from_config(
+            direct_input=environment_variables,
+            config_path=REMOTE_FUNCTION_ENVIRONMENT_VARIABLES,
+            default_value={},
+            sagemaker_session=self.sagemaker_session,
+        )
+        self.environment_variables.update(
+            {"AWS_DEFAULT_REGION": self.sagemaker_session.boto_region_name}
+        )
+
+        # The following will be overridden by the _Job.compile method.
+        # However, it needs to be kept here for feature store SDK.
+        # TODO: update the feature store SDK to set the HMAC key there.
+        self.environment_variables.update({"REMOTE_FUNCTION_SECRET_KEY": secrets.token_hex(32)})
+
+        if spark_config and image_uri:
+            raise ValueError("spark_config and image_uri cannot be specified at the same time!")
+
+        if spark_config and job_conda_env:
+            raise ValueError("Remote Spark jobs do not support job_conda_env.")
+
+        if spark_config and dependencies == "auto_capture":
+            raise ValueError(
+                "Remote Spark jobs do not support automatically capturing dependencies."
+            )
+
+        _image_uri = resolve_value_from_config(
+            direct_input=image_uri,
+            config_path=REMOTE_FUNCTION_IMAGE_URI,
+            sagemaker_session=self.sagemaker_session,
+        )
+
+        if spark_config:
+            self.image_uri = self._get_default_spark_image(self.sagemaker_session)
+            logger.info(
+                "Set the image uri as %s because value of spark_config is "
+                "indicating this is a remote spark job.",
+                self.image_uri,
+            )
+        elif _image_uri:
+            self.image_uri = _image_uri
+        else:
+            self.image_uri = self._get_default_image(self.sagemaker_session)
+
+        self.dependencies = resolve_value_from_config(
+            direct_input=dependencies,
+            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,
+            default_value=False,
+            sagemaker_session=self.sagemaker_session,
+        )
+
+        self.custom_file_filter = resolve_custom_file_filter_from_config_file(
+            custom_file_filter, self.sagemaker_session
+        )
+
+        self.instance_type = resolve_value_from_config(
+            direct_input=instance_type,
+            config_path=REMOTE_FUNCTION_INSTANCE_TYPE,
+            sagemaker_session=self.sagemaker_session,
+        )
+        if not self.instance_type:
+            raise ValueError("instance_type is a required parameter!")
+
+        self.instance_count = instance_count
+        self.volume_size = volume_size
+        self.max_runtime_in_seconds = max_runtime_in_seconds
+        self.max_retry_attempts = max_retry_attempts
+        self.keep_alive_period_in_seconds = keep_alive_period_in_seconds
+        self.spark_config = spark_config
+        self.use_spot_instances = use_spot_instances
+        self.max_wait_time_in_seconds = max_wait_time_in_seconds
+        self.job_conda_env = resolve_value_from_config(
+            direct_input=job_conda_env,
+            config_path=REMOTE_FUNCTION_JOB_CONDA_ENV,
+            sagemaker_session=self.sagemaker_session,
+        )
+        self.job_name_prefix = job_name_prefix
+        self.encrypt_inter_container_traffic = resolve_value_from_config(
+            direct_input=encrypt_inter_container_traffic,
+            config_path=REMOTE_FUNCTION_ENABLE_INTER_CONTAINER_TRAFFIC_ENCRYPTION,
+            default_value=False,
+            sagemaker_session=self.sagemaker_session,
+        )
+        self.enable_network_isolation = False
+
+        _role = resolve_value_from_config(
+            direct_input=role,
+            config_path=REMOTE_FUNCTION_ROLE_ARN,
+            sagemaker_session=self.sagemaker_session,
+        )
+        if _role:
+            self.role = expand_role(self.sagemaker_session.boto_session, _role)
+        else:
+            self.role = get_execution_role(self.sagemaker_session)
+
+        self.s3_root_uri = resolve_value_from_config(
+            direct_input=s3_root_uri,
+            config_path=REMOTE_FUNCTION_S3_ROOT_URI,
+            default_value=s3_path_join(
+                "s3://",
+                self.sagemaker_session.default_bucket(),
+                self.sagemaker_session.default_bucket_prefix,
+            ),
+            sagemaker_session=self.sagemaker_session,
+        )
+
+        self.s3_kms_key = resolve_value_from_config(
+            direct_input=s3_kms_key,
+            config_path=REMOTE_FUNCTION_S3_KMS_KEY_ID,
+            sagemaker_session=self.sagemaker_session,
+        )
+        self.volume_kms_key = resolve_value_from_config(
+            direct_input=volume_kms_key,
+            config_path=REMOTE_FUNCTION_VOLUME_KMS_KEY_ID,
+            sagemaker_session=self.sagemaker_session,
+        )
+
+        _subnets = resolve_value_from_config(
+            direct_input=subnets,
+            config_path=REMOTE_FUNCTION_VPC_CONFIG_SUBNETS,
+            sagemaker_session=self.sagemaker_session,
+        )
+        _security_group_ids = resolve_value_from_config(
+            direct_input=security_group_ids,
+            config_path=REMOTE_FUNCTION_VPC_CONFIG_SECURITY_GROUP_IDS,
+            sagemaker_session=self.sagemaker_session,
+        )
+        vpc_config = vpc_utils.to_dict(subnets=_subnets, security_group_ids=_security_group_ids)
+        self.vpc_config = vpc_utils.sanitize(vpc_config)
+
+        tags = format_tags(tags)
+        self.tags = _append_sagemaker_config_tags(
+            self.sagemaker_session, tags, REMOTE_FUNCTION_TAGS
+        )
+
+        self.disable_output_compression = disable_output_compression
+        self.use_torchrun = use_torchrun
+        self.use_mpirun = use_mpirun
+        self.nproc_per_node = nproc_per_node
+
+    @staticmethod
+    def _get_default_image(session):
+        """Return Studio notebook image, if in Studio env. Else, base python.
+
+        Args:
+            session (Session): Boto session.
+
+        Returns:
+            Default SageMaker base python image.
+        """
+
+        if (
+            "SAGEMAKER_INTERNAL_IMAGE_URI" in os.environ
+            and os.environ["SAGEMAKER_INTERNAL_IMAGE_URI"]
+        ):
+            return os.environ["SAGEMAKER_INTERNAL_IMAGE_URI"]
+
+        py_version = str(sys.version_info[0]) + str(sys.version_info[1])
+
+        if py_version not in ["310", "38"]:
+            raise ValueError(
+                "Default image is supported only for Python versions 3.8 and 3.10. If you "
+                "are using any other python version, you must provide a compatible image_uri."
+            )
+
+        region = session.boto_region_name
+        image_uri = get_base_python_image_uri(region=region, py_version=py_version)
+
+        return image_uri
+
+    @staticmethod
+    def _get_default_spark_image(session):
+        """Return the Spark image.
+
+        Args:
+            session (Session): Boto session.
+
+        Returns:
+            SageMaker Spark container image uri.
+        """
+
+        region = session.boto_region_name
+
+        py_version = str(sys.version_info[0]) + str(sys.version_info[1])
+
+        if py_version not in ["39"]:
+            raise ValueError(
+                "The SageMaker Spark image for remote job only supports Python version 3.9. "
+            )
+
+        image_uri = image_uris.retrieve(
+            framework=SPARK_NAME,
+            region=region,
+            version=DEFAULT_SPARK_VERSION,
+            instance_type=None,
+            py_version=f"py{py_version}",
+            container_version=DEFAULT_SPARK_CONTAINER_VERSION,
+        )
+
+        return image_uri
+
+
+class _Job:
+    """Helper class that interacts with the SageMaker training service."""
+
+    def __init__(self, job_name: str, s3_uri: str, sagemaker_session: Session, hmac_key: str):
+        """Initialize a _Job object.
+
+        Args:
+            job_name (str): The training job name.
+            s3_uri (str): The training job output S3 uri.
+            sagemaker_session (Session): SageMaker boto session.
+            hmac_key (str): Remote function secret key.
+        """
+        self.job_name = job_name
+        self.s3_uri = s3_uri
+        self.sagemaker_session = sagemaker_session
+        self.hmac_key = hmac_key
+        self._last_describe_response = None
+
+    @staticmethod
+    def from_describe_response(describe_training_job_response, sagemaker_session):
+        """Construct a _Job from a describe_training_job_response object.
+
+        Args:
+            describe_training_job_response (Dict): Describe training job response.
+            sagemaker_session (Session): SageMaker boto session.
+
+        Returns:
+            the _Job object.
+        """
+        job_name = describe_training_job_response["TrainingJobName"]
+        s3_uri = describe_training_job_response["OutputDataConfig"]["S3OutputPath"]
+        hmac_key = describe_training_job_response["Environment"]["REMOTE_FUNCTION_SECRET_KEY"]
+
+        job = _Job(job_name, s3_uri, sagemaker_session, hmac_key)
+        job._last_describe_response = describe_training_job_response
+        return job
+
+    @staticmethod
+    def start(job_settings: _JobSettings, func, func_args, func_kwargs, run_info=None):
+        """Start a training job.
+
+        Args:
+            job_settings (_JobSettings): the job settings.
+            func: the function to be executed.
+            func_args: the positional arguments to the function.
+            func_kwargs: the keyword arguments to the function
+
+        Returns:
+            the _Job object.
+        """
+        job_name = _Job._get_job_name(job_settings, func)
+        s3_base_uri = s3_path_join(job_settings.s3_root_uri, job_name)
+
+        training_job_request = _Job.compile(
+            job_settings=job_settings,
+            job_name=job_name,
+            s3_base_uri=s3_base_uri,
+            func=func,
+            func_args=func_args,
+            func_kwargs=func_kwargs,
+            run_info=run_info,
+        )
+
+        logger.info("Creating job: %s", job_name)
+
+        job_settings.sagemaker_session.sagemaker_client.create_training_job(**training_job_request)
+
+        return _Job(
+            job_name,
+            s3_base_uri,
+            job_settings.sagemaker_session,
+            training_job_request["Environment"]["REMOTE_FUNCTION_SECRET_KEY"],
+        )
+
+    @staticmethod
+    def compile(
+        job_settings: _JobSettings,
+        job_name: str,
+        s3_base_uri: str,
+        func: Callable,
+        func_args: tuple,
+        func_kwargs: dict,
+        run_info=None,
+        serialized_data: _SerializedData = None,
+    ) -> dict:
+        """Build the artifacts and generate the training job request."""
+        from sagemaker.core.workflow.properties import Properties
+        from sagemaker.core.workflow.parameters import Parameter
+        from sagemaker.core.workflow.functions import Join
+        from sagemaker.core.workflow.execution_variables import (
+            ExecutionVariables,
+            ExecutionVariable,
+        )
+        from sagemaker.core.workflow.utilities import load_step_compilation_context
+
+        step_compilation_context = load_step_compilation_context()
+
+        jobs_container_entrypoint = JOBS_CONTAINER_ENTRYPOINT[:]
+
+        # generate hmac key for integrity check
+        if step_compilation_context is None:
+            hmac_key = secrets.token_hex(32)
+        else:
+            hmac_key = step_compilation_context.function_step_secret_token
+
+        # serialize function and arguments
+        if step_compilation_context is None:
+            stored_function = StoredFunction(
+                sagemaker_session=job_settings.sagemaker_session,
+                s3_base_uri=s3_base_uri,
+                hmac_key=hmac_key,
+                s3_kms_key=job_settings.s3_kms_key,
+            )
+            stored_function.save(func, *func_args, **func_kwargs)
+        else:
+            stored_function = StoredFunction(
+                sagemaker_session=job_settings.sagemaker_session,
+                s3_base_uri=s3_base_uri,
+                hmac_key=hmac_key,
+                s3_kms_key=job_settings.s3_kms_key,
+                context=Context(
+                    step_name=step_compilation_context.step_name,
+                    func_step_s3_dir=step_compilation_context.pipeline_build_time,
+                ),
+            )
+
+            stored_function.save_pipeline_step_function(serialized_data)
+
+        stopping_condition = {
+            "MaxRuntimeInSeconds": job_settings.max_runtime_in_seconds,
+        }
+        if job_settings.max_wait_time_in_seconds is not None:
+            stopping_condition["MaxWaitTimeInSeconds"] = job_settings.max_wait_time_in_seconds
+
+        request_dict = dict(
+            TrainingJobName=job_name,
+            RoleArn=job_settings.role,
+            StoppingCondition=stopping_condition,
+            RetryStrategy={"MaximumRetryAttempts": job_settings.max_retry_attempts},
+        )
+
+        _update_job_request_with_checkpoint_config(func_args, func_kwargs, request_dict)
+
+        if job_settings.tags:
+            request_dict["Tags"] = job_settings.tags
+
+        # generate other build artifacts including workspace, requirements.txt
+        request_dict["InputDataConfig"] = _generate_input_data_config(
+            job_settings=job_settings, s3_base_uri=s3_base_uri
+        )
+
+        if step_compilation_context:
+            # Path format: base/step_name/build_timestamp/execution_id/results
+            # This matches the path construction in stored_function.py
+            s3_output_path = Join(
+                on="/",
+                values=[
+                    s3_base_uri,
+                    step_compilation_context.step_name,
+                    step_compilation_context.pipeline_build_time,
+                    ExecutionVariables.PIPELINE_EXECUTION_ID,
+                    "results",
+                ],
+            )
+            output_config = {"S3OutputPath": s3_output_path}
+        else:
+            output_config = {"S3OutputPath": s3_base_uri}
+        if job_settings.s3_kms_key is not None:
+            output_config["KmsKeyId"] = job_settings.s3_kms_key
+        if job_settings.disable_output_compression:
+            output_config["CompressionType"] = "NONE"
+        request_dict["OutputDataConfig"] = output_config
+
+        container_args = ["--s3_base_uri", s3_base_uri]
+        container_args.extend(["--region", job_settings.sagemaker_session.boto_region_name])
+        container_args.extend(
+            ["--client_python_version", RuntimeEnvironmentManager()._current_python_version()]
+        )
+        container_args.extend(
+            [
+                "--client_sagemaker_pysdk_version",
+                RuntimeEnvironmentManager()._current_sagemaker_pysdk_version(),
+            ]
+        )
+        container_args.extend(
+            [
+                "--dependency_settings",
+                _DependencySettings.from_dependency_file_path(
+                    job_settings.dependencies
+                ).to_string(),
+            ]
+        )
+        if job_settings.use_torchrun:
+            container_args.extend(["--distribution", "torchrun"])
+        elif job_settings.use_mpirun:
+            container_args.extend(["--distribution", "mpirun"])
+        if job_settings.nproc_per_node is not None and int(job_settings.nproc_per_node) > 0:
+            container_args.extend(["--user_nproc_per_node", str(job_settings.nproc_per_node)])
+        if job_settings.s3_kms_key:
+            container_args.extend(["--s3_kms_key", job_settings.s3_kms_key])
+
+        if job_settings.job_conda_env:
+            container_args.extend(["--job_conda_env", job_settings.job_conda_env])
+
+        if step_compilation_context:
+            # TODO: remove the duplicates in the list
+            container_args.extend(["--pipeline_step_name", step_compilation_context.step_name])
+            container_args.extend(
+                ["--pipeline_execution_id", ExecutionVariables.PIPELINE_EXECUTION_ID]
+            )
+            container_args.extend(
+                ["--func_step_s3_dir", step_compilation_context.pipeline_build_time]
+            )
+            container_args.extend(["--property_references"])
+            container_args.extend(
+                [
+                    ExecutionVariables.PIPELINE_EXECUTION_ID.expr["Get"],
+                    ExecutionVariables.PIPELINE_EXECUTION_ID.to_string(),
+                ]
+            )
+            for arg in func_args + tuple(func_kwargs.values()):
+                if isinstance(arg, (Parameter, ExecutionVariable, Properties)):
+                    container_args.extend([arg.expr["Get"], arg.to_string()])
+
+                # Lazy import to avoid circular dependency
+                try:
+                    from sagemaker.mlops.workflow.function_step import DelayedReturn
+
+                    if isinstance(arg, DelayedReturn):
+                        # The uri is a Properties object
+                        uri = get_step(arg)._properties.OutputDataConfig.S3OutputPath
+                        container_args.extend([uri.expr["Get"], uri.to_string()])
+                except ImportError:
+                    # MLOps not installed, skip DelayedReturn handling
+                    pass
+
+        if run_info is not None:
+            container_args.extend(["--run_in_context", json.dumps(dataclasses.asdict(run_info))])
+        elif _RunContext.get_current_run() is not None:
+            container_args.extend(
+                ["--run_in_context", _convert_run_to_json(_RunContext.get_current_run())]
+            )
+
+        algorithm_spec = dict(
+            TrainingImage=job_settings.image_uri,
+            TrainingInputMode="File",
+            ContainerEntrypoint=jobs_container_entrypoint,
+            ContainerArguments=container_args,
+        )
+
+        request_dict["AlgorithmSpecification"] = algorithm_spec
+
+        resource_config = dict(
+            VolumeSizeInGB=job_settings.volume_size,
+            InstanceCount=job_settings.instance_count,
+            InstanceType=job_settings.instance_type,
+        )
+        if job_settings.volume_kms_key is not None:
+            resource_config["VolumeKmsKeyId"] = job_settings.volume_kms_key
+        if job_settings.keep_alive_period_in_seconds is not None:
+            resource_config["KeepAlivePeriodInSeconds"] = job_settings.keep_alive_period_in_seconds
+
+        request_dict["ResourceConfig"] = resource_config
+
+        if job_settings.enable_network_isolation is not None:
+            request_dict["EnableNetworkIsolation"] = job_settings.enable_network_isolation
+
+        if job_settings.encrypt_inter_container_traffic is not None:
+            request_dict["EnableInterContainerTrafficEncryption"] = (
+                job_settings.encrypt_inter_container_traffic
+            )
+
+        if job_settings.vpc_config:
+            request_dict["VpcConfig"] = job_settings.vpc_config
+
+        request_dict["EnableManagedSpotTraining"] = job_settings.use_spot_instances
+
+        request_dict["Environment"] = job_settings.environment_variables
+        request_dict["Environment"].update({"REMOTE_FUNCTION_SECRET_KEY": hmac_key})
+
+        extended_request = _extend_spark_config_to_request(request_dict, job_settings, s3_base_uri)
+        extended_request = _extend_mpirun_to_request(extended_request, job_settings)
+        extended_request = _extend_torchrun_to_request(extended_request, job_settings)
+
+        return extended_request
+
+    def describe(self):
+        """Describe the underlying sagemaker training job.
+
+        Returns:
+            Dict: Describe training job response.
+        """
+        if self._last_describe_response is not None and self._last_describe_response[
+            "TrainingJobStatus"
+        ] in ["Completed", "Failed", "Stopped"]:
+            return self._last_describe_response
+
+        self._last_describe_response = (
+            self.sagemaker_session.sagemaker_client.describe_training_job(
+                TrainingJobName=self.job_name
+            )
+        )
+
+        return self._last_describe_response
+
+    def stop(self):
+        """Stop the underlying sagemaker training job."""
+        self.sagemaker_session.sagemaker_client.stop_training_job(TrainingJobName=self.job_name)
+
+    def wait(self, timeout: int = None):
+        """Wait for the underlying sagemaker job to finish and displays its logs .
+
+        This method blocks on the sagemaker job completing for up to the timeout value (if
+        specified). If timeout is ``None``, this method will block until the job is completed.
+
+        Args:
+            timeout (int): Timeout in seconds to wait until the job is completed. ``None`` by
+                default.
+
+        Returns: None
+        """
+
+        self._last_describe_response = _logs_for_job(
+            sagemaker_session=self.sagemaker_session,
+            job_name=self.job_name,
+            wait=True,
+            timeout=timeout,
+        )
+
+    @staticmethod
+    def _get_job_name(job_settings, func):
+        """Get the underlying SageMaker job name from job_name_prefix or func.
+
+        Args:
+            job_settings (_JobSettings): the job settings.
+            func: the function to be executed.
+
+        Returns:
+            str : the training job name.
+        """
+        from sagemaker.core.workflow.utilities import load_step_compilation_context
+
+        step_complication_context = load_step_compilation_context()
+
+        job_name_prefix = job_settings.job_name_prefix
+        if not job_name_prefix:
+            job_name_prefix = func.__name__
+            # remove all special characters in the beginning of function name
+            job_name_prefix = re.sub(r"^[^a-zA-Z0-9]+", "", job_name_prefix)
+            # convert all remaining special characters to '-'
+            job_name_prefix = re.sub(r"[^a-zA-Z0-9-]", "-", job_name_prefix)
+
+        if step_complication_context:
+            return job_name_prefix
+        return name_from_base(job_name_prefix)
+
+
+def _prepare_and_upload_runtime_scripts(
+    spark_config: SparkConfig,
+    s3_base_uri: str,
+    s3_kms_key: str,
+    sagemaker_session: Session,
+    use_torchrun: bool = False,
+    use_mpirun: bool = False,
+):
+    """Copy runtime scripts to a folder and upload to S3.
+
+    In case of remote function, s3_base_uri is s3_root_uri + function_name.
+    In case of pipeline, s3_base_uri is s3_root_uri + pipeline_name. The runtime scripts are
+    uploaded only once per pipeline.
+
+    Args:
+        spark_config (SparkConfig): remote Spark job configurations.
+
+        s3_base_uri (str): S3 location that the runtime scripts will be uploaded to.
+
+        s3_kms_key (str): kms key used to encrypt the files uploaded to S3.
+
+        sagemaker_session (str): SageMaker boto client session.
+
+        use_torchrun (bool): Whether to use torchrun or not.
+
+        use_mpirun (bool): Whether to use mpirun or not.
+
+        nproc_per_node (Optional[int]): Number of processes per node
+    """
+
+    from sagemaker.core.workflow.utilities import load_step_compilation_context
+
+    step_compilation_context = load_step_compilation_context()
+
+    if step_compilation_context and not step_compilation_context.upload_runtime_scripts:
+        return s3_path_join(s3_base_uri, RUNTIME_SCRIPTS_CHANNEL_NAME)
+
+    with _tmpdir() as bootstrap_scripts:
+
+        # write entrypoint script to tmpdir
+        entrypoint_script_path = os.path.join(bootstrap_scripts, ENTRYPOINT_SCRIPT_NAME)
+        entry_point_script = ENTRYPOINT_SCRIPT
+        if spark_config:
+            entry_point_script = SPARK_ENTRYPOINT_SCRIPT
+            spark_script_path = os.path.join(
+                os.path.dirname(__file__), "runtime_environment", SPARK_APP_SCRIPT_NAME
+            )
+            shutil.copy2(spark_script_path, bootstrap_scripts)
+
+        if use_torchrun:
+            entry_point_script = ENTRYPOINT_TORCHRUN_SCRIPT
+
+        if use_mpirun:
+            entry_point_script = ENTRYPOINT_MPIRUN_SCRIPT
+
+        with open(entrypoint_script_path, "w", newline="\n") as file:
+            file.writelines(entry_point_script)
+
+        bootstrap_script_path = os.path.join(
+            os.path.dirname(__file__), "runtime_environment", BOOTSTRAP_SCRIPT_NAME
+        )
+        mpi_utils_path = os.path.join(
+            os.path.dirname(__file__), "runtime_environment", MPI_UTILS_SCRIPT_NAME
+        )
+        runtime_manager_script_path = os.path.join(
+            os.path.dirname(__file__), "runtime_environment", RUNTIME_MANAGER_SCRIPT_NAME
+        )
+
+        # copy runtime scripts to tmpdir
+        shutil.copy2(bootstrap_script_path, bootstrap_scripts)
+        shutil.copy2(mpi_utils_path, bootstrap_scripts)
+        shutil.copy2(runtime_manager_script_path, bootstrap_scripts)
+
+        upload_path = S3Uploader.upload(
+            bootstrap_scripts,
+            s3_path_join(s3_base_uri, RUNTIME_SCRIPTS_CHANNEL_NAME),
+            s3_kms_key,
+            sagemaker_session,
+        )
+
+        if step_compilation_context:
+            step_compilation_context.upload_runtime_scripts = False
+        return upload_path
+
+
+def _generate_input_data_config(job_settings: _JobSettings, s3_base_uri: str):
+    """Generates input data config"""
+    from sagemaker.core.workflow.utilities import load_step_compilation_context
+
+    step_compilation_context = load_step_compilation_context()
+
+    bootstrap_scripts_s3uri = _prepare_and_upload_runtime_scripts(
+        spark_config=job_settings.spark_config,
+        s3_base_uri=s3_base_uri,
+        s3_kms_key=job_settings.s3_kms_key,
+        sagemaker_session=job_settings.sagemaker_session,
+        use_torchrun=job_settings.use_torchrun,
+        use_mpirun=job_settings.use_mpirun,
+    )
+
+    input_data_config = [
+        dict(
+            ChannelName=RUNTIME_SCRIPTS_CHANNEL_NAME,
+            DataSource={
+                "S3DataSource": {
+                    "S3Uri": bootstrap_scripts_s3uri,
+                    "S3DataType": "S3Prefix",
+                }
+            },
+        )
+    ]
+
+    local_dependencies_path = RuntimeEnvironmentManager().snapshot(job_settings.dependencies)
+
+    if step_compilation_context:
+        with _tmpdir() as tmp_dir:
+            script_and_dependencies_s3uri = _prepare_dependencies_and_pre_execution_scripts(
+                local_dependencies_path=local_dependencies_path,
+                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,
+                tmp_dir=tmp_dir,
+            )
+
+            if script_and_dependencies_s3uri:
+                input_data_config.append(
+                    dict(
+                        ChannelName=SCRIPT_AND_DEPENDENCIES_CHANNEL_NAME,
+                        DataSource={
+                            "S3DataSource": {
+                                "S3Uri": script_and_dependencies_s3uri,
+                                "S3DataType": "S3Prefix",
+                            }
+                        },
+                    )
+                )
+
+    user_workspace_s3uri = _prepare_and_upload_workspace(
+        local_dependencies_path=local_dependencies_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,
+        custom_file_filter=job_settings.custom_file_filter,
+    )
+
+    if user_workspace_s3uri:
+        input_data_config.append(
+            dict(
+                ChannelName=(
+                    REMOTE_FUNCTION_WORKSPACE
+                    if not step_compilation_context
+                    else step_compilation_context.pipeline_build_time
+                ),
+                DataSource={
+                    "S3DataSource": {
+                        "S3Uri": user_workspace_s3uri,
+                        "S3DataType": "S3Prefix",
+                    }
+                },
+            )
+        )
+
+    return input_data_config
+
+
+def _prepare_dependencies_and_pre_execution_scripts(
+    local_dependencies_path: str,
+    pre_execution_commands: List[str],
+    pre_execution_script_local_path: str,
+    s3_base_uri: str,
+    s3_kms_key: str,
+    sagemaker_session: Session,
+    tmp_dir: str,
+):
+    """Prepare pre-execution scripts and dependencies and upload them to s3.
+
+    If pre execution commands are provided, a new bash file will be created
+      with those commands in tmp directory.
+    If pre execution script is provided, it copies that file from local file path
+      to tmp directory.
+    If local dependencies file is provided, it copies that file from local file path
+      to tmp directory.
+    If under pipeline context, tmp directory with copied dependencies and scripts is
+      uploaded to S3.
+    """
+    from sagemaker.core.workflow.utilities import load_step_compilation_context
+
+    if not (local_dependencies_path or pre_execution_commands or pre_execution_script_local_path):
+        return None
+
+    if local_dependencies_path:
+        dst_path = shutil.copy2(local_dependencies_path, tmp_dir)
+        logger.info("Copied dependencies file at '%s' to '%s'", local_dependencies_path, dst_path)
+
+    if pre_execution_commands or pre_execution_script_local_path:
+        pre_execution_script = os.path.join(tmp_dir, 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.copy2(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,
+            )
+
+    step_compilation_context = load_step_compilation_context()
+    if step_compilation_context:
+        upload_path = S3Uploader.upload(
+            tmp_dir,
+            s3_path_join(
+                s3_base_uri,
+                step_compilation_context.step_name,
+                step_compilation_context.pipeline_build_time,
+                SCRIPT_AND_DEPENDENCIES_CHANNEL_NAME,
+            ),
+            s3_kms_key,
+            sagemaker_session,
+        )
+        logger.info(
+            "Successfully uploaded dependencies and pre execution scripts to '%s'", upload_path
+        )
+        return upload_path
+    return None
+
+
+def _prepare_and_upload_workspace(
+    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,
+    custom_file_filter: Optional[Union[Callable[[str, List], List], CustomFileFilter]] = None,
+) -> str:
+    """Prepare and upload the workspace to S3.
+
+    Under pipeline context, only workdir is packaged in the workspace folder and uploaded to s3.
+    Under remote function context, workdir along with pre execution scripts and dependencies
+      are packaged together into the workspace folder and uploaded to S3.
+    """
+    from sagemaker.core.workflow.utilities import load_step_compilation_context
+
+    step_compilation_context = load_step_compilation_context()
+
+    if not (
+        local_dependencies_path
+        or include_local_workdir
+        or pre_execution_commands
+        or pre_execution_script_local_path
+    ):
+        return None
+
+    func_step_s3_dir = None
+    if step_compilation_context:
+        func_step_s3_dir = step_compilation_context.pipeline_build_time
+        if not include_local_workdir:
+            return None
+        if not step_compilation_context.upload_workspace:
+            return s3_path_join(s3_base_uri, REMOTE_FUNCTION_WORKSPACE, func_step_s3_dir)
+
+    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_dir, JOB_REMOTE_FUNCTION_WORKSPACE)
+
+        if include_local_workdir:
+            copy_workdir(tmp_workspace, custom_file_filter)
+            logger.info("Copied user workspace to '%s'", tmp_workspace)
+
+        if not os.path.isdir(tmp_workspace):
+            # create the directory if no workdir_path was provided in the input.
+            os.mkdir(tmp_workspace)
+
+        if not step_compilation_context:
+            _prepare_dependencies_and_pre_execution_scripts(
+                local_dependencies_path=local_dependencies_path,
+                pre_execution_commands=pre_execution_commands,
+                pre_execution_script_local_path=pre_execution_script_local_path,
+                s3_base_uri=s3_base_uri,
+                s3_kms_key=s3_kms_key,
+                sagemaker_session=sagemaker_session,
+                tmp_dir=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(
+            workspace_archive_path,
+            s3_path_join(s3_base_uri, REMOTE_FUNCTION_WORKSPACE, func_step_s3_dir),
+            s3_kms_key,
+            sagemaker_session,
+        )
+        logger.info("Successfully uploaded workdir to '%s'", upload_path)
+        if step_compilation_context:
+            step_compilation_context.upload_workspace = False
+        return upload_path
+
+
+def _convert_run_to_json(run: Run) -> str:
+    """Convert current run into json string"""
+    run_info = _RunInfo(run.experiment_name, run.run_name)
+    return json.dumps(dataclasses.asdict(run_info))
+
+
+def _prepare_and_upload_spark_dependent_files(
+    spark_config: SparkConfig,
+    s3_base_uri: str,
+    s3_kms_key: str,
+    sagemaker_session: Session,
+) -> Tuple:
+    """Upload the Spark dependencies to S3 if present.
+
+    Args:
+        spark_config (SparkConfig): The remote Spark job configurations.
+        s3_base_uri (str): The S3 location that the Spark dependencies will be uploaded to.
+        s3_kms_key (str): The kms key used to encrypt the files uploaded to S3.
+        sagemaker_session (str): SageMaker boto client session.
+    """
+    if not spark_config:
+        return None, None, None, None
+
+    submit_jars_s3_paths = _upload_spark_submit_deps(
+        spark_config.submit_jars,
+        SPARK_SUBMIT_JARS_WORKSPACE,
+        s3_base_uri,
+        s3_kms_key,
+        sagemaker_session,
+    )
+    submit_py_files_s3_paths = _upload_spark_submit_deps(
+        spark_config.submit_py_files,
+        SPARK_SUBMIT_PY_FILES_WORKSPACE,
+        s3_base_uri,
+        s3_kms_key,
+        sagemaker_session,
+    )
+    submit_files_s3_path = _upload_spark_submit_deps(
+        spark_config.submit_files,
+        SPARK_SUBMIT_FILES_WORKSPACE,
+        s3_base_uri,
+        s3_kms_key,
+        sagemaker_session,
+    )
+    config_file_s3_uri = _upload_serialized_spark_configuration(
+        s3_base_uri, s3_kms_key, spark_config.configuration, sagemaker_session
+    )
+
+    return submit_jars_s3_paths, submit_py_files_s3_paths, submit_files_s3_path, config_file_s3_uri
+
+
+def _upload_spark_submit_deps(
+    submit_deps: List[str],
+    workspace_name: str,
+    s3_base_uri: str,
+    s3_kms_key: str,
+    sagemaker_session: Session,
+) -> str:
+    """Upload the Spark submit dependencies to S3.
+
+    Args:
+        submit_deps (List[str]): A list of path which points to the Spark dependency files.
+          The path can be either a local path or S3 uri. For example ``/local/deps.jar`` or
+          ``s3://<your-bucket>/deps.jar``.
+
+        workspace_name (str): workspace name for Spark dependency.
+        s3_base_uri (str): S3 location that the Spark dependencies will be uploaded to.
+        s3_kms_key (str): kms key used to encrypt the files uploaded to S3.
+        sagemaker_session (str): SageMaker boto client session.
+
+    Returns:
+        str : The concatenated path of all dependencies which will be passed to Spark.
+    """
+    spark_opt_s3_uris = []
+    if not submit_deps:
+        return None
+
+    if not workspace_name or not s3_base_uri:
+        raise ValueError("workspace_name or s3_base_uri may not be empty.")
+
+    for dep_path in submit_deps:
+        dep_url = urlparse(dep_path)
+
+        if dep_url.scheme in ["s3", "s3a"]:
+            spark_opt_s3_uris.append(dep_path)
+        elif not dep_url.scheme or dep_url.scheme == "file":
+            if not os.path.isfile(dep_path):
+                raise ValueError(f"submit_deps path {dep_path} is not a valid local file.")
+
+            upload_path = S3Uploader.upload(
+                local_path=dep_path,
+                desired_s3_uri=s3_path_join(s3_base_uri, workspace_name),
+                kms_key=s3_kms_key,
+                sagemaker_session=sagemaker_session,
+            )
+
+            spark_opt_s3_uris.append(upload_path)
+            logger.info("Uploaded the local file %s to %s", dep_path, upload_path)
+    return str.join(",", spark_opt_s3_uris)
+
+
+def _upload_serialized_spark_configuration(
+    s3_base_uri: str, s3_kms_key: str, configuration: Dict, sagemaker_session: Session
+) -> str:
+    """Upload the Spark configuration json to S3"""
+    if not configuration:
+        return None
+
+    serialized_configuration = BytesIO(json.dumps(configuration).encode("utf-8"))
+    config_file_s3_uri = s3_path_join(s3_base_uri, SPARK_CONF_WORKSPACE, SPARK_CONF_FILE_NAME)
+
+    S3Uploader.upload_string_as_file_body(
+        body=serialized_configuration,
+        desired_s3_uri=config_file_s3_uri,
+        kms_key=s3_kms_key,
+        sagemaker_session=sagemaker_session,
+    )
+
+    logger.info("Uploaded spark configuration json %s to %s", configuration, config_file_s3_uri)
+
+    return config_file_s3_uri
+
+
+def _extend_mpirun_to_request(
+    request_dict: Dict,
+    job_settings: _JobSettings,
+) -> Dict:
+    """Extend the create training job request with mpirun configuration.
+
+    Args:
+        request_dict (Dict): create training job request dict.
+        job_settings (_JobSettings): the job settings.
+    """
+    use_mpirun = job_settings.use_mpirun
+    instance_count = job_settings.instance_count
+
+    if not use_mpirun:
+        return request_dict
+
+    if instance_count == 1:
+        return request_dict
+
+    extended_request = request_dict.copy()
+
+    for input_channel in extended_request["InputDataConfig"]:
+        s3_data_source = input_channel["DataSource"].get("S3DataSource", None)
+        if s3_data_source:
+            s3_data_source["S3DataDistributionType"] = "FullyReplicated"
+
+    return extended_request
+
+
+def _extend_torchrun_to_request(
+    request_dict: Dict,
+    job_settings: _JobSettings,
+) -> Dict:
+    """Extend the create training job request with torchrun configuration.
+
+    Args:
+        request_dict (Dict): create training job request dict.
+        job_settings (_JobSettings): the job settings.
+    """
+    use_torchrun = job_settings.use_torchrun
+    instance_count = job_settings.instance_count
+
+    if not use_torchrun:
+        return request_dict
+
+    if instance_count == 1:
+        return request_dict
+
+    extended_request = request_dict.copy()
+
+    for input_channel in extended_request["InputDataConfig"]:
+        s3_data_source = input_channel["DataSource"].get("S3DataSource", None)
+        if s3_data_source:
+            s3_data_source["S3DataDistributionType"] = "FullyReplicated"
+
+    return extended_request
+
+
+def _extend_spark_config_to_request(
+    request_dict: Dict,
+    job_settings: _JobSettings,
+    s3_base_uri: str,
+) -> Dict:
+    """Extend the create training job request with spark configurations.
+
+    Args:
+        request_dict (Dict): create training job request dict.
+        job_settings (_JobSettings): the job settings.
+        s3_base_uri (str): S3 location that the Spark dependencies will be uploaded to.
+    """
+    spark_config = job_settings.spark_config
+
+    if not spark_config:
+        return request_dict
+
+    extended_request = request_dict.copy()
+    container_entrypoint = extended_request["AlgorithmSpecification"]["ContainerEntrypoint"]
+
+    (
+        submit_jars_s3_paths,
+        submit_py_files_s3_paths,
+        submit_files_s3_path,
+        config_file_s3_uri,
+    ) = _prepare_and_upload_spark_dependent_files(
+        spark_config=spark_config,
+        s3_base_uri=s3_base_uri,
+        s3_kms_key=job_settings.s3_kms_key,
+        sagemaker_session=job_settings.sagemaker_session,
+    )
+
+    input_data_config = extended_request["InputDataConfig"]
+
+    if config_file_s3_uri:
+        input_data_config.append(
+            dict(
+                ChannelName=SPARK_CONF_CHANNEL_NAME,
+                DataSource={
+                    "S3DataSource": {
+                        "S3Uri": config_file_s3_uri,
+                        "S3DataType": "S3Prefix",
+                    }
+                },
+            )
+        )
+
+    for input_channel in extended_request["InputDataConfig"]:
+        s3_data_source = input_channel["DataSource"].get("S3DataSource", None)
+        if s3_data_source:
+            s3_data_source["S3DataDistributionType"] = "FullyReplicated"
+
+    if spark_config.spark_event_logs_uri:
+        container_entrypoint.extend(
+            ["--spark-event-logs-s3-uri", spark_config.spark_event_logs_uri]
+        )
+
+    if submit_jars_s3_paths:
+        container_entrypoint.extend(["--jars", submit_jars_s3_paths])
+
+    if submit_py_files_s3_paths:
+        container_entrypoint.extend(["--py-files", submit_py_files_s3_paths])
+
+    if submit_files_s3_path:
+        container_entrypoint.extend(["--files", submit_files_s3_path])
+
+    if spark_config:
+        container_entrypoint.extend([SPARK_APP_SCRIPT_PATH])
+
+    return extended_request
+
+
+def _update_job_request_with_checkpoint_config(args, kwargs, request_dict):
+    """Extend job request with checkpoint config based on CheckpointLocation in function args.
+
+    Args:
+        args (tuple): The positional arguments of the remote function.
+        kwargs (Dict): The keyword arguments of the remote function.
+        request_dict (Dict): create training job request dict.
+    """
+    checkpoint_location_index_in_args = None
+    checkpoint_location_key_in_kwargs = None
+    checkpoint_location_count = 0
+
+    for index, arg in enumerate(args):
+        if isinstance(arg, CheckpointLocation):
+            checkpoint_location_index_in_args = index
+            checkpoint_location_count += 1
+
+    for key, value in kwargs.items():
+        if isinstance(value, CheckpointLocation):
+            checkpoint_location_key_in_kwargs = key
+            checkpoint_location_count += 1
+
+    if checkpoint_location_count < 1:
+        return
+
+    if checkpoint_location_count > 1:
+        raise ValueError(
+            "Remote function cannot have more than one argument of type CheckpointLocation."
+        )
+
+    if checkpoint_location_index_in_args is not None:
+        checkpoint_location_arg = args[checkpoint_location_index_in_args]
+    else:
+        checkpoint_location_arg = kwargs[checkpoint_location_key_in_kwargs]
+
+    checkpoint_s3_uri = checkpoint_location_arg._s3_uri
+    checkpoint_local_path = checkpoint_location_arg._local_path
+
+    request_dict["CheckpointConfig"] = {
+        "LocalPath": checkpoint_local_path,
+        "S3Uri": checkpoint_s3_uri,
+    }
+
+
+@dataclasses.dataclass
+class _RunInfo:
+    """Data class to hold information of the run object from context."""
+
+    experiment_name: str
+    run_name: str
+
+
+def _get_initial_job_state(description, status_key, wait):
+    """Placeholder docstring"""
+    status = description[status_key]
+    job_already_completed = status in ("Completed", "Failed", "Stopped")
+    return LogState.TAILING if wait and not job_already_completed else LogState.COMPLETE
+
+
+def _logs_for_job(  # noqa: C901 - suppress complexity warning for this method
+    sagemaker_session, job_name, wait=False, poll=10, log_type="All", timeout=None
+):
+    """Display logs for a given training job, optionally tailing them until job is complete.
+
+    If the output is a tty or a Jupyter cell, it will be color-coded
+    based on which instance the log entry is from.
+
+    Args:
+        sagemaker_session (sagemaker.core.helper.session.Session): A SageMaker Session
+            object, used for SageMaker interactions.
+        job_name (str): Name of the training job to display the logs for.
+        wait (bool): Whether to keep looking for new log entries until the job completes
+            (default: False).
+        poll (int): The interval in seconds between polling for new log entries and job
+            completion (default: 5).
+        log_type ([str]): A list of strings specifying which logs to print. Acceptable
+            strings are "All", "None", "Training", or "Rules". To maintain backwards
+            compatibility, boolean values are also accepted and converted to strings.
+        timeout (int): Timeout in seconds to wait until the job is completed. ``None`` by
+            default.
+    Returns:
+        Last call to sagemaker DescribeTrainingJob
+    Raises:
+        exceptions.CapacityError: If the training job fails with CapacityError.
+        exceptions.UnexpectedStatusException: If waiting and the training job fails.
+    """
+    sagemaker_client = sagemaker_session.sagemaker_client
+    request_end_time = time.time() + timeout if timeout else None
+    description = _wait_until(
+        lambda: sagemaker_client.describe_training_job(TrainingJobName=job_name)
+    )
+    print(secondary_training_status_message(description, None), end="")
+
+    instance_count, stream_names, positions, client, log_group, dot, color_wrap = _logs_init(
+        sagemaker_session.boto_session, description, job="Training"
+    )
+
+    state = _get_initial_job_state(description, "TrainingJobStatus", wait)
+
+    # The loop below implements a state machine that alternates between checking the job status
+    # and reading whatever is available in the logs at this point. Note, that if we were
+    # called with wait == False, we never check the job status.
+    #
+    # If wait == TRUE and job is not completed, the initial state is TAILING
+    # If wait == FALSE, the initial state is COMPLETE (doesn't matter if the job really is
+    # complete).
+    #
+    # The state table:
+    #
+    # STATE               ACTIONS                        CONDITION             NEW STATE
+    # ----------------    ----------------               -----------------     ----------------
+    # TAILING             Read logs, Pause, Get status   Job complete          JOB_COMPLETE
+    #                                                    Else                  TAILING
+    # JOB_COMPLETE        Read logs, Pause               Any                   COMPLETE
+    # COMPLETE            Read logs, Exit                                      N/A
+    #
+    # Notes:
+    # - The JOB_COMPLETE state forces us to do an extra pause and read any items that got to
+    #   Cloudwatch after the job was marked complete.
+    last_describe_job_call = time.time()
+    last_description = description
+    last_debug_rule_statuses = None
+    last_profiler_rule_statuses = None
+
+    while True:
+        _flush_log_streams(
+            stream_names,
+            instance_count,
+            client,
+            log_group,
+            job_name,
+            positions,
+            dot,
+            color_wrap,
+        )
+        if timeout and time.time() > request_end_time:
+            print("Timeout Exceeded. {} seconds elapsed.".format(timeout))
+            break
+
+        if state == LogState.COMPLETE:
+            break
+
+        time.sleep(poll)
+
+        if state == LogState.JOB_COMPLETE:
+            state = LogState.COMPLETE
+        elif time.time() - last_describe_job_call >= 30:
+            description = sagemaker_client.describe_training_job(TrainingJobName=job_name)
+            last_describe_job_call = time.time()
+
+            if secondary_training_status_changed(description, last_description):
+                print()
+                print(secondary_training_status_message(description, last_description), end="")
+                last_description = description
+
+            status = description["TrainingJobStatus"]
+
+            if status in ("Completed", "Failed", "Stopped"):
+                print()
+                state = LogState.JOB_COMPLETE
+
+            # Print prettified logs related to the status of SageMaker Debugger rules.
+            debug_rule_statuses = description.get("DebugRuleEvaluationStatuses", {})
+            if (
+                debug_rule_statuses
+                and _rule_statuses_changed(debug_rule_statuses, last_debug_rule_statuses)
+                and (log_type in {"All", "Rules"})
+            ):
+                for status in debug_rule_statuses:
+                    rule_log = (
+                        f"{status['RuleConfigurationName']}: {status['RuleEvaluationStatus']}"
+                    )
+                    print(rule_log)
+
+                last_debug_rule_statuses = debug_rule_statuses
+
+            # Print prettified logs related to the status of SageMaker Profiler rules.
+            profiler_rule_statuses = description.get("ProfilerRuleEvaluationStatuses", {})
+            if (
+                profiler_rule_statuses
+                and _rule_statuses_changed(profiler_rule_statuses, last_profiler_rule_statuses)
+                and (log_type in {"All", "Rules"})
+            ):
+                for status in profiler_rule_statuses:
+                    rule_log = (
+                        f"{status['RuleConfigurationName']}: {status['RuleEvaluationStatus']}"
+                    )
+                    print(rule_log)
+
+                last_profiler_rule_statuses = profiler_rule_statuses
+
+    if wait:
+        _check_job_status(job_name, description, "TrainingJobStatus")
+        if dot:
+            print()
+        # Customers are not billed for hardware provisioning, so billable time is less than
+        # total time
+        training_time = description.get("TrainingTimeInSeconds")
+        billable_time = description.get("BillableTimeInSeconds")
+        if training_time is not None:
+            print("Training seconds:", training_time * instance_count)
+        if billable_time is not None:
+            print("Billable seconds:", billable_time * instance_count)
+            if description.get("EnableManagedSpotTraining"):
+                saving = (1 - float(billable_time) / training_time) * 100
+                print("Managed Spot Training savings: {:.1f}%".format(saving))
+    return last_description
+
+
+def _check_job_status(job, desc, status_key_name):
+    """Check to see if the job completed successfully.
+
+    If not, construct and raise a exceptions. (UnexpectedStatusException).
+
+    Args:
+        job (str): The name of the job to check.
+        desc (dict[str, str]): The result of ``describe_training_job()``.
+        status_key_name (str): Status key name to check for.
+
+    Raises:
+        exceptions.CapacityError: If the training job fails with CapacityError.
+        exceptions.UnexpectedStatusException: If the training job fails.
+    """
+    status = desc[status_key_name]
+    # If the status is capital case, then convert it to Camel case
+    status = _STATUS_CODE_TABLE.get(status, status)
+
+    if status == "Stopped":
+        logger.warning(
+            "Job ended with status 'Stopped' rather than 'Completed'. "
+            "This could mean the job timed out or stopped early for some other reason: "
+            "Consider checking whether it completed as you expect."
+        )
+    elif status != "Completed":
+        reason = desc.get("FailureReason", "(No reason provided)")
+        job_type = status_key_name.replace("JobStatus", " job")
+        troubleshooting = (
+            "https://docs.aws.amazon.com/sagemaker/latest/dg/"
+            "sagemaker-python-sdk-troubleshooting.html"
+        )
+        message = (
+            "Error for {job_type} {job_name}: {status}. Reason: {reason}. "
+            "Check troubleshooting guide for common errors: {troubleshooting}"
+        ).format(
+            job_type=job_type,
+            job_name=job,
+            status=status,
+            reason=reason,
+            troubleshooting=troubleshooting,
+        )
+        if "CapacityError" in str(reason):
+            raise exceptions.CapacityError(
+                message=message,
+                allowed_statuses=["Completed", "Stopped"],
+                actual_status=status,
+            )
+        raise exceptions.UnexpectedStatusException(
+            message=message,
+            allowed_statuses=["Completed", "Stopped"],
+            actual_status=status,
+        )
+
+
+def _flush_log_streams(
+    stream_names, instance_count, client, log_group, job_name, positions, dot, color_wrap
+):
+    """Placeholder docstring"""
+    if len(stream_names) < instance_count:
+        # Log streams are created whenever a container starts writing to stdout/err, so this list
+        # may be dynamic until we have a stream for every instance.
+        try:
+            streams = client.describe_log_streams(
+                logGroupName=log_group,
+                logStreamNamePrefix=job_name + "/",
+                orderBy="LogStreamName",
+                limit=min(instance_count, 50),
+            )
+            stream_names = [s["logStreamName"] for s in streams["logStreams"]]
+
+            while "nextToken" in streams:
+                streams = client.describe_log_streams(
+                    logGroupName=log_group,
+                    logStreamNamePrefix=job_name + "/",
+                    orderBy="LogStreamName",
+                    limit=50,
+                )
+
+                stream_names.extend([s["logStreamName"] for s in streams["logStreams"]])
+
+            positions.update(
+                [
+                    (s, sagemaker_logs.Position(timestamp=0, skip=0))
+                    for s in stream_names
+                    if s not in positions
+                ]
+            )
+        except ClientError as e:
+            # On the very first training job run on an account, there's no log group until
+            # the container starts logging, so ignore any errors thrown about that
+            err = e.response.get("Error", {})
+            if err.get("Code", None) != "ResourceNotFoundException":
+                raise
+
+    if len(stream_names) > 0:
+        if dot:
+            print("")
+            dot = False
+        for idx, event in sagemaker_logs.multi_stream_iter(
+            client, log_group, stream_names, positions
+        ):
+            color_wrap(idx, event["message"])
+            ts, count = positions[stream_names[idx]]
+            if event["timestamp"] == ts:
+                positions[stream_names[idx]] = sagemaker_logs.Position(timestamp=ts, skip=count + 1)
+            else:
+                positions[stream_names[idx]] = sagemaker_logs.Position(
+                    timestamp=event["timestamp"], skip=1
+                )
+    else:
+        dot = True
+        print(".", end="")
+        sys.stdout.flush()
+
+
+def _rule_statuses_changed(current_statuses, last_statuses):
+    """Checks the rule evaluation statuses for SageMaker Debugger and Profiler rules."""
+    if not last_statuses:
+        return True
+
+    for current, last in zip(current_statuses, last_statuses):
+        if (current["RuleConfigurationName"] == last["RuleConfigurationName"]) and (
+            current["RuleEvaluationStatus"] != last["RuleEvaluationStatus"]
+        ):
+            return True
+
+    return False
+
+
+def _get_initial_job_state(description, status_key, wait):
+    """Placeholder docstring"""
+    status = description[status_key]
+    job_already_completed = status in ("Completed", "Failed", "Stopped")
+    return LogState.TAILING if wait and not job_already_completed else LogState.COMPLETE
+
+
+def _logs_init(boto_session, description, job):
+    """Placeholder docstring"""
+    if job == "Training":
+        if "InstanceGroups" in description["ResourceConfig"]:
+            instance_count = 0
+            for instanceGroup in description["ResourceConfig"]["InstanceGroups"]:
+                instance_count += instanceGroup["InstanceCount"]
+        else:
+            instance_count = description["ResourceConfig"]["InstanceCount"]
+    elif job == "Transform":
+        instance_count = description["TransformResources"]["InstanceCount"]
+    elif job == "Processing":
+        instance_count = description["ProcessingResources"]["ClusterConfig"]["InstanceCount"]
+    elif job == "AutoML":
+        instance_count = 0
+
+    stream_names = []  # The list of log streams
+    positions = {}  # The current position in each stream, map of stream name -> position
+
+    # Increase retries allowed (from default of 4), as we don't want waiting for a training job
+    # to be interrupted by a transient exception.
+    config = botocore.config.Config(retries={"max_attempts": 15})
+    client = boto_session.client("logs", config=config)
+    log_group = "/aws/sagemaker/" + job + "Jobs"
+
+    dot = False
+
+    from sagemaker.core.logs import ColorWrap
+
+    color_wrap = ColorWrap()
+
+    return instance_count, stream_names, positions, client, log_group, dot, color_wrap