# 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.
"""ModelTrainer class module."""
from __future__ import absolute_import
from enum import Enum
import os
import json
import re
import shlex
import shutil
from tempfile import TemporaryDirectory
from typing import Optional, List, Union, Dict, Any, ClassVar
import yaml
from botocore.exceptions import ClientError, NoCredentialsError, NoRegionError
from graphene.utils.str_converters import to_camel_case, to_snake_case
from sagemaker.core.config.config_manager import SageMakerConfig
from sagemaker.core import resources
from sagemaker.core.resources import TrainingJob
from sagemaker.core import shapes
from sagemaker.core.shapes import AlgorithmSpecification, ModelPackageConfig
from sagemaker.core.utils.utils import serialize
from sagemaker.core.apiutils._boto_functions import to_pascal_case
from pydantic import BaseModel, ConfigDict, PrivateAttr, validate_call
from sagemaker.core.config.config_schema import (
_simple_path,
SAGEMAKER,
MODEL_TRAINER,
MODULES,
PYTHON_SDK,
TRAINING_JOB_ENVIRONMENT_PATH,
TRAINING_JOB_ENABLE_NETWORK_ISOLATION_PATH,
TRAINING_JOB_VPC_CONFIG_PATH,
TRAINING_JOB_SUBNETS_PATH,
TRAINING_JOB_SECURITY_GROUP_IDS_PATH,
TRAINING_JOB_OUTPUT_DATA_CONFIG_PATH,
TRAINING_JOB_RESOURCE_CONFIG_PATH,
TRAINING_JOB_ROLE_ARN_PATH,
TRAINING_JOB_TAGS_PATH,
)
from sagemaker.core.helper.session_helper import Session
from sagemaker.train import configs
from sagemaker.train.configs import (
Compute,
StoppingCondition,
RetryStrategy,
SourceCode,
TrainingImageConfig,
Channel,
DataSource,
S3DataSource,
FileSystemDataSource,
Networking,
Tag,
InfraCheckConfig,
RemoteDebugConfig,
SessionChainingConfig,
InputData,
MetricDefinition,
)
from sagemaker.train.distributed import Torchrun, DistributedConfig
from sagemaker.train.utils import (
_default_s3_uri,
_get_unique_name,
_is_valid_path,
_is_valid_s3_uri,
safe_serialize,
)
from sagemaker.train.types import DataSourceType
from sagemaker.train.constants import (
SM_CODE,
SM_CODE_CONTAINER_PATH,
SM_DRIVERS,
SM_DRIVERS_LOCAL_PATH,
SM_RECIPE,
SM_RECIPE_YAML,
SM_RECIPE_CONTAINER_PATH,
TRAIN_SCRIPT,
DEFAULT_CONTAINER_ENTRYPOINT,
DEFAULT_CONTAINER_ARGUMENTS,
SOURCE_CODE_JSON,
DISTRIBUTED_JSON,
)
from sagemaker.train.templates import (
TRAIN_SCRIPT_TEMPLATE,
EXECUTE_BASE_COMMANDS,
EXEUCTE_DISTRIBUTED_DRIVER,
EXECUTE_BASIC_SCRIPT_DRIVER,
INSTALL_AUTO_REQUIREMENTS,
INSTALL_REQUIREMENTS,
)
from sagemaker.core.telemetry.telemetry_logging import _telemetry_emitter, TelemetryParamType
from sagemaker.core.telemetry.constants import Feature
from sagemaker.core.utils.utils import logger
from sagemaker.train.sm_recipes.utils import (
_get_args_from_recipe,
_determine_device_type,
_is_nova_recipe,
_is_llmft_recipe,
_load_base_recipe,
)
from sagemaker.core.jumpstart.configs import JumpStartConfig
from sagemaker.core.jumpstart.document import get_hub_content_and_document
from sagemaker.core.jumpstart.utils import get_eula_url
from sagemaker.train.defaults import TrainDefaults, JumpStartTrainDefaults
from sagemaker.core.workflow.pipeline_context import PipelineSession, runnable_by_pipeline
from sagemaker.core.helper.pipeline_variable import StrPipeVar
from sagemaker.train.local.local_container import _LocalContainer
[docs]
class Mode(Enum):
"""Enum class for training mode."""
LOCAL_CONTAINER = "LOCAL_CONTAINER"
SAGEMAKER_TRAINING_JOB = "SAGEMAKER_TRAINING_JOB"
def _log_actionable_client_error(client_error: "ClientError") -> None:
"""Log remediation guidance for common CreateTrainingJob service errors.
This never swallows the error: the caller always re-raises the original
exception. It only adds guidance so the failure is actionable without
guesswork. ResourceLimitExceeded in particular is not retryable until the
quota is raised, so blind retry loops keep failing.
Args:
client_error (ClientError): The error raised by the CreateTrainingJob call.
"""
error = getattr(client_error, "response", None) or {}
error_code = error.get("Error", {}).get("Code", "")
error_message = error.get("Error", {}).get("Message", "")
if error_code == "ResourceLimitExceeded":
quota_match = re.search(r"service limit '([^']+)'", error_message)
quota_hint = f" ('{quota_match.group(1)}')" if quota_match else ""
logger.error(
"CreateTrainingJob was rejected because a SageMaker service quota%s "
"is insufficient in this account and Region. Retrying will keep "
"failing until the quota is raised. Request an increase in the "
"Service Quotas console: "
"https://console.aws.amazon.com/servicequotas/home/services/sagemaker/quotas "
"(or via 'aws service-quotas request-service-quota-increase "
"--service-code sagemaker'). Alternatively, use a different instance "
"type or reduce the instance count.",
quota_hint,
)
elif error_code in ("AccessDeniedException", "AccessDenied"):
logger.error(
"CreateTrainingJob was denied by IAM. Verify that your caller identity "
"and the training role have the sagemaker:CreateTrainingJob and "
"iam:PassRole permissions for the role passed to the ModelTrainer."
)
[docs]
class ModelTrainer(BaseModel):
"""Class that trains a model using AWS SageMaker.
Example:
.. code:: python
from sagemaker.train import ModelTrainer
from sagemaker.train.configs import SourceCode, Compute, InputData
ignore_patterns = ['.env', '.git', '__pycache__', '.DS_Store', 'data']
source_code = SourceCode(source_dir="source", entry_script="train.py", ignore_patterns=ignore_patterns)
training_image = "123456789012.dkr.ecr.us-west-2.amazonaws.com/my-training-image"
model_trainer = ModelTrainer(
training_image=training_image,
source_code=source_code,
)
train_data = InputData(channel_name="train", data_source="s3://bucket/train")
model_trainer.train(input_data_config=[train_data])
training_job = model_trainer._latest_training_job
Parameters:
training_mode (Mode):
The training mode. Valid values are "Mode.LOCAL_CONTAINER" or
"Mode.SAGEMAKER_TRAINING_JOB".
sagemaker_session (Optiona(Session)):
The SageMaker Session. This object can be used to manage underlying
boto3 clients and to specify aritfact upload paths via the ``default_bucket``
and ``default_bucket_prefix`` attributes.
If not specified, a new session will be created.
role (Optional(str)):
The IAM role ARN for the training job.
If not specified, the default SageMaker execution role will be used.
base_job_name (Optional[str]):
The base name for the training job.
If not specified, a default name will be generated using the algorithm name
or training image name.
input_s3_key_prefix (Optional[str]):
The leading S3 key prefix under which locally-staged input data (source code,
SDK drivers, and local input channels) is uploaded. This is useful for grouping
a training job's artifacts under a custom prefix such as a pipeline name, e.g.
``s3://<default_bucket_path>/<input_s3_key_prefix>/<job_name>/input/``. If not
specified, ``base_job_name`` is used as the leading prefix (unchanged behavior).
source_code (Optional[SourceCode]):
The source code configuration. This is used to configure the source code for
running the training job.
distributed (Optional[DistributedConfig]):
The distributed runner for the training job. This is used to configure
a distributed training job. If specifed, ``source_code`` must also
be provided.
compute (Optional[Compute]):
The compute configuration. This is used to specify the compute resources for
the training job. If not specified, will default to 1 instance of ml.m5.xlarge.
networking (Optional[Networking]):
The networking configuration. This is used to specify the networking settings
for the training job.
stopping_condition (Optional[StoppingCondition]):
The stopping condition. This is used to specify the different stopping
conditions for the training job.
If not specified, will default to 1 hour max run time.
algorithm_name (Optional[str]):
The SageMaker marketplace algorithm name/arn to use for the training job.
algorithm_name cannot be specified if training_image is specified.
training_image (Optional[str]):
The training image URI to use for the training job container.
training_image cannot be specified if algorithm_name is specified.
To find available sagemaker distributed images,
see: https://docs.aws.amazon.com/sagemaker/latest/dg-ecr-paths/sagemaker-algo-docker-registry-paths
training_image_config (Optional[TrainingImageConfig]):
Training image Config. This is the configuration to use an image from a private
Docker registry for a training job.
output_data_config (Optional[OutputDataConfig]):
The output data configuration. This is used to specify the output data location
for the training job.
If not specified in the session, will default to
``s3://<default_bucket>/<default_prefix>/<base_job_name>/``.
input_data_config (Optional[List[Union[Channel, InputData]]]):
The input data config for the training job.
Takes a list of Channel or InputData objects. An InputDataSource can be an S3 URI
string, local file path string, S3DataSource object, or FileSystemDataSource object.
checkpoint_config (Optional[CheckpointConfig]):
Contains information about the output location for managed spot training checkpoint
data.
training_input_mode (Optional[str]):
The input mode for the training job. Valid values are "Pipe", "File", "FastFile".
Defaults to "File".
environment (Optional[Dict[str, str]]):
The environment variables for the training job.
hyperparameters (Optional[Union[Dict[str, Any], str]):
The hyperparameters for the training job. Can be a dictionary of hyperparameters
or a path to hyperparameters json/yaml file.
tags (Optional[List[Tag]]):
An array of key-value pairs. You can use tags to categorize your AWS resources
in different ways, for example, by purpose, owner, or environment.
local_container_root (Optional[str]):
The local root directory to store artifacts from a training job launched in
"LOCAL_CONTAINER" mode.
model_package_config (Optional[ModelPackageConfig]):
The model package configuration for the training job. Used for Nova RMP
(Restricted Model Package) jobs to specify SourceModelPackageArn and
ModelPackageGroupArn.
"""
model_config = ConfigDict(
arbitrary_types_allowed=True, validate_assignment=True, extra="forbid"
)
training_mode: Mode = Mode.SAGEMAKER_TRAINING_JOB
sagemaker_session: Optional[Session] = None
role: Optional[str] = None
base_job_name: Optional[str] = None
input_s3_key_prefix: Optional[str] = None
source_code: Optional[SourceCode] = None
distributed: Optional[DistributedConfig] = None
compute: Optional[Compute] = None
networking: Optional[Networking] = None
stopping_condition: Optional[StoppingCondition] = None
training_image: Optional[StrPipeVar] = None
training_image_config: Optional[TrainingImageConfig] = None
algorithm_name: Optional[StrPipeVar] = None
output_data_config: Optional[shapes.OutputDataConfig] = None
input_data_config: Optional[List[Union[Channel, InputData]]] = None
checkpoint_config: Optional[shapes.CheckpointConfig] = None
training_input_mode: Optional[StrPipeVar] = "File"
environment: Optional[Dict[str, StrPipeVar]] = {}
hyperparameters: Optional[Union[Dict[str, Any], str]] = {}
tags: Optional[List[Tag]] = None
local_container_root: Optional[str] = os.getcwd()
model_package_config: Optional[ModelPackageConfig] = None
# Created Artifacts
_latest_training_job: Optional[resources.TrainingJob] = PrivateAttr(default=None)
# Private TrainingJob Parameters
_tensorboard_output_config: Optional[shapes.TensorBoardOutputConfig] = PrivateAttr(default=None)
_retry_strategy: Optional[RetryStrategy] = PrivateAttr(default=None)
_infra_check_config: Optional[InfraCheckConfig] = PrivateAttr(default=None)
_session_chaining_config: Optional[SessionChainingConfig] = PrivateAttr(default=None)
_remote_debug_config: Optional[RemoteDebugConfig] = PrivateAttr(default=None)
_metric_definitions: Optional[List[MetricDefinition]] = PrivateAttr(default=None)
_is_nova_recipe: Optional[bool] = PrivateAttr(default=None)
_is_llmft_recipe: Optional[bool] = PrivateAttr(default=None)
# Private Attributes for Recipes
_temp_recipe_train_dir: Optional[TemporaryDirectory] = PrivateAttr(default=None)
# Private Attributes for JumpStart
_jumpstart_config: Optional[JumpStartConfig] = PrivateAttr(default=None)
# Private Attributes for AWS_Batch
_temp_code_dir: Optional[TemporaryDirectory] = PrivateAttr(default=None)
CONFIGURABLE_ATTRIBUTES: ClassVar[List[str]] = [
"role",
"base_job_name",
"input_s3_key_prefix",
"source_code",
"compute",
"networking",
"stopping_condition",
"training_image",
"training_image_config",
"algorithm_name",
"output_data_config",
"checkpoint_config",
"training_input_mode",
"environment",
"hyperparameters",
]
SERIALIZABLE_CONFIG_ATTRIBUTES: ClassVar[Any] = {
"source_code": SourceCode,
"compute": Compute,
"networking": Networking,
"stopping_condition": StoppingCondition,
"training_image_config": TrainingImageConfig,
"output_data_config": configs.OutputDataConfig,
"checkpoint_config": configs.CheckpointConfig,
}
config_mgr: SageMakerConfig = SageMakerConfig()
@property
def output_data(self) -> str:
"""The S3 URI of the training job's ``output.tar.gz`` archive.
This is the non-model output archive that SageMaker uploads from the
container's ``/opt/ml/output/data`` directory after training completes,
distinct from the model artifact (``model.tar.gz``). The URI follows the
layout ``{s3_output_path}/{training_job_name}/output/output.tar.gz``.
The S3 output path is resolved from the latest training job's
``output_data_config`` and falls back to the trainer's own
``output_data_config`` when the job does not report one.
Note:
This is a computed S3 URI derived from the job's configuration; it is
not verified to exist. The archive is only present once the job has
completed successfully and produced output data.
Returns:
str: The fully-qualified S3 URI to the job's ``output.tar.gz``.
Raises:
ValueError: If no training job has been created yet (``train`` has
not been called), or if no S3 output path can be resolved.
"""
training_job = self._latest_training_job
if training_job is None:
raise ValueError(
"No training job is associated with this ModelTrainer. "
"Call train() before accessing output_data."
)
# Prefer the S3 output path reported by the job, falling back to the
# trainer's own config. ``Unassigned`` and ``None`` are both falsy.
s3_output_path = None
job_output_config = getattr(training_job, "output_data_config", None)
if job_output_config:
s3_output_path = getattr(job_output_config, "s3_output_path", None)
if not s3_output_path and self.output_data_config is not None:
s3_output_path = getattr(self.output_data_config, "s3_output_path", None)
if not s3_output_path:
raise ValueError(
"Unable to resolve the output S3 path for this training job. "
"Ensure output_data_config is set on the ModelTrainer or the "
"training job."
)
s3_output_path = s3_output_path.rstrip("/")
return f"{s3_output_path}/{training_job.training_job_name}/output/output.tar.gz"
def _populate_intelligent_defaults(self):
"""Function to populate all the possible default configs
Model Trainer specific configs take precedence over the generic training job ones.
"""
self._populate_intelligent_defaults_from_model_trainer_space()
self._populate_intelligent_defaults_from_training_job_space()
def _populate_intelligent_defaults_from_training_job_space(self):
"""Function to populate all the possible default configs from Training Job Space"""
if not self.environment:
self.environment = self.config_mgr.resolve_value_from_config(
config_path=TRAINING_JOB_ENVIRONMENT_PATH, sagemaker_session=self.sagemaker_session
)
default_enable_network_isolation = self.config_mgr.resolve_value_from_config(
config_path=TRAINING_JOB_ENABLE_NETWORK_ISOLATION_PATH,
sagemaker_session=self.sagemaker_session,
)
default_vpc_config = self.config_mgr.resolve_value_from_config(
config_path=TRAINING_JOB_VPC_CONFIG_PATH, sagemaker_session=self.sagemaker_session
)
if not self.networking:
if default_enable_network_isolation is not None or default_vpc_config is not None:
self.networking = Networking(
enable_network_isolation=default_enable_network_isolation,
subnets=self.config_mgr.resolve_value_from_config(
config_path=TRAINING_JOB_SUBNETS_PATH
),
security_group_ids=self.config_mgr.resolve_value_from_config(
config_path=TRAINING_JOB_SECURITY_GROUP_IDS_PATH
),
)
else:
if self.networking.enable_network_isolation is None:
self.networking.enable_network_isolation = default_enable_network_isolation
if self.networking.subnets is None:
self.networking.subnets = self.config_mgr.resolve_value_from_config(
config_path=TRAINING_JOB_SUBNETS_PATH
)
if self.networking.security_group_ids is None:
self.networking.security_group_ids = self.config_mgr.resolve_value_from_config(
config_path=TRAINING_JOB_SECURITY_GROUP_IDS_PATH
)
if not self.output_data_config:
default_output_data_config = self.config_mgr.resolve_value_from_config(
config_path=TRAINING_JOB_OUTPUT_DATA_CONFIG_PATH
)
if default_output_data_config:
self.output_data_config = configs.OutputDataConfig(
**self._convert_keys_to_snake(default_output_data_config)
)
if not self.compute:
default_resource_config = self.config_mgr.resolve_value_from_config(
config_path=TRAINING_JOB_RESOURCE_CONFIG_PATH
)
if default_resource_config:
self.compute = Compute(**self._convert_keys_to_snake(default_resource_config))
if not self.role:
self.role = self.config_mgr.resolve_value_from_config(
config_path=TRAINING_JOB_ROLE_ARN_PATH
)
if not self.tags:
self.tags = self.config_mgr.resolve_value_from_config(
config_path=TRAINING_JOB_TAGS_PATH
)
def _convert_keys_to_snake(self, config: dict) -> dict:
"""Utility helper function that converts the keys of a dictionary into snake case"""
return {to_snake_case(key): value for key, value in config.items()}
def _populate_intelligent_defaults_from_model_trainer_space(self):
"""Function to populate all the possible default configs from Model Trainer Space"""
for configurable_attribute in self.CONFIGURABLE_ATTRIBUTES:
if getattr(self, configurable_attribute) is None:
default_config = self.config_mgr.resolve_value_from_config(
config_path=_simple_path(
SAGEMAKER,
PYTHON_SDK,
MODULES,
MODEL_TRAINER,
to_camel_case(configurable_attribute),
),
sagemaker_session=self.sagemaker_session,
)
if default_config is not None:
if configurable_attribute in self.SERIALIZABLE_CONFIG_ATTRIBUTES:
default_config = self.SERIALIZABLE_CONFIG_ATTRIBUTES.get(
configurable_attribute
)(
**default_config # pylint: disable=E1134
)
setattr(self, configurable_attribute, default_config)
def __del__(self):
"""Destructor method to clean up the temporary directory."""
# Clean up the temporary directory if it exists and class was initialized
if hasattr(self, "__pydantic_fields_set__"):
if self._temp_recipe_train_dir is not None:
self._temp_recipe_train_dir.cleanup()
if self._temp_code_dir is not None:
self._temp_code_dir.cleanup()
def _validate_training_image_and_algorithm_name(
self, training_image: Optional[str], algorithm_name: Optional[str]
):
"""Validate that only one of 'training_image' or 'algorithm_name' is provided."""
if not training_image and not algorithm_name:
raise ValueError(
"Atleast one of 'training_image' or 'algorithm_name' must be provided.",
)
if training_image and algorithm_name:
raise ValueError(
"Only one of 'training_image' or 'algorithm_name' must be provided.",
)
def _validate_distributed_config(
self,
source_code: Optional[SourceCode],
distributed: Optional[DistributedConfig],
):
"""Validate the distribution configuration."""
if distributed and not source_code.entry_script:
raise ValueError(
"Must provide 'entry_script' if 'distribution' " + "is provided in 'source_code'.",
)
def _validate_source_code(self, source_code: Optional[SourceCode]):
"""Validate the source code configuration."""
if source_code:
if source_code.requirements or source_code.entry_script:
source_dir = source_code.source_dir
requirements = source_code.requirements
entry_script = source_code.entry_script
if not source_dir:
raise ValueError(
"If 'requirements' or 'entry_script' is provided in 'source_code', "
"'source_dir' must also be provided."
)
if not (
_is_valid_path(source_dir, path_type="Directory")
or _is_valid_s3_uri(source_dir, path_type="Directory")
or (
_is_valid_path(source_dir, path_type="File")
and source_dir.endswith(".tar.gz")
)
or (
_is_valid_s3_uri(source_dir, path_type="File")
and source_dir.endswith(".tar.gz")
)
):
raise ValueError(
f"Invalid 'source_dir' path: {source_dir}. "
"Must be a valid local directory, "
"s3 uri or path to tar.gz file stored locally or in s3."
)
if requirements:
if not source_dir.endswith(".tar.gz"):
if not _is_valid_path(
f"{source_dir}/{requirements}", path_type="File"
) and not _is_valid_s3_uri(
f"{source_dir}/{requirements}", path_type="File"
):
raise ValueError(
f"Invalid 'requirements': {requirements}. "
"Must be a valid file within the 'source_dir'.",
)
if entry_script:
if not source_dir.endswith(".tar.gz"):
if not _is_valid_path(
f"{source_dir}/{entry_script}", path_type="File"
) and not _is_valid_s3_uri(
f"{source_dir}/{entry_script}", path_type="File"
):
raise ValueError(
f"Invalid 'entry_script': {entry_script}. "
"Must be a valid file within the 'source_dir'.",
)
@staticmethod
def _validate_and_fetch_hyperparameters_file(hyperparameters_file: str):
"""Validate and fetch hyperparameters from a file."""
if not os.path.exists(hyperparameters_file):
raise ValueError(f"Hyperparameters file not found: {hyperparameters_file}")
logger.info(f"Loading hyperparameters from file: {hyperparameters_file}")
with open(hyperparameters_file, "r") as f:
contents = f.read()
try:
hyperparameters = json.loads(contents)
logger.debug("Hyperparameters loaded as JSON")
except json.JSONDecodeError:
try:
logger.info(f"contents: {contents}")
hyperparameters = yaml.safe_load(contents)
if not isinstance(hyperparameters, dict):
raise ValueError("YAML contents must be a valid mapping")
logger.info(f"hyperparameters: {hyperparameters}")
logger.debug("Hyperparameters loaded as YAML")
except (yaml.YAMLError, ValueError):
raise ValueError(
f"Invalid hyperparameters file: {hyperparameters_file}. "
"Must be a valid JSON or YAML file."
)
return hyperparameters
[docs]
def model_post_init(self, __context: Any):
"""Post init method to perform custom validation and set default values."""
self._validate_training_image_and_algorithm_name(self.training_image, self.algorithm_name)
self._validate_source_code(self.source_code)
self._validate_distributed_config(self.source_code, self.distributed)
if self.hyperparameters and isinstance(self.hyperparameters, str):
self.hyperparameters = self._validate_and_fetch_hyperparameters_file(
hyperparameters_file=self.hyperparameters
)
if self.training_mode == Mode.SAGEMAKER_TRAINING_JOB:
self.sagemaker_session = TrainDefaults.get_sagemaker_session(self.sagemaker_session)
self.role = TrainDefaults.get_role(
role=self.role, sagemaker_session=self.sagemaker_session
)
self.base_job_name = TrainDefaults.get_base_job_name(
base_job_name=self.base_job_name,
algorithm_name=self.algorithm_name,
training_image=self.training_image,
)
self.compute = TrainDefaults.get_compute(compute=self.compute)
self.stopping_condition = TrainDefaults.get_stopping_condition(
stopping_condition=self.stopping_condition
)
if self.training_mode == Mode.SAGEMAKER_TRAINING_JOB:
self.output_data_config = TrainDefaults.get_output_data_config(
base_job_name=self.base_job_name,
output_data_config=self.output_data_config,
sagemaker_session=self.sagemaker_session,
)
if self.training_image:
from sagemaker.core.helper.pipeline_variable import PipelineVariable
if isinstance(self.training_image, PipelineVariable):
logger.info(
"Training image URI: (PipelineVariable - resolved at pipeline execution)"
)
else:
logger.info(f"Training image URI: {self.training_image}")
def _create_training_job_args(
self,
input_data_config: Optional[List[Union[Channel, InputData]]] = None,
boto3: bool = False,
) -> Dict[str, Any]:
"""Create the training job arguments.
Args:
input_data_config (Optional[List[Union[Channel, InputData]]]):
input_data_config (Optional[List[Union[Channel, InputData]]]):
The input data config for the training job.
Takes a list of Channel objects or a dictionary of channel names to DataSourceType.
DataSourceType can be an S3 URI string, local file path string,
S3DataSource object, or FileSystemDataSource object.
boto3 (bool): Whether to return the arguments in boto3 format. Defaults to False.
By default, the arguments are returned in the format used by the SageMaker Core.
Returns:
Dict[str, Any]: The training job arguments.
"""
self._populate_intelligent_defaults()
current_training_job_name = _get_unique_name(self.base_job_name)
# Use input_s3_key_prefix (e.g. a pipeline name) as the leading key prefix when set;
# otherwise fall back to base_job_name to preserve the default upload layout.
input_data_key_prefix_root = self.input_s3_key_prefix or self.base_job_name
input_data_key_prefix = f"{input_data_key_prefix_root}/{current_training_job_name}/input"
final_input_data_config = self.input_data_config.copy() if self.input_data_config else []
if input_data_config:
# merge the inputs with method parameter taking precedence
existing_channels = {input.channel_name: input for input in final_input_data_config}
new_channels = []
for new_input in input_data_config:
if new_input.channel_name in existing_channels:
existing_channels[new_input.channel_name] = new_input
else:
new_channels.append(new_input)
final_input_data_config = list(existing_channels.values()) + new_channels
# Assign SDK-managed channels to instance groups when the user uses them (see helper).
managed_channel_instance_group_names = self._resolve_managed_channel_instance_groups(
final_input_data_config
)
if self._is_nova_recipe or self._is_llmft_recipe:
for input_data in final_input_data_config:
if input_data.channel_name == SM_RECIPE:
raise ValueError(
"Cannot use reserved channel name 'recipe' as an input channel name "
" for Nova or LLMFT Recipe"
)
recipe_file_path = os.path.join(self._temp_recipe_train_dir.name, SM_RECIPE_YAML)
recipe_channel = self.create_input_data_channel(
channel_name=SM_RECIPE,
data_source=recipe_file_path,
key_prefix=input_data_key_prefix,
instance_group_names=managed_channel_instance_group_names,
)
final_input_data_config.append(recipe_channel)
if self._is_nova_recipe or self._is_llmft_recipe:
self.hyperparameters.update(
{"sagemaker_recipe_local_path": SM_RECIPE_CONTAINER_PATH}
)
if final_input_data_config:
final_input_data_config = self._get_input_data_config(
final_input_data_config, input_data_key_prefix
)
if self.checkpoint_config and not self.checkpoint_config.s3_uri:
self.checkpoint_config.s3_uri = _default_s3_uri(
self.sagemaker_session,
f"{self.base_job_name}/{current_training_job_name}/checkpoints",
)
if self._tensorboard_output_config and not self._tensorboard_output_config.s3_output_path:
self._tensorboard_output_config.s3_output_path = _default_s3_uri(
self.sagemaker_session, self.base_job_name
)
string_hyper_parameters = {}
if self.hyperparameters:
for hyper_parameter, value in self.hyperparameters.items():
string_hyper_parameters[hyper_parameter] = safe_serialize(value)
container_entrypoint = None
container_arguments = None
if self.source_code:
if self.training_mode == Mode.LOCAL_CONTAINER:
self._temp_code_dir = TemporaryDirectory(
prefix=os.path.join(self.local_container_root + "/")
)
else:
self._temp_code_dir = TemporaryDirectory()
# Copy everything under container_drivers/ to a temporary directory
shutil.copytree(SM_DRIVERS_LOCAL_PATH, self._temp_code_dir.name, dirs_exist_ok=True)
# Copy the CodeArtifact-aware install_requirements script from sagemaker-core
# so it's available in the container at /opt/ml/input/data/sm_drivers/scripts/
import sagemaker.core.utils.install_requirements as _ir_mod
shutil.copy2(
_ir_mod.__file__,
os.path.join(self._temp_code_dir.name, "scripts", "install_requirements.py"),
)
# If distributed is provided, overwrite code under <root>/drivers
if self.distributed:
distributed_driver_dir = self.distributed.driver_dir
driver_dir = os.path.join(self._temp_code_dir.name, "distributed_drivers")
shutil.copytree(distributed_driver_dir, driver_dir, dirs_exist_ok=True)
# If source code is provided, create a channel for the source code
# The source code will be mounted at /opt/ml/input/data/code in the container
if self.source_code.source_dir:
source_code_channel = self.create_input_data_channel(
channel_name=SM_CODE,
data_source=self.source_code.source_dir,
key_prefix=input_data_key_prefix,
ignore_patterns=self.source_code.ignore_patterns,
instance_group_names=managed_channel_instance_group_names,
)
final_input_data_config.append(source_code_channel)
self._prepare_train_script(
tmp_dir=self._temp_code_dir,
source_code=self.source_code,
distributed=self.distributed,
)
if isinstance(self.distributed, Torchrun) and self.distributed.smp:
mp_parameters = self.distributed.smp._to_mp_hyperparameters()
string_hyper_parameters.update(mp_parameters)
self._write_source_code_json(tmp_dir=self._temp_code_dir, source_code=self.source_code)
self._write_distributed_json(tmp_dir=self._temp_code_dir, distributed=self.distributed)
# Create an input channel for drivers packaged by the sdk.
# Do NOT apply the user's source_code.ignore_patterns here: this directory is
# SDK-owned driver content (e.g. scripts/environment.py), and user patterns such as
# "scripts" or "environment" would strip files the container bootstrap requires,
# causing "sm_drivers/scripts/environment.py: No such file or directory" (issue #5493).
sm_drivers_channel = self.create_input_data_channel(
channel_name=SM_DRIVERS,
data_source=self._temp_code_dir.name,
key_prefix=input_data_key_prefix,
instance_group_names=managed_channel_instance_group_names,
)
final_input_data_config.append(sm_drivers_channel)
# If source_code is provided, we will always use
# the default container entrypoint and arguments
# to execute the sm_train.sh script.
# Any commands generated from the source_code will be
# executed from the sm_train.sh script.
container_entrypoint = DEFAULT_CONTAINER_ENTRYPOINT
container_arguments = DEFAULT_CONTAINER_ARGUMENTS
algorithm_specification = AlgorithmSpecification(
algorithm_name=self.algorithm_name,
training_image=self.training_image,
training_input_mode=self.training_input_mode,
training_image_config=self.training_image_config,
container_entrypoint=container_entrypoint,
container_arguments=container_arguments,
metric_definitions=self._metric_definitions,
)
resource_config = self.compute._to_resource_config()
vpc_config = self.networking._to_vpc_config() if self.networking else None
# Convert tags to dictionaries if they are Tag objects
tags_as_dicts = None
if self.tags:
tags_as_dicts = []
for tag in self.tags:
if hasattr(tag, "model_dump"):
tags_as_dicts.append(tag.model_dump())
elif isinstance(tag, dict):
tags_as_dicts.append(tag)
else:
# Fallback for any other tag-like object
tags_as_dicts.append(
{"key": getattr(tag, "key", ""), "value": getattr(tag, "value", "")}
)
# Build training request with snake_case keys (Python SDK convention)
training_request = {
"training_job_name": current_training_job_name,
"algorithm_specification": algorithm_specification,
"hyper_parameters": string_hyper_parameters,
"input_data_config": (final_input_data_config if final_input_data_config else None),
"resource_config": resource_config,
"vpc_config": vpc_config,
"role_arn": self.role,
"tags": tags_as_dicts,
"stopping_condition": self.stopping_condition,
"output_data_config": self.output_data_config,
"checkpoint_config": self.checkpoint_config,
"environment": self.environment,
"enable_managed_spot_training": self.compute.enable_managed_spot_training,
"enable_inter_container_traffic_encryption": (
self.networking.enable_inter_container_traffic_encryption
if self.networking
else None
),
"enable_network_isolation": (
self.networking.enable_network_isolation if self.networking else None
),
"remote_debug_config": self._remote_debug_config,
"tensor_board_output_config": self._tensorboard_output_config,
"retry_strategy": self._retry_strategy,
"infra_check_config": self._infra_check_config,
"session_chaining_config": self._session_chaining_config,
}
if self.model_package_config:
training_request["model_package_config"] = self.model_package_config
if boto3 or isinstance(self.sagemaker_session, PipelineSession):
if isinstance(self.sagemaker_session, PipelineSession):
training_request.pop("training_job_name", None)
# ``serialize`` renames the fields of Tag models to the boto-style
# {"Key": ..., "Value": ...} form the service requires, but passes plain dicts
# through untouched, so hand it the validated Tag objects. ``TrainingJob.create``
# takes the dicts instead, which it revalidates into its own Tag model.
training_request["tags"] = self.tags
# Convert snake_case to PascalCase for AWS API
pipeline_request = {to_pascal_case(k): v for k, v in training_request.items()}
serialized_request = serialize(pipeline_request)
return serialized_request
return training_request
[docs]
@_telemetry_emitter(
feature=Feature.MODEL_TRAINER,
func_name="model_trainer.train",
telemetry_params=[
("training_mode", TelemetryParamType.ATTR_VALUE),
("training_input_mode", TelemetryParamType.ATTR_VALUE),
("networking", TelemetryParamType.ATTR_EXISTS),
("stopping_condition", TelemetryParamType.ATTR_EXISTS),
("distributed", TelemetryParamType.ATTR_EXISTS),
("source_code", TelemetryParamType.ATTR_EXISTS),
("checkpoint_config", TelemetryParamType.ATTR_EXISTS),
],
)
@runnable_by_pipeline
@validate_call
def train(
self,
input_data_config: Optional[List[Union[Channel, InputData]]] = None,
wait: Optional[bool] = True,
logs: Optional[bool] = True,
dry_run: bool = False,
):
"""Train a model using AWS SageMaker.
Args:
input_data_config (Optional[List[Union[Channel, InputData]]]):
The input data config for the training job.
Takes a list of Channel objects or a dictionary of channel names to DataSourceType.
DataSourceType can be an S3 URI string, local file path string,
S3DataSource object, or FileSystemDataSource object.
wait (Optional[bool]):
Whether to wait for the training job to complete before returning.
Defaults to True.
logs (Optional[bool]):
Whether to display the training container logs while training.
Defaults to True.
dry_run (bool):
If True, runs all validation (configuration, input resolution, hyperparameters)
without submitting a job. Returns None on success, raises on validation failure.
Defaults to False.
"""
training_request = self._create_training_job_args(input_data_config=input_data_config)
job_name = training_request.get("training_job_name")
if job_name:
logger.info(f"Training Job Name: {job_name}")
if dry_run:
logger.info("Dry-run validation passed. No job submitted.")
return None
# Handle PipelineSession
if self.training_mode == Mode.SAGEMAKER_TRAINING_JOB:
if isinstance(self.sagemaker_session, PipelineSession):
self.sagemaker_session._intercept_create_request(training_request, None, "train")
return
try:
training_job = TrainingJob.create(
session=self.sagemaker_session.boto_session, **training_request
)
except ClientError as ce:
_log_actionable_client_error(ce)
raise
except NoRegionError:
logger.error(
"No AWS Region configured. Set one before calling train(): export "
"AWS_DEFAULT_REGION=<region>, add 'region = <region>' to ~/.aws/config, "
"or pass a boto session with a region to the ModelTrainer's "
"sagemaker_session."
)
raise
except NoCredentialsError:
logger.error(
"No AWS credentials found. Configure credentials before calling train(): "
"run 'aws configure' or 'aws sso login', set the AWS_ACCESS_KEY_ID/"
"AWS_SECRET_ACCESS_KEY environment variables, or use an IAM role. See "
"https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html"
)
raise
self._latest_training_job = training_job
if wait:
training_job.wait(logs=logs)
if logs and not wait:
logger.warning(
"Not displaing the training container logs as 'wait' is set to False."
)
else:
if self.compute is not None and getattr(self.compute, "instance_preferences", None):
raise ValueError(
"Local mode training does not support 'instance_preferences'. "
"Set a single 'instance_type' on Compute for local mode."
)
local_container = _LocalContainer(
training_job_name=training_request["training_job_name"],
instance_type=training_request["resource_config"].instance_type,
instance_count=training_request["resource_config"].instance_count,
image=training_request["algorithm_specification"].training_image,
container_root=self.local_container_root,
sagemaker_session=self.sagemaker_session,
container_entrypoint=training_request[
"algorithm_specification"
].container_entrypoint,
container_arguments=training_request["algorithm_specification"].container_arguments,
input_data_config=training_request["input_data_config"] or [],
hyper_parameters=training_request["hyper_parameters"],
environment=training_request["environment"],
)
local_container.train(wait)
if self._temp_code_dir is not None:
self._temp_code_dir.cleanup()
def _resolve_staging_bucket(self) -> tuple[str, str]:
"""Resolve the S3 bucket and key prefix for staging training artifacts.
Uses iam:SimulatePrincipalPolicy to check whether the training role
can access the default SageMaker bucket. If not, falls back to the
s3_output_path bucket and prefix.
Returns:
tuple: (bucket_name, prefix_override) where prefix_override is None
when using the default bucket, or the output path prefix when
falling back.
"""
default_bucket = self.sagemaker_session.default_bucket()
if not self.role:
logger.debug(
"No training role specified; skipping bucket access check. "
"Using default bucket '%s' for artifact staging.",
default_bucket,
)
return default_bucket, None
try:
iam_client = self.sagemaker_session.boto_session.client("iam")
result = iam_client.simulate_principal_policy(
PolicySourceArn=self.role,
ActionNames=["s3:PutObject"],
ResourceArns=[f"arn:aws:s3:::{default_bucket}/*"],
)
decisions = result.get("EvaluationResults", [])
if decisions and decisions[0].get("EvalDecision") != "allowed":
# Training role can't access default bucket — fall back to output path
if self.output_data_config and hasattr(self.output_data_config, "s3_output_path"):
output_path = self.output_data_config.s3_output_path
if output_path and output_path.startswith("s3://"):
from urllib.parse import urlparse
parsed = urlparse(output_path)
if parsed.netloc:
prefix = parsed.path.strip("/")
logger.info(
f"Training role '{self.role}' cannot access default bucket "
f"'{default_bucket}'. Using output path bucket "
f"'{parsed.netloc}/{prefix}' for artifact staging."
)
return parsed.netloc, prefix
except Exception as e:
logger.debug(f"Could not verify role access to default bucket: {e}")
return default_bucket, None
def _resolve_managed_channel_instance_groups(
self, input_data_config: List[Union[Channel, InputData]]
) -> Optional[List[str]]:
"""Resolve the instance groups to assign to SDK-managed channels.
SageMaker requires that, on a heterogeneous cluster, either all channels are
assigned to instance groups or none are. The SDK injects channels (recipe, code,
sm_drivers) that users cannot configure, so when a user assigns instance groups to
any of their channels, the managed channels must be assigned too. We assign them to
the full set of instance groups defined on the compute config, since the code and
drivers must be present on every node regardless of instance group.
Args:
input_data_config (List[Union[Channel, InputData]]): The user-provided channels.
Returns:
Optional[List[str]]: The full list of instance group names to assign to
managed channels, or ``None`` if instance groups are not in use.
"""
instance_groups = getattr(self.compute, "instance_groups", None) if self.compute else None
if not instance_groups:
return None
def _channel_has_instance_groups(channel: Union[Channel, InputData]) -> bool:
# Channel nests its S3DataSource under data_source; InputData may hold it directly.
s3_data_source = None
if isinstance(channel, Channel):
if channel.data_source:
s3_data_source = channel.data_source.s3_data_source
elif isinstance(channel, InputData):
if isinstance(channel.data_source, S3DataSource):
s3_data_source = channel.data_source
return bool(getattr(s3_data_source, "instance_group_names", None))
if not any(_channel_has_instance_groups(channel) for channel in input_data_config):
return None
return [group.instance_group_name for group in instance_groups]
def _resolve_upload_extra_args(self, kms_key: Optional[str] = None) -> Optional[dict]:
"""Build S3 ``ExtraArgs`` for encrypting local source/data uploads with KMS.
Uses the explicit ``kms_key`` if given, otherwise falls back to the
``kms_key_id`` on this trainer's ``output_data_config`` (V2 ``Estimator`` parity:
staged user code is encrypted with the output KMS key). Returns ``None`` when no
usable string key is configured, so behavior is unchanged by default (GH #5956).
A pipeline-variable key is ignored here because uploads happen at SDK compile time.
"""
resolved_key = kms_key
if resolved_key is None and self.output_data_config is not None:
resolved_key = getattr(self.output_data_config, "kms_key_id", None)
if isinstance(resolved_key, str) and resolved_key:
return {"ServerSideEncryption": "aws:kms", "SSEKMSKeyId": resolved_key}
return None
def _get_input_data_config(
self,
input_data_channels: Optional[List[Union[Channel, InputData]]],
key_prefix: Optional[str] = None,
) -> List[Channel]:
"""Get the input data configuration for the training job.
Args:
input_data_channels (Optional[List[Union[Channel, InputData]]]):
The input data config for the training job.
Takes a list of Channel or InputData objects. An InputDataSource can be an S3 URI
string, local file path string, S3DataSource object, or FileSystemDataSource object.
"""
if input_data_channels is None:
return []
channels = []
for input_data in input_data_channels:
if isinstance(input_data, Channel):
channels.append(input_data)
elif isinstance(input_data, InputData):
channel = self.create_input_data_channel(
input_data.channel_name, input_data.data_source, key_prefix=key_prefix
)
if input_data.content_type:
channel.content_type = input_data.content_type
channels.append(channel)
else:
raise ValueError(
f"Invalid input data channel: {input_data}. "
"Must be a Channel or InputDataSource."
)
return channels
def _write_source_code_json(self, tmp_dir: TemporaryDirectory, source_code: SourceCode):
"""Write the source code configuration to a JSON file."""
file_path = os.path.join(tmp_dir.name, SOURCE_CODE_JSON)
with open(file_path, "w") as f:
dump = source_code.model_dump() if source_code else {}
f.write(json.dumps(dump))
def _write_distributed_json(
self,
tmp_dir: TemporaryDirectory,
distributed: Optional[DistributedConfig] = None,
):
"""Write the distributed runner configuration to a JSON file."""
file_path = os.path.join(tmp_dir.name, DISTRIBUTED_JSON)
with open(file_path, "w") as f:
dump = distributed.model_dump() if distributed else {}
f.write(json.dumps(dump))
def _prepare_train_script(
self,
tmp_dir: TemporaryDirectory,
source_code: SourceCode,
distributed: Optional[DistributedConfig] = None,
):
"""Prepare the training script to be executed in the training job container.
Args:
source_code (SourceCode): The source code configuration.
"""
base_command = ""
if source_code.command:
if source_code.entry_script:
logger.warning(
"Both 'command' and 'entry_script' are provided in the SourceCode. "
"Defaulting to 'command'."
)
command_parts = source_code.command.split()
if source_code.args:
# Shell-quote each arg so values with spaces or special characters are
# passed through as single, literal arguments to the executed command.
command_parts.extend(shlex.quote(str(arg)) for arg in source_code.args)
base_command = " ".join(command_parts)
if self.hyperparameters:
logger.warning(
"Hyperparameters are set but are not passed as command-line arguments "
"when 'command' is used in the SourceCode. They are available inside the "
"training container via the 'SM_HPS' environment variable, e.g. "
"`hps = json.loads(os.environ['SM_HPS'])`. To pass arguments to your "
"'command', use the SourceCode 'args' field."
)
install_requirements = ""
if source_code.requirements:
if self._jumpstart_config and source_code.requirements == "auto":
install_requirements = INSTALL_AUTO_REQUIREMENTS
else:
install_requirements = INSTALL_REQUIREMENTS.format(
requirements_file=source_code.requirements
)
working_dir = ""
if source_code.source_dir:
working_dir = f"cd {SM_CODE_CONTAINER_PATH} \n"
if source_code.source_dir.endswith(".tar.gz"):
tarfile_name = os.path.basename(source_code.source_dir)
working_dir += f"tar -xzf {tarfile_name} \n"
if base_command:
execute_driver = EXECUTE_BASE_COMMANDS.format(base_command=base_command)
elif distributed:
execute_driver = EXEUCTE_DISTRIBUTED_DRIVER.format(
driver_name=distributed.__class__.__name__,
driver_script=distributed.driver_script,
)
elif source_code.entry_script and not source_code.command and not distributed:
if not source_code.entry_script.endswith((".py", ".sh")):
raise ValueError(
f"Unsupported entry script: {source_code.entry_script}."
+ "Only .py and .sh scripts are supported."
)
execute_driver = EXECUTE_BASIC_SCRIPT_DRIVER
else:
# This should never be reached, as the source_code should have been validated.
raise ValueError(
f"Unsupported SourceCode or DistributedConfig: {source_code}, {distributed}."
"Please provide a valid configuration with atleast one of 'command'"
" or 'entry_script'."
)
train_script = TRAIN_SCRIPT_TEMPLATE.format(
working_dir=working_dir,
install_requirements=install_requirements,
execute_driver=execute_driver,
)
with open(os.path.join(tmp_dir.name, TRAIN_SCRIPT), "w", newline="\n") as f:
f.write(train_script)
[docs]
@classmethod
def from_recipe(
cls,
training_recipe: str,
compute: Compute,
recipe_overrides: Optional[Dict[str, Any]] = None,
networking: Optional[Networking] = None,
stopping_condition: Optional[StoppingCondition] = None,
requirements: Optional[str] = None,
training_image: Optional[str] = None,
training_image_config: Optional[TrainingImageConfig] = None,
output_data_config: Optional[shapes.OutputDataConfig] = None,
input_data_config: Optional[List[Union[Channel, InputData]]] = None,
checkpoint_config: Optional[shapes.CheckpointConfig] = None,
training_input_mode: Optional[str] = "File",
environment: Optional[Dict[str, str]] = None,
hyperparameters: Optional[Union[Dict[str, Any], str]] = {},
tags: Optional[List[Tag]] = None,
sagemaker_session: Optional[Session] = None,
role: Optional[str] = None,
base_job_name: Optional[str] = None,
model_package_config: Optional[ModelPackageConfig] = None,
) -> "ModelTrainer": # noqa: D412
"""Create a ModelTrainer from a training recipe.
Example:
.. code:: python
from sagemaker.train import ModelTrainer
from sagemaker.train.configs import Compute
recipe_overrides = {
"run": {
"results_dir": "/opt/ml/model",
},
"model": {
"data": {
"use_synthetic_data": True
}
}
}
compute = Compute(
instance_type="ml.p5.48xlarge",
keep_alive_period_in_seconds=3600
)
model_trainer = ModelTrainer.from_recipe(
training_recipe="fine-tuning/deepseek/hf_deepseek_r1_distilled_llama_8b_seq8k_gpu_fine_tuning",
recipe_overrides=recipe_overrides,
compute=compute,
)
model_trainer.train(wait=False)
Args:
training_recipe (str):
The training recipe to use for training the model. This must be the name of
a sagemaker training recipe or a path to a local training recipe .yaml file.
For available training recipes, see: https://github.com/aws/sagemaker-hyperpod-recipes/
compute (Compute):
The compute configuration. This is used to specify the compute resources for
the training job. Specifying instance_type is required for training recipes.
Must be a GPU or Tranium instance type.
recipe_overrides (Optional[Dict[str, Any]]):
The recipe overrides. This is used to override the default recipe parameters.
networking (Optional[Networking]):
The networking configuration. This is used to specify the networking settings
for the training job.
stopping_condition (Optional[StoppingCondition]):
The stopping condition. This is used to specify the different stopping
conditions for the training job.
If not specified, will default to 1 hour max run time.
requirements (Optional[str]):
The path to a requirements file to install in the training job container.
training_image (Optional[str]):
The training image URI to use for the training job container. If not specified,
the training image will be determined from the recipe.
training_image_config (Optional[TrainingImageConfig]):
Training image Config. This is the configuration to use an image from a private
Docker registry for a training job.
output_data_config (Optional[OutputDataConfig]):
The output data configuration. This is used to specify the output data location
for the training job.
If not specified, will default to ``s3://<default_bucket>/<base_job_name>/output/``.
input_data_config (Optional[List[Union[Channel, InputData]]]):
The input data config for the training job.
Takes a list of Channel or InputData objects. An InputDataSource can be an S3 URI
string, local file path string, S3DataSource object, or FileSystemDataSource object.
checkpoint_config (Optional[CheckpointConfig]):
Contains information about the output location for managed spot training checkpoint
data.
training_input_mode (Optional[str]):
The input mode for the training job. Valid values are "Pipe", "File", "FastFile".
Defaults to "File".
environment (Optional[Dict[str, str]]):
The environment variables for the training job.
tags (Optional[List[Tag]]):
An array of key-value pairs. You can use tags to categorize your AWS resources
in different ways, for example, by purpose, owner, or environment.
sagemaker_session (Optional[Session]):
The SageMakerCore session.
If not specified, a new session will be created.
role (Optional[str]):
The IAM role ARN for the training job.
If not specified, the default SageMaker execution role will be used.
base_job_name (Optional[str]):
The base name for the training job.
If not specified, a default name will be generated using the algorithm name
or training image.
model_package_config (Optional[ModelPackageConfig]):
The model package configuration. Used for Nova RMP jobs to specify
SourceModelPackageArn and ModelPackageGroupArn. If also specified in recipe,
direct param wins on conflict.
"""
if getattr(compute, "instance_preferences", None):
raise ValueError(
"Training recipes do not support ``instance_preferences``. "
"Set a single ``instance_type`` in Compute when using training recipes."
)
if compute.instance_type is None:
raise ValueError("Must set ``instance_type`` in Compute when using training recipes.")
device_type = _determine_device_type(compute.instance_type)
recipe = _load_base_recipe(
training_recipe=training_recipe, recipe_overrides=recipe_overrides
)
is_nova = _is_nova_recipe(recipe=recipe)
is_llmft = _is_llmft_recipe(recipe=recipe)
if device_type == "cpu" and not (is_nova or is_llmft):
raise ValueError(
"Training recipes are not supported for CPU instances. "
"Please provide a GPU or Tranium instance type."
)
if training_image is None and (is_nova or is_llmft):
raise ValueError("training_image must be provided when using recipe for Nova or LLMFT")
if training_image_config and training_image is None:
raise ValueError("training_image must be provided when using training_image_config.")
sagemaker_session = TrainDefaults.get_sagemaker_session(sagemaker_session)
# The training recipe is used to prepare the following args:
# - source_code
# - training_image
# - distributed
# - compute
# - hyperparameters
model_trainer_args, tmp_dir = _get_args_from_recipe(
training_recipe=recipe,
recipe_overrides=recipe_overrides,
requirements=requirements,
compute=compute,
region_name=sagemaker_session.boto_region_name,
role=role,
)
if training_image is not None:
model_trainer_args["training_image"] = training_image
if hyperparameters and not is_nova:
logger.warning(
"Hyperparameters are not supported for general and LLMFT training recipes. "
+ "Ignoring hyperparameters input."
)
if is_nova:
if hyperparameters and isinstance(hyperparameters, str):
hyperparameters = cls._validate_and_fetch_hyperparameters_file(hyperparameters)
model_trainer_args["hyperparameters"].update(hyperparameters)
elif hyperparameters and isinstance(hyperparameters, dict):
model_trainer_args["hyperparameters"].update(hyperparameters)
# Pop recipe model_package_config dict before unpacking into constructor
recipe_mpc_dict = model_trainer_args.pop("model_package_config", None) or {}
model_trainer = cls(
sagemaker_session=sagemaker_session,
role=role,
base_job_name=base_job_name,
networking=networking,
stopping_condition=stopping_condition,
training_image_config=training_image_config,
output_data_config=output_data_config,
input_data_config=input_data_config,
checkpoint_config=checkpoint_config,
training_input_mode=training_input_mode,
environment=environment,
tags=tags,
**model_trainer_args,
)
# Merge ModelPackageConfig: recipe dict + direct Pydantic param (direct wins)
direct_mpc_dict = (
model_package_config.model_dump(exclude_unset=True) if model_package_config else {}
)
merged = {**recipe_mpc_dict, **direct_mpc_dict}
if merged:
model_trainer.model_package_config = ModelPackageConfig(**merged)
model_trainer._is_nova_recipe = is_nova
model_trainer._is_llmft_recipe = is_llmft
model_trainer._temp_recipe_train_dir = tmp_dir
model_trainer._resolved_recipe_cache = None
model_trainer._training_recipe = training_recipe
model_trainer._recipe_overrides = recipe_overrides
return model_trainer
[docs]
def get_resolved_recipe(self) -> Dict[str, Any]:
"""Return the fully resolved recipe configuration.
Shows the final merged result of base recipe + recipe_overrides after
OmegaConf interpolation resolution. Callable before or after train().
Works for all recipe types (Nova, LLMFT, general).
Returns:
Dict[str, Any]: Deep copy of the resolved recipe configuration.
Raises:
ValueError: If recipe resolution fails.
AttributeError: If called on a ModelTrainer not created via from_recipe().
"""
if not hasattr(self, "_training_recipe"):
raise AttributeError(
"get_resolved_recipe() is only available on ModelTrainer instances "
"created via ModelTrainer.from_recipe()."
)
if self._resolved_recipe_cache is not None:
import copy
return copy.deepcopy(self._resolved_recipe_cache)
from omegaconf import OmegaConf
from sagemaker.train.sm_recipes.utils import _register_custom_resolvers
import copy
recipe = _load_base_recipe(
training_recipe=self._training_recipe,
recipe_overrides=self._recipe_overrides,
)
_register_custom_resolvers()
try:
OmegaConf.resolve(recipe)
except Exception as e:
logger.warning(
f"Could not fully resolve recipe interpolations: {e}. "
f"Returning partially resolved recipe."
)
resolved = OmegaConf.to_container(recipe, resolve=True)
self._resolved_recipe_cache = resolved
return copy.deepcopy(resolved)
[docs]
@classmethod
def from_jumpstart_config(
cls,
jumpstart_config: JumpStartConfig,
source_code: Optional[SourceCode] = None,
compute: Optional[Compute] = None,
networking: Optional[Networking] = None,
stopping_condition: Optional[StoppingCondition] = None,
training_image: Optional[str] = None,
training_image_config: Optional[TrainingImageConfig] = None,
output_data_config: Optional[shapes.OutputDataConfig] = None,
input_data_config: Optional[List[Union[Channel, InputData]]] = None,
checkpoint_config: Optional[shapes.CheckpointConfig] = None,
training_input_mode: Optional[str] = "File",
environment: Optional[Dict[str, str]] = {},
hyperparameters: Optional[Union[Dict[str, Any], str]] = {},
tags: Optional[List[Tag]] = None,
sagemaker_session: Optional[Session] = None,
role: Optional[str] = None,
base_job_name: Optional[str] = None,
) -> "ModelTrainer": # noqa: D412
"""Create a ModelTrainer from a JumpStart Model ID.
.. code:: python
from sagemaker.train import ModelTrainer
from sagemaker.train.configs import InputData
from sagemaker.core.jumpstart import JumpStartConfig
jumpstart_config = JumpStartConfig(model_id="xxxxxxx")
model_trainer = ModelTrainer.from_jumpstart_config(
jumpstart_config=jumpstart_config
)
training_data = InputData(channel_name="training", data_source="s3://bucket/path")
model_trainer.train(input_data_config=[training_data])
Args:
jumpstart_config (JumpStart):
The JumpStart model configuration. This is used to specify the model ID,
version, and other parameters for the training job.
source_code (Optional[SourceCode]):
The source code configuration. This is used to configure the source code for
running the training job.
compute (Optional[Compute]):
The compute configuration. This is used to specify the compute resources for
the training job.
networking (Optional[Networking]):
The networking configuration. This is used to specify the networking settings
for the training job.
stopping_condition (Optional[StoppingCondition]):
The stopping condition. This is used to specify the different stopping
conditions for the training job.
If not specified, will default to 1 hour max run time.
training_image (Optional[str]):
The training image URI to use for the training job container. If not specified,
the training image will be determined from the recipe.
training_image_config (Optional[TrainingImageConfig]):
Training image Config. This is the configuration to use an image from a private
Docker registry for a training job.
output_data_config (Optional[OutputDataConfig]):
The output data configuration. This is used to specify the output data location
for the training job. If not specified, will default
to ``s3://<default_bucket>/<default_prefix>/<base_job_name>/``.
input_data_config (Optional[List[Union[Channel, InputData]]]):
The input data config for the training job.
Takes a list of Channel or InputData objects. An InputDataSource can be an S3 URI
string, local file path string, S3DataSource object, or FileSystemDataSource object.
checkpoint_config (Optional[CheckpointConfig]):
Contains information about the output location for managed spot training checkpoint
data.
training_input_mode (Optional[str]):
The input mode for the training job. Valid values are "Pipe", "File", "FastFile".
Defaults to "File".
environment (Optional[Dict[str, str]]):
The environment variables for the training job.
hyperparameters (Optional[Union[Dict[str, Any], str]]):
The hyperparameters for the training job. Can be a dictionary of hyperparameters
or a path to hyperparameters json/yaml file.
tags (Optional[List[Tag]]):
An array of key-value pairs. You can use tags to categorize your AWS resources
in different ways, for example, by purpose, owner, or environment.
sagemaker_session (Optional[Session]):
The SageMaker Session. This object can be used to manage underlying
boto3 clients and to specify aritfact upload paths via the ``default_bucket``
and ``default_bucket_prefix`` attributes.
If not specified, a new session will be created.
role (Optional[str]):
The IAM role ARN for the training job.
If not specified, the default SageMaker execution role will be used.
base_job_name (Optional[str]):
The base name for the training job.
If not specified, a default name will be generated using the algorithm name
or training image name.
"""
sagemaker_session = TrainDefaults.get_sagemaker_session(sagemaker_session=sagemaker_session)
role = TrainDefaults.get_role(role=role, sagemaker_session=sagemaker_session)
_, document = get_hub_content_and_document(
jumpstart_config=jumpstart_config, sagemaker_session=sagemaker_session
)
# Basic Validation
if not document.TrainingSupported:
raise ValueError(
f"Training is not supported for the model ID: {jumpstart_config.model_id}.\n"
"Please check that the model ID is available for training."
)
if compute and getattr(compute, "instance_preferences", None):
raise ValueError(
"JumpStart training does not support ``instance_preferences``. "
"Set a single ``instance_type`` in Compute for JumpStart models."
)
if compute and document.SupportedTrainingInstanceTypes:
# Optional[List] is guarded by the enclosing ``if``; pylint cannot see that.
if (
compute.instance_type
not in document.SupportedTrainingInstanceTypes # pylint: disable=unsupported-membership-test
):
raise ValueError(
"Training is not supported for model ID with instance type: "
f" {compute.instance_type}.\n"
"This model ID is only supported for the following instance types:\n"
f"{document.SupportedTrainingInstanceTypes}.\n"
)
if document.GatedBucket:
eula_url = get_eula_url(sagemaker_session=sagemaker_session, document=document)
if not jumpstart_config.accept_eula:
raise ValueError(
f"Model {jumpstart_config.model_id} is a gated model "
"and requires accepting the EULA via the `accept_eula` parameter.\n"
f"See {eula_url} for terms of use."
)
logger.warning(f"Model {jumpstart_config.model_id} is a gated model ")
print(f"See {eula_url} for terms of use.")
compute = JumpStartTrainDefaults.get_compute(
jumpstart_config=jumpstart_config,
compute=compute,
sagemaker_session=sagemaker_session,
)
networking = JumpStartTrainDefaults.get_networking(
jumpstart_config=jumpstart_config,
networking=networking,
sagemaker_session=sagemaker_session,
)
training_image = JumpStartTrainDefaults.get_training_image(
jumpstart_config=jumpstart_config,
compute=compute,
training_image=training_image,
sagemaker_session=sagemaker_session,
)
base_job_name = JumpStartTrainDefaults.get_base_job_name(
jumpstart_config=jumpstart_config,
base_job_name=base_job_name,
)
environment = JumpStartTrainDefaults.get_enviornment(
jumpstart_config=jumpstart_config,
compute=compute,
environment=environment,
sagemaker_session=sagemaker_session,
)
if hyperparameters and isinstance(hyperparameters, str):
hyperparameters = cls._validate_and_fetch_hyperparameters_file(hyperparameters)
hyperparameters = JumpStartTrainDefaults.get_hyperparameters(
jumpstart_config=jumpstart_config,
compute=compute,
hyperparameters=hyperparameters,
environment=environment,
sagemaker_session=sagemaker_session,
)
source_code = JumpStartTrainDefaults.get_source_code(
jumpstart_config=jumpstart_config,
source_code=source_code,
sagemaker_session=sagemaker_session,
)
input_data_config = JumpStartTrainDefaults.get_training_dataset_input(
jumpstart_config=jumpstart_config,
input_data_config=input_data_config,
sagemaker_session=sagemaker_session,
)
input_data_config = JumpStartTrainDefaults.get_model_artifact_input(
jumpstart_config=jumpstart_config,
compute=compute,
input_data_config=input_data_config,
environment=environment,
sagemaker_session=sagemaker_session,
)
output_data_config = JumpStartTrainDefaults.get_output_data_config(
jumpstart_config=jumpstart_config,
base_job_name=base_job_name,
output_data_config=output_data_config,
sagemaker_session=sagemaker_session,
)
tags = JumpStartTrainDefaults.get_tags(
jumpstart_config=jumpstart_config,
tags=tags,
sagemaker_session=sagemaker_session,
)
model_trainer = cls(
source_code=source_code,
compute=compute,
networking=networking,
stopping_condition=stopping_condition,
training_image=training_image,
training_image_config=training_image_config,
output_data_config=output_data_config,
input_data_config=input_data_config,
checkpoint_config=checkpoint_config,
training_input_mode=training_input_mode,
environment=environment,
hyperparameters=hyperparameters,
tags=tags,
sagemaker_session=sagemaker_session,
role=role,
base_job_name=base_job_name,
)
model_trainer._jumpstart_config = jumpstart_config
return model_trainer
[docs]
def with_tensorboard_output_config(
self, tensorboard_output_config: Optional[shapes.TensorBoardOutputConfig] = None
) -> "ModelTrainer": # noqa: D412
"""Set the TensorBoard output configuration.
Example:
.. code:: python
from sagemaker.train import ModelTrainer
model_trainer = ModelTrainer(
...
).with_tensorboard_output_config()
Args:
tensorboard_output_config (sagemaker.train.configs.TensorBoardOutputConfig):
The TensorBoard output configuration.
"""
self._tensorboard_output_config = (
tensorboard_output_config or configs.TensorBoardOutputConfig()
)
return self
[docs]
def with_retry_strategy(self, retry_strategy: RetryStrategy) -> "ModelTrainer": # noqa: D412
"""Set the retry strategy for the training job.
Example:
..code:: python
from sagemaker.train import ModelTrainer
from sagemaker.train.configs import RetryStrategy
retry_strategy = RetryStrategy(maximum_retry_attempts=3)
model_trainer = ModelTrainer(
...
).with_retry_strategy(retry_strategy)
Args:
retry_strategy (sagemaker.train.configs.RetryStrategy):
The retry strategy for the training job.
"""
self._retry_strategy = retry_strategy
return self
[docs]
def with_infra_check_config(
self, infra_check_config: Optional[InfraCheckConfig] = None
) -> "ModelTrainer": # noqa: D412
"""Set the infra check configuration for the training job.
Example:
..code:: python
from sagemaker.train import ModelTrainer
model_trainer = ModelTrainer(
...
).with_infra_check_config(infra_check_config)
Args:
infra_check_config (sagemaker.train.configs.InfraCheckConfig):
The infra check configuration for the training job.
"""
self._infra_check_config = infra_check_config or InfraCheckConfig(enable_infra_check=True)
return self
[docs]
def with_session_chaining_config(
self, session_chaining_config: Optional[SessionChainingConfig] = None
) -> "ModelTrainer": # noqa: D412
"""Set the session chaining configuration for the training job.
Example:
..code:: python
from sagemaker.train import ModelTrainer
model_trainer = ModelTrainer(
...
).with_session_chaining_config()
Args:
session_chaining_config (sagemaker.train.configs.SessionChainingConfig):
The session chaining configuration for the training job.
"""
self._session_chaining_config = session_chaining_config or SessionChainingConfig(
enable_session_tag_chaining=True
)
return self
[docs]
def with_remote_debug_config(
self, remote_debug_config: RemoteDebugConfig
) -> "ModelTrainer": # noqa: D412
"""Set the remote debug configuration for the training job.
Example:
..code:: python
from sagemaker.train import ModelTrainer
model_trainer = ModelTrainer(
...
).with_remote_debug_config()
Args:
remote_debug_config (sagemaker.train.configs.RemoteDebugConfig):
The remote debug configuration for the training job.
"""
self._remote_debug_config = remote_debug_config or RemoteDebugConfig(
enable_remote_debug=True
)
return self
[docs]
def with_checkpoint_config(
self, checkpoint_config: Optional[shapes.CheckpointConfig] = None
) -> "ModelTrainer": # noqa: D412
"""Set the checkpoint configuration for the training job.
Example:
.. code:: python
from sagemaker.train import ModelTrainer
model_trainer = ModelTrainer(
...
).with_checkpoint_config()
Args:
checkpoint_config (sagemaker.modules.configs.CheckpointConfig):
The checkpoint configuration for the training job.
"""
self.checkpoint_config = checkpoint_config or configs.CheckpointConfig()
return self
[docs]
def with_metric_definitions(
self, metric_definitions: List[MetricDefinition]
) -> "ModelTrainer": # noqa: D412
"""Set the metric definitions for the training job.
Example:
.. code:: python
from sagemaker.modules.train import ModelTrainer
from sagemaker.modules.configs import MetricDefinition
metric_definitions = [
MetricDefinition(
name="loss",
regex="Loss: (.*?)",
)
]
model_trainer = ModelTrainer(
...
).with_metric_definitions(metric_definitions)
Args:
metric_definitions (List[MetricDefinition]):
The metric definitions for the training job.
"""
self._metric_definitions = metric_definitions
return self