# 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.
"""This module provides the configuration classes used in ``sagemaker.modules``.
Some of these classes are re-exported from ``sagemaker.core.shapes``. For convinence,
users can import these classes directly from ``sagemaker.modules.configs``.
For more documentation on ``sagemaker.core.shapes``, see:
- https://sagemaker-core.readthedocs.io/en/stable/#sagemaker-core-shapes
"""
from __future__ import absolute_import
from typing import Optional, Union, List
from pydantic import BaseModel, model_validator, ConfigDict
import sagemaker.core.shapes as shapes
from sagemaker.core.helper.pipeline_variable import StrPipeVar, IntPipeVar, BoolPipeVar
# TODO: Can we add custom logic to some of these to set better defaults?
from sagemaker.core.shapes import (
StoppingCondition,
RetryStrategy,
Channel,
ShuffleConfig,
DataSource,
S3DataSource,
FileSystemDataSource,
TrainingImageConfig,
TrainingRepositoryAuthConfig,
Tag,
InfraCheckConfig,
RemoteDebugConfig,
SessionChainingConfig,
InstanceGroup,
HubAccessConfig,
ModelAccessConfig,
MetricDefinition,
DatasetSource,
)
from sagemaker.core.training.utils import convert_unassigned_to_none
__all__ = [
"BaseConfig",
"SourceCode",
"StoppingCondition",
"RetryStrategy",
"OutputDataConfig",
"Channel",
"ShuffleConfig",
"DataSource",
"S3DataSource",
"FileSystemDataSource",
"TrainingImageConfig",
"TrainingRepositoryAuthConfig",
"Tag",
"InfraCheckConfig",
"RemoteDebugConfig",
"SessionChainingConfig",
"InstanceGroup",
"TensorBoardOutputConfig",
"CheckpointConfig",
"HubAccessConfig",
"ModelAccessConfig",
"Compute",
"Networking",
"InputData",
"MetricDefinition",
"DatasetSource",
]
[docs]
class BaseConfig(BaseModel):
"""BaseConfig"""
model_config = ConfigDict(validate_assignment=True, extra="forbid")
[docs]
class SourceCode(BaseConfig):
"""SourceCode.
The SourceCode class allows the user to specify the source code location, dependencies,
entry script, or commands to be executed in the training job container.
Parameters:
source_dir (Optional[StrPipeVar]):
The local directory, s3 uri, or path to tar.gz file stored locally or in s3 that
contains the source code to be used in the training job container.
requirements (Optional[StrPipeVar]):
The path within ``source_dir`` to a ``requirements.txt`` file. If specified, the listed
requirements will be installed in the training job container.
entry_script (Optional[StrPipeVar]):
The path within ``source_dir`` to the entry script that will be executed in the training
job container. If not specified, command must be provided.
command (Optional[StrPipeVar]):
The command(s) to execute in the training job container. Example: "python my_script.py".
If not specified, entry_script must be provided.
ignore_patterns: (Optional[List[str]]) :
The ignore patterns to ignore specific files/folders when uploading to S3. If not specified,
default to: ['.env', '.git', '__pycache__', '.DS_Store', '.cache', '.ipynb_checkpoints'].
"""
source_dir: Optional[StrPipeVar] = None
requirements: Optional[StrPipeVar] = None
entry_script: Optional[StrPipeVar] = None
command: Optional[StrPipeVar] = None
ignore_patterns: Optional[List[str]] = [
".env",
".git",
"__pycache__",
".DS_Store",
".cache",
".ipynb_checkpoints",
]
class OutputDataConfig(shapes.OutputDataConfig):
"""OutputDataConfig.
Provides the configuration for the output data location of the training job
(will not be carried over to any model repository or deployment).
Parameters:
s3_output_path (Optional[StrPipeVar]):
The S3 URI where the output data will be stored. This is the location where the
training job will save its output data, such as model artifacts and logs.
kms_key_id (Optional[StrPipeVar]):
The Amazon Web Services Key Management Service (Amazon Web Services KMS) key that
SageMaker uses to encrypt the model artifacts at rest using Amazon S3 server-side
encryption.
compression_type (Optional[StrPipeVar]):
The model output compression type. Select None to output an uncompressed model,
recommended for large model outputs. Defaults to gzip.
"""
s3_output_path: Optional[StrPipeVar] = None
kms_key_id: Optional[StrPipeVar] = None
compression_type: Optional[StrPipeVar] = None
[docs]
class Compute(shapes.ResourceConfig):
"""Compute.
The Compute class is a subclass of ``sagemaker.core.shapes.ResourceConfig``
and allows the user to specify the compute resources for the training job.
Parameters:
instance_type (Optional[StrPipeVar]):
The ML compute instance type. For information about available instance types,
see https://aws.amazon.com/sagemaker/pricing/.
instance_count (Optional[IntPipeVar]): The number of ML compute instances to use. For distributed
training, provide a value greater than 1.
volume_size_in_gb (Optional[IntPipeVar]):
The size of the ML storage volume that you want to provision. ML storage volumes store
model artifacts and incremental states. Training algorithms might also use the ML
storage volume for scratch space. Default: 30
volume_kms_key_id (Optional[StrPipeVar]):
The Amazon Web Services KMS key that SageMaker uses to encrypt data on the storage
volume attached to the ML compute instance(s) that run the training job.
keep_alive_period_in_seconds (Optional[IntPipeVar]):
The duration of time in seconds to retain configured resources in a warm pool for
subsequent training jobs.
instance_groups (Optional[List[InstanceGroup]]):
A list of instance groups for heterogeneous clusters to be used in the training job.
training_plan_arn (Optional[StrPipeVar]):
The Amazon Resource Name (ARN) of the training plan to use for this resource configuration.
enable_managed_spot_training (Optional[BoolPipeVar]):
To train models using managed spot training, choose True. Managed spot training
provides a fully managed and scalable infrastructure for training machine learning
models. this option is useful when training jobs can be interrupted and when there
is flexibility when the training job is run.
"""
volume_size_in_gb: Optional[int] = 30
enable_managed_spot_training: Optional[bool] = None
@model_validator(mode="after")
def _model_validator(self) -> "Compute":
"""Convert Unassigned values to None."""
return convert_unassigned_to_none(self)
def _to_resource_config(self) -> shapes.ResourceConfig:
"""Convert to a sagemaker.core.shapes.ResourceConfig object."""
compute_config_dict = self.model_dump()
resource_config_fields = set(shapes.ResourceConfig.__annotations__.keys())
filtered_dict = {
k: v
for k, v in compute_config_dict.items()
if k in resource_config_fields and v is not None
}
if not filtered_dict:
return None
return shapes.ResourceConfig(**filtered_dict)
[docs]
class Networking(shapes.VpcConfig):
"""Networking.
The Networking class is a subclass of ``sagemaker.core.shapes.VpcConfig`` and
allows the user to specify the networking configuration for the training job.
Parameters:
security_group_ids (Optional[List[StrPipeVar]]):
The VPC security group IDs, in the form sg-xxxxxxxx. Specify the
security groups for the VPC that is specified in the Subnets field.
subnets (Optional[List[StrPipeVar]]):
The ID of the subnets in the VPC to which you want to connect your
training job or model.
enable_network_isolation (Optional[bool]):
Isolates the training container. No inbound or outbound network calls can be made,
except for calls between peers within a training cluster for distributed training.
If you enable network isolation for training jobs that are configured to use a VPC,
SageMaker downloads and uploads customer data and model artifacts through the
specified VPC, but the training container does not have network access.
enable_inter_container_traffic_encryption (Optional[bool]):
To encrypt all communications between ML compute instances in distributed training
choose True. Encryption provides greater security for distributed training, but
training might take longer. How long it takes depends on the amount of
communication between compute instances, especially if you use a deep learning
algorithm in distributed training.
"""
security_group_ids: Optional[list[StrPipeVar]] = None
subnets: Optional[list[StrPipeVar]] = None
enable_network_isolation: Optional[bool] = None
enable_inter_container_traffic_encryption: Optional[bool] = None
@model_validator(mode="after")
def _model_validator(self) -> "Networking":
"""Convert Unassigned values to None."""
return convert_unassigned_to_none(self)
def _to_vpc_config(self) -> shapes.VpcConfig:
"""Convert to a sagemaker.core.shapes.VpcConfig object."""
compute_config_dict = self.model_dump()
vpc_config_fields = set(shapes.VpcConfig.__annotations__.keys())
filtered_dict = {
k: v for k, v in compute_config_dict.items() if k in vpc_config_fields and v is not None
}
if not filtered_dict:
return None
return shapes.VpcConfig(**filtered_dict)
[docs]
class OutputDataConfig(shapes.OutputDataConfig):
"""OutputDataConfig.
The OutputDataConfig class is a subclass of ``sagemaker.core.shapes.OutputDataConfig``
and allows the user to specify the output data configuration for the training job
(will not be carried over to any model repository or deployment).
Parameters:
s3_output_path (Optional[StrPipeVar]):
The S3 URI where the output data will be stored. This is the location where the
training job will save its output data, such as model artifacts and logs.
kms_key_id (Optional[StrPipeVar]):
The Amazon Web Services Key Management Service (Amazon Web Services KMS) key that
SageMaker uses to encrypt the model artifacts at rest using Amazon S3 server-side
encryption.
compression_type (Optional[StrPipeVar]):
The model output compression type. Select `NONE` to output an uncompressed model,
recommended for large model outputs. Defaults to `GZIP`.
"""
s3_output_path: Optional[StrPipeVar] = None
kms_key_id: Optional[StrPipeVar] = None
compression_type: Optional[StrPipeVar] = None
[docs]
class TensorBoardOutputConfig(shapes.TensorBoardOutputConfig):
"""TensorBoardOutputConfig.
The TensorBoardOutputConfig class is a subclass of ``sagemaker.core.shapes.TensorBoardOutputConfig``
and allows the user to specify the storage locations for the Amazon SageMaker
Debugger TensorBoard.
Parameters:
s3_output_path (Optional[StrPipeVar]):
Path to Amazon S3 storage location for TensorBoard output. If not specified, will
default to
``s3://<default_bucket>/<default_prefix>/<base_job_name>/<job_name>/tensorboard-output``
local_path (Optional[StrPipeVar]):
Path to local storage location for tensorBoard output. Defaults to /opt/ml/output/tensorboard.
"""
s3_output_path: Optional[StrPipeVar] = None
local_path: Optional[StrPipeVar] = "/opt/ml/output/tensorboard"
[docs]
class CheckpointConfig(shapes.CheckpointConfig):
"""CheckpointConfig.
The CheckpointConfig class is a subclass of ``sagemaker.core.shapes.CheckpointConfig``
and allows the user to specify the checkpoint configuration for the training job.
Parameters:
s3_uri (Optional[StrPipeVar]):
Path to Amazon S3 storage location for the Checkpoint data. If not specified, will
default to
``s3://<default_bucket>/<default_prefix>/<base_job_name>/<job_name>/checkpoints``
local_path (Optional[StrPipeVar]):
The local directory where checkpoints are written. The default directory is /opt/ml/checkpoints.
"""
s3_uri: Optional[StrPipeVar] = None
local_path: Optional[StrPipeVar] = "/opt/ml/checkpoints"