sagemaker.train.base_trainer

sagemaker.train.base_trainer#

Classes

BaseTrainer([sagemaker_session, role, ...])

Abstract base class for all SageMaker training workflows.
class sagemaker.train.base_trainer.BaseTrainer(sagemaker_session: Session | None = None, role: str | None = None, base_job_name: str | None = None, tags: List[Tag] | None = None, hyperparameters: Dict[str, Any] | None = None, output_data_config: OutputDataConfig | None = None, input_data_config: List[Channel | InputData] | None = None, environment: Dict[str, str] | None = None, training_image: str | None = None, base_model_name: str | None = None, disable_output_compression: bool | None = False, notifications: Dict[str, Any] | None = None)[source]#

Bases: ABC

Abstract base class for all SageMaker training workflows.

This class provides the common interface and shared functionality for all trainer implementations including SFT, DPO, RLVR, and RLAIF trainers. It defines the standard parameters and abstract methods that concrete trainer classes must implement.

Parameters:

  • sagemaker_session (Optional[Session]) – The SageMaker session for managing API calls and resources. If not specified, a default session will be created.

  • role (Optional[str]) – The IAM role ARN for the training job execution. If not specified, the default SageMaker execution role will be used.

  • base_job_name (Optional[str]) – The base name for training jobs. A unique suffix will be appended. If not specified, a default name will be generated based on the trainer type.

  • tags (Optional[List[Tag]]) – List of tags to apply to the training job for resource management and billing.

  • hyperparameters (Optional[Dict[str, Any]]) – Dictionary of hyperparameters for the training job. Trainer-specific defaults will be applied if not specified.

  • output_data_config (Optional[shapes.OutputDataConfig]) – Configuration for training job outputs including S3 paths and encryption. If not specified, default output configuration will be used.

  • input_data_config (Optional[List[Union[Channel, InputData]]]) – List of input data channels for the training job. Can include training and validation datasets.

  • environment (Optional[Dict[str, str]]) – Environment variables to set in the training container.

  • training_image (Optional[str]) – Custom training container image URI. If not provided, the image is auto-resolved from the model’s recipe metadata in SageMaker Hub.

  • notifications (Optional[Dict[str, Any]]) – Configuration for SNS notifications on job status changes. Requires ‘sns_topic_arn’. Optional keys: ‘events’ [“Completed”, “Failed”, “Stopped”], ‘event_bus_arn’, and ‘job_name_prefix’. If not specified, no notifications are sent.

  • notification_rule_arn (str) – String of the EventBridge rule that is set up when enabling job notifications.

base_job_name: str | None = None#
delete_notification_rule(rule_arn: str, event_bus_arn: str | None = None) → str[source]#

Delete an SDK-created EventBridge notification rule.

Parameters:

  • rule_arn – The ARN of the rule to delete.

  • event_bus_arn – Optional EventBridge bus ARN. Defaults to “default”.

Returns:

The name of the deleted rule.

environment: Dict[str, str] | None = None#
get_resolved_recipe() → Dict[str, Any][source]#

Return the fully resolved recipe configuration.

Shows the final merged result of base defaults + user recipe + overrides after interpolation resolution and validation. Callable before or after train().

When neither recipe nor overrides were provided at construction time but hyperparameters have been set directly (e.g. trainer.hyperparameters.x = val), those user-set values are treated as implicit overrides so the resolved recipe still reflects the user’s intent.

Returns:

Deep copy of the resolved recipe configuration.

Return type:

dict

Raises:

ValueError – If no recipe, overrides, or direct hyperparameter assignments were provided.

hyperparameters: Dict[str, Any] | None = None#
input_data_config: List[Channel | InputData] | None = None#
latest_training_job: TrainingJob | None = None#
list_notification_rules(event_bus_arn: str | None = None) → List[Dict[str, str]][source]#

List all SDK-created EventBridge notification rules.

Returns:

List of dicts with ‘name’, ‘arn’, and ‘state’ for each rule.

classmethod list_supported_models(session=None) → List[str][source]#

Return the models that support this trainer’s fine-tuning technique.

Queries SageMakerPublicHub for all models whose RecipeCollection contains a FineTuning recipe for this trainer’s customization technique (cls._customization_technique, e.g. "SFT", "DPO", "RLVR", "RLAIF", "CPT").

Parameters:

session – Optional boto3 session.

Returns:

Sorted list of hub content model names supporting the technique.

output_data_config: OutputDataConfig | None = None#
role: str | None = None#
sagemaker_session: Session | None = None#
show_metrics(metrics: List[str] | None = None, starting_step: int | None = None, ending_step: int | None = None, start_time: Any | None = None, end_time: Any | None = None) → Any[source]#

Plot training metrics from CloudWatch logs (Nova) or MLflow (OSS).

For Nova models, parses CloudWatch logs for training_loss, lr, and reward_score. For non-Nova (OSS) models, pulls metrics from MLflow (requires mlflow_resource_arn to be configured on the trainer or auto-resolved).

Parameters:

  • metrics – Optional list of metric names to plot. If None, plots all available metrics for the training technique.

  • starting_step – Only plot metrics from this global step onwards.

  • ending_step – Only plot metrics up to this global step.

  • start_time – Optional start time for log retrieval. Accepts a datetime object or epoch milliseconds (int). When not provided, auto-resolved from the training job’s start time.

  • end_time – Optional end time for log retrieval. Accepts a datetime object or epoch milliseconds (int). When not provided, defaults to now.

Returns:

pandas.DataFrame containing the extracted metrics.

Raises:

  • NotImplementedError – If the training technique does not support metric extraction (e.g., DPO).

  • PermissionError – If CloudWatch logs cannot be read because the caller’s credentials are expired/invalid or lack CloudWatch Logs permissions.

  • ValueError – If no training job has been run yet, no logs/metrics are found, or MLflow is not configured for OSS models.

stream_logs(poll: int = 5, start_time: Any | None = None, tail_lines: int | None = None) → None[source]#

Stream CloudWatch logs in real-time (like kubectl logs -f).

Continuously polls for new log events and prints them as they arrive. Blocks until the training job reaches a terminal state (SMTJ) or the user interrupts with Ctrl+C (HyperPod).

Parameters:

  • poll – Polling interval in seconds between log fetches. Defaults to 5.

  • start_time – Optional start time to stream logs from. Accepts a datetime object or epoch milliseconds (int). Useful when attaching to a job that’s already running. If not provided, auto-resolved from the training job’s start time (SMTJ) or defaults to now (HyperPod).

  • tail_lines – Optional maximum number of most recent log lines to print. Logs are returned in chronological order; when specified, only the last tail_lines entries are shown (similar to kubectl logs --tail). Useful for quickly checking the latest output of long-running jobs without scrolling through the full history. If not provided, streams all logs until the job completes.

Raises:

ValueError – If no training job has been run yet.

tags: List[Tag] | None = None#
abstract train(input_data_config: List[InputData], wait: bool = True, logs: bool = True, wait_timeout: int | None = None, dry_run: bool = False)[source]#

Common training method that calls the specific implementation.

training_image: str | None = None#
Contents