# 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 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.core.remote_function.invoke_function \\n"
$conda_exe run -n $conda_env python -m sagemaker.core.remote_function.invoke_function "$@"
else
printf "INFO: No conda env provided. Invoking remote function\\n"
printf "INFO: python -m sagemaker.core.remote_function.invoke_function \\n"
python -m sagemaker.core.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.core.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.core.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.core.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.core.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.core.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.core.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.core.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.core.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()
[docs]
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}
)
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):
"""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.
"""
self.job_name = job_name
self.s3_uri = s3_uri
self.sagemaker_session = sagemaker_session
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"]
job = _Job(job_name, s3_uri, sagemaker_session)
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,
)
@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[:]
# 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,
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,
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
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 _decrement_version(version_str: str) -> str:
"""Decrement a version string by one minor or patch version.
Rules:
- If patch version is 0 (e.g., 3.2.0), decrement minor: 3.2.0 -> 3.1.0
- If patch version is not 0 (e.g., 3.1.2), decrement patch: 3.1.2 -> 3.1.1
Args:
version_str: Version string (e.g., "3.2.0")
Returns:
Decremented version string
"""
from packaging import version as pkg_version
try:
parsed = pkg_version.parse(version_str)
major = parsed.major
minor = parsed.minor
patch = parsed.micro
if patch == 0:
# Decrement minor version
minor = max(0, minor - 1)
else:
# Decrement patch version
patch = max(0, patch - 1)
return f"{major}.{minor}.{patch}"
except Exception:
return version_str
def _resolve_version_from_specifier(specifier_str: str) -> str:
"""Resolve the version to check based on upper bounds.
Upper bounds take priority. If upper bound is <4.0.0, it's safe (V3 only).
If no upper bound exists, it's safe (unbounded).
If the decremented upper bound is less than a lower bound, use the lower bound.
Args:
specifier_str: Version specifier string (e.g., ">=3.2.0", "<3.2.0", "==3.1.0")
Returns:
The resolved version string to check, or None if safe
"""
import re
from packaging import version as pkg_version
# Handle exact version pinning (==)
match = re.search(r'==\s*([\d.]+)', specifier_str)
if match:
return match.group(1)
# Extract lower bounds for comparison
lower_bounds = []
for match in re.finditer(r'>=\s*([\d.]+)', specifier_str):
lower_bounds.append(match.group(1))
# Handle upper bounds - find the most restrictive one
upper_bounds = []
# Find all <= bounds
for match in re.finditer(r'<=\s*([\d.]+)', specifier_str):
upper_bounds.append(('<=', match.group(1)))
# Find all < bounds
for match in re.finditer(r'<\s*([\d.]+)', specifier_str):
upper_bounds.append(('<', match.group(1)))
if upper_bounds:
# Sort by version to find the most restrictive (lowest) upper bound
upper_bounds.sort(key=lambda x: pkg_version.parse(x[1]))
operator, version = upper_bounds[0]
# Special case: if upper bound is <4.0.0, it's safe (V3 only)
try:
parsed_upper = pkg_version.parse(version)
if operator == '<' and parsed_upper.major == 4 and parsed_upper.minor == 0 and parsed_upper.micro == 0:
# <4.0.0 means V3 only, which is safe
return None
except Exception:
pass
resolved_version = version
if operator == '<':
resolved_version = _decrement_version(version)
# If we have a lower bound and the resolved version is less than it, use the lower bound
if lower_bounds:
try:
resolved_parsed = pkg_version.parse(resolved_version)
for lower_bound_str in lower_bounds:
lower_parsed = pkg_version.parse(lower_bound_str)
if resolved_parsed < lower_parsed:
resolved_version = lower_bound_str
except Exception:
pass
return resolved_version
# For lower bounds only (>=, >), we don't check
return None
def _check_sagemaker_version_compatibility(sagemaker_requirement: str) -> None:
"""Check if the sagemaker version requirement uses incompatible hashing.
Raises ValueError if the requirement would install a version that uses HMAC hashing
(which is incompatible with the current SHA256-based integrity checks).
Args:
sagemaker_requirement: The sagemaker requirement string (e.g., "sagemaker>=3.2.0")
Raises:
ValueError: If the requirement would install a version using HMAC hashing
"""
import re
from packaging import version as pkg_version
match = re.search(r'sagemaker\s*(.+)$', sagemaker_requirement.strip(), re.IGNORECASE)
if not match:
return
specifier_str = match.group(1).strip()
# Resolve the version that would be installed
resolved_version_str = _resolve_version_from_specifier(specifier_str)
if not resolved_version_str:
# No upper bound or exact version, so we can't determine if it's bad
return
try:
resolved_version = pkg_version.parse(resolved_version_str)
except Exception:
return
# Define HMAC thresholds for each major version
v2_hmac_threshold = pkg_version.parse("2.256.0")
v3_hmac_threshold = pkg_version.parse("3.2.0")
# Check if the resolved version uses HMAC hashing
uses_hmac = False
if resolved_version.major == 2 and resolved_version < v2_hmac_threshold:
uses_hmac = True
elif resolved_version.major == 3 and resolved_version < v3_hmac_threshold:
uses_hmac = True
if uses_hmac:
raise ValueError(
f"The sagemaker version specified in requirements.txt ({sagemaker_requirement}) "
f"could install a version using HMAC-based integrity checks which are incompatible "
f"with the current SHA256-based integrity checks. Please update to "
f"sagemaker>=2.256.0,<3.0.0 (for V2) or sagemaker>=3.2.0,<4.0.0 (for V3)."
)
def _ensure_sagemaker_dependency(local_dependencies_path: str) -> str:
"""Ensure sagemaker>=3.2.0 is in the dependencies.
This function ensures that the remote environment has a compatible version of sagemaker
that includes the fix for the HMAC key security issue. Versions < 3.2.0 use HMAC-based
integrity checks which require the REMOTE_FUNCTION_SECRET_KEY environment variable.
Versions >= 3.2.0 use SHA256-based integrity checks which are secure and don't require
the secret key.
If no dependencies are provided, creates a temporary requirements.txt with sagemaker.
If dependencies are provided, appends sagemaker if not already present.
Args:
local_dependencies_path: Path to user's dependencies file or None
Returns:
Path to the dependencies file (created or modified)
Raises:
ValueError: If user has pinned sagemaker to a version using HMAC hashing
"""
import tempfile
SAGEMAKER_MIN_VERSION = "sagemaker>=3.2.0,<4.0.0"
if local_dependencies_path is None:
fd, req_file = tempfile.mkstemp(suffix=".txt", prefix="sagemaker_requirements_")
os.close(fd)
with open(req_file, "w") as f:
f.write(f"{SAGEMAKER_MIN_VERSION}\n")
logger.info("Created temporary requirements.txt at %s with %s", req_file, SAGEMAKER_MIN_VERSION)
return req_file
if local_dependencies_path.endswith(".txt"):
with open(local_dependencies_path, "r") as f:
content = f.read()
if "sagemaker" in content.lower():
for line in content.split('\n'):
if 'sagemaker' in line.lower():
_check_sagemaker_version_compatibility(line.strip())
break
else:
with open(local_dependencies_path, "a") as f:
f.write(f"\n{SAGEMAKER_MIN_VERSION}\n")
logger.info("Appended %s to requirements.txt", SAGEMAKER_MIN_VERSION)
return local_dependencies_path
def _generate_input_data_config(job_settings, s3_base_uri):
"""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)
# Ensure sagemaker dependency is included to prevent version mismatch issues
# Resolves issue where computing hash for integrity check changed in 3.2.0
local_dependencies_path = _ensure_sagemaker_dependency(local_dependencies_path)
job_settings.dependencies = local_dependencies_path
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