SageMaker Train

SageMaker Train#

Training capabilities including model training, hyperparameter tuning, and distributed training.

Model Training#

SageMaker Python SDK Train Module.

Distributed Training#

Distributed module.

class sagemaker.train.distributed.DistributedConfig[source]#

Bases: BaseConfig, ABC

Abstract base class for distributed training configurations.

This class defines the interface that all distributed training configurations must implement. It provides a standardized way to specify driver scripts and their locations for distributed training jobs.

abstract property driver_dir: str#

Directory containing the driver script.

This property should return the path to the directory containing the driver script, relative to the container’s working directory.

Returns:: Path to directory containing the driver script
Return type:: str

abstract property driver_script: str#

Name of the driver script.

This property should return the name of the Python script that implements the distributed training driver logic.

Returns:: Name of the driver script file
Return type:: str

model_config: ClassVar[ConfigDict] = {'extra': 'forbid', 'validate_assignment': True}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class sagemaker.train.distributed.MPI(*, process_count_per_node: int | None = None, mpi_additional_options: List[str] | None = None)[source]#

Bases: DistributedConfig

MPI.

The MPI class configures a job that uses mpirun in the backend to launch distributed training.

Parameters:

process_count_per_node (int) – The number of processes to run on each node in the training job. Will default to the number of GPUs available in the container.
mpi_additional_options (Optional[str]) – The custom MPI options to use for the training job.

property driver_dir: str#

Directory containing the driver script.

Returns:: Path to directory containing the driver script
Return type:: str

property driver_script: str#

Name of the driver script.

Returns:: Name of the driver script
Return type:: str

model_config: ClassVar[ConfigDict] = {'extra': 'forbid', 'validate_assignment': True}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

mpi_additional_options: List[str] | None#

process_count_per_node: int | None#

class sagemaker.train.distributed.SMP(*, hybrid_shard_degree: int | None = None, sm_activation_offloading: bool | None = None, activation_loading_horizon: int | None = None, fsdp_cache_flush_warnings: bool | None = None, allow_empty_shards: bool | None = None, tensor_parallel_degree: int | None = None, context_parallel_degree: int | None = None, expert_parallel_degree: int | None = None, random_seed: int | None = None)[source]#

Bases: BaseConfig

SMP.

This class is used for configuring the SageMaker Model Parallelism v2 parameters. For more information on the model parallelism parameters, see: https://docs.aws.amazon.com/sagemaker/latest/dg/distributed-model-parallel-v2-reference.html#distributed-model-parallel-v2-reference-init-config

Parameters:

hybrid_shard_degree (Optional[int]) – Specifies a sharded parallelism degree for the model.
sm_activation_offloading (Optional[bool]) – Specifies whether to enable the SMP activation offloading implementation.
activation_loading_horizon (Optional[int]) – An integer specifying the activation offloading horizon type for FSDP. This is the maximum number of checkpointed or offloaded layers whose inputs can be in the GPU memory simultaneously.
fsdp_cache_flush_warnings (Optional[bool]) – Detects and warns if cache flushes happen in the PyTorch memory manager, because they can degrade computational performance.
allow_empty_shards (Optional[bool]) – Whether to allow empty shards when sharding tensors if tensor is not divisible. This is an experimental fix for crash during checkpointing in certain scenarios. Disabling this falls back to the original PyTorch behavior.
tensor_parallel_degree (Optional[int]) – Specifies a tensor parallelism degree. The value must be between 1 and world_size.
context_parallel_degree (Optional[int]) – Specifies the context parallelism degree. The value must be between 1 and world_size , and must be <= hybrid_shard_degree.
expert_parallel_degree (Optional[int]) – Specifies a expert parallelism degree. The value must be between 1 and world_size.
random_seed (Optional[int]) – A seed number for the random operations in distributed modules by SMP tensor parallelism or expert parallelism.

activation_loading_horizon: int | None#

allow_empty_shards: bool | None#

context_parallel_degree: int | None#

expert_parallel_degree: int | None#

fsdp_cache_flush_warnings: bool | None#

hybrid_shard_degree: int | None#

model_config: ClassVar[ConfigDict] = {'extra': 'forbid', 'validate_assignment': True}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

random_seed: int | None#

sm_activation_offloading: bool | None#

tensor_parallel_degree: int | None#

class sagemaker.train.distributed.Torchrun(*, process_count_per_node: int | None = None, smp: SMP | None = None)[source]#

Bases: DistributedConfig

Torchrun.

The Torchrun class configures a job that uses torchrun or torch.distributed.launch in the backend to launch distributed training.

Parameters:

process_count_per_node (int) – The number of processes to run on each node in the training job. Will default to the number of GPUs available in the container.
smp (Optional[SMP]) – The SageMaker Model Parallelism v2 parameters.

property driver_dir: str#

Directory containing the driver script.

Returns:: Path to directory containing the driver script
Return type:: str

property driver_script: str#

Name of the driver script.

Returns:: Name of the driver script file
Return type:: str

model_config: ClassVar[ConfigDict] = {'extra': 'forbid', 'validate_assignment': True}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

process_count_per_node: int | None#

smp: SMP | None#

Multi-Turn RL Training#

MultiTurnRLTrainer — trainer for Agentic Reinforcement Fine-Tuning (Multi-Turn RL) jobs.

Bases: BaseTrainer

Trainer for Agentic Reinforcement Fine-Tuning (Multi-Turn RL) jobs.

Uses CreateJob API (not CreateTrainingJob) with a JobConfigDocument JSON string.

Example

from sagemaker.train.multi_turn_rl_trainer import MultiTurnRLTrainer

trainer = MultiTurnRLTrainer(: model=”huggingface-reasoning-qwen3-32b”, agent_env=”arn:aws:bedrock-agentcore::us-west-2:123456789012:runtime/AGENTID”, training_dataset=”s3://my-bucket/”, output_model_package_group=”arn:aws:sagemaker:us-west-2:123456789012:model-package-group/grp”, mlflow_app_arn=”arn:aws:sagemaker:us-west-2:123456789012:mlflow-app/srv”, s3_output_path=”s3://my-bucket/output/”, accept_eula=True,

) job = trainer.train()

Parameters:

model – JumpStart model ID string or JumpStart hub content Model ARN.
agent_env – Bedrock AgentCore ARN, agent runtime ID, Lambda ARN, or CustomAgentLambda. When a bare agent runtime ID is provided (e.g. "myRuntime-aBcDeFgHiJ"), it is resolved to the full ARN via GetAgentRuntime.
training_dataset – S3 URI, DataSet object, or DataSet ARN string (optional). Must be provided at __init__ or train() time.
mlflow_app_arn – MLflow app ARN or MlflowApp object (optional). If not specified, uses the default MLflow experience.
s3_output_path – S3 path for output artifacts (optional). If not specified, defaults to s3://sagemaker-<region>-<account>/output.
output_model_package_group – ModelPackageGroup object or ARN string (optional).
intermediate_checkpoint_model_package_group – ModelPackageGroup object or ARN string for intermediate checkpoints (optional). If not provided, auto-creates {model_name}-mtrl-checkpoint-mpg. Must differ from output_model_package_group.
validation_dataset – S3 URI, DataSet object, or DataSet ARN string (optional).
bedrock_agentcore_qualifier – Bedrock AgentCore qualifier (default: "DEFAULT").
mlflow_experiment_name – MLflow experiment name (optional).
mlflow_run_name – MLflow run name (optional).
networking – VpcConfig for the job (optional).
kms_key_arn – KMS key ID for output encryption (optional).
accept_eula – Boolean for EULA acceptance (optional).
**kwargs – Passed to BaseTrainer (sagemaker_session, role, base_job_name, tags).
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.

classmethod attach(job_name: str, session=None) → AgentRFTJob[source]#

Attach to an existing Agentic RFT job by name.

Parameters:

job_name – The name of the job.
session – Optional boto3 session.

Returns:

AgentRFTJob wrapping the existing job.

static list_bedrock_agentcore_runtimes(session=None) → list[dict][source]#

List Bedrock AgentCore runtimes.

Parameters:: session – Optional boto3 session.
Returns:: List of dicts, each with keys name, runtime_id, arn, and status.

static list_supported_models(session=None) → list[str][source]#

Return the list of supported model names.

Queries SageMakerPublicHub to discover all models with MTRL recipes in their RecipeCollection.

Parameters:: session – Optional boto3 session.
Returns:: List of hub content model names supporting MTRL.

property output_model_package_arn: str | None#: The output model package ARN from the latest completed training job.

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

Stream CloudWatch logs for the latest MTRL training job.

Polls the correct log group (/aws/sagemaker/Job/AgentRFT) and checks job status via the Job API.

Parameters:

poll – Seconds between CloudWatch polling cycles (1-300).
start_time – Stream from this timestamp. Accepts datetime or epoch milliseconds (int). If None, streams from the beginning.

Raises:

ValueError – If no training job has been launched yet or poll is out of range.

train(training_dataset: str | DataSet | None = None, wait: bool = True, dry_run: bool = False) → AgentRFTJob | None[source]#

Launch an Agentic RFT job.

Parameters:

training_dataset – Training dataset override.
wait – If True (default), block until job reaches terminal status. Ignored under a PipelineSession, where there is no job to wait on – runnable_by_pipeline forces it to False and warns.
dry_run – If True, runs validation without submitting a job. Returns None on success. Ignored (with a warning) under a PipelineSession, where nothing is submitted either way.

Returns:

AgentRFTJob instance for tracking the job. None if dry_run=True. Under a PipelineSession the caller receives the _StepArguments runnable_by_pipeline returns: the assembled request is captured for JobStep composition instead of submitted.

AgentRFTJob — wrapper around sagemaker-core Job for AgentRFT job category.

class sagemaker.train.agent_rft_job.AgentRFTJob(job: Job)[source]#

Bases: object

Wrapper around sagemaker-core Job for AgentRFT job category.

Delegates lifecycle methods to the underlying Job and adds typed convenience properties by parsing the JobConfigDocument JSON string.

Parameters:: job – The sagemaker-core Job instance to wrap.

JOB_CATEGORY = 'AgentRFT'#

property agent_config: dict | None#: Full AgentConfig section from JobConfigDocument.

property billable_token_usage: dict | None#

Billable token usage from ServiceOutput.

Returns dict with keys: TrainTokenCount, PrefillTokenCount, SampleTokenCount.

property creation_time#: Return the job creation time.

delete()[source]#: Delete the job via DeleteJob API.

description: str | None#

property end_time#: Return the job end time.

property failure_reason: str | None#: Return the failure reason, if any.

classmethod from_job(job: Job) → AgentRFTJob[source]#: Create an AgentRFTJob from a sagemaker-core Job instance.

classmethod get(job_name: str, session=None) → AgentRFTJob[source]#

Attach to an existing AgentRFT job by name.

Parameters:

job_name – The name of the job.
session – Optional boto3 session.

Returns:

AgentRFTJob wrapping the existing job.

classmethod get_all(session=None, **kwargs)[source]#

List all AgentRFT jobs.

Delegates to Job.get_all with job_category pre-filled. Additional keyword arguments (e.g. creation_time_after, name_contains, sort_by, sort_order, status_equals) are forwarded.

Parameters:

session – Optional boto3 session.
**kwargs – Additional filter arguments forwarded to Job.get_all.

Yields:

AgentRFTJob instances.

get_mlflow_url() → str | None[source]#

Generate a fresh presigned MLflow URL for this job’s experiment/run.

In Jupyter notebooks, also renders a clickable link.

Returns:: Presigned URL string, or None if MLflow is not configured.

get_training_metrics() → list[dict][source]#

Fetch per-step MTRL training metrics from MLflow.

Retrieves rollout/reward/mean, rollout/turns/mean, training/total_tokens, and training/num_trajectories for each training step and prints a summary table.

Returns:: List of dicts, one per step, with keys step, rollout/reward/mean, rollout/turns/mean, training/total_tokens, and training/num_trajectories.

property job_arn: str#: Return the training job ARN.

property job_name: str#: Return the training job name.

property job_status: str#: Return the current job status.

property last_modified_time#: Return the job last-modified time.

property mlflow_details: dict | None#

Return MLflow experiment/run details from ServiceOutput.

Returns dict with keys: ExperimentName, RunName, ExperimentId, RunId.

property output_model_package_arn: str | None#: ARN of the output model package from ServiceOutput, or None.

property progress_info: dict | None#

Training progress from ServiceOutput.

Supports two formats: - Epoch-based: dict with MaxEpoch, StepsPerEpoch, CurrentEpoch, CurrentStep. - Step-only: dict with MaxSteps, CurrentStep.

Returns None if not available.

refresh()[source]#: Refresh job state from DescribeJob API.

property s3_output_path: str | None#: S3 output path from OutputDataConfig.

property secondary_status: str#: Return the current secondary status.

property secondary_status_transitions: list#: Return the list of secondary status transitions.

stop()[source]#: Stop the job via StopJob API.

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

Stream CloudWatch logs for this job in real-time.

Polls /aws/sagemaker/Job/AgentRFT and exits when the job reaches a terminal status or the user interrupts with Ctrl+C.

Parameters:

poll – Seconds between CloudWatch polling cycles (1-300).
start_time – Stream from this timestamp. Accepts datetime or epoch milliseconds (int). If None, streams from the beginning.

Raises:

ValueError – If poll is out of range.

property training_config: dict | None#: Full TrainingConfig section from JobConfigDocument.

wait(poll: int = 5, timeout: int | None = 3000, max_log_lines: int = 20)[source]#

Wait for job to reach terminal status.

Parameters:

poll – Seconds between polls.
timeout – Maximum seconds to wait.
max_log_lines – Maximum number of log lines to display. Defaults to 20.

wait_for_delete()[source]#: Wait for job deletion to complete.

CustomAgentLambda — Lambda-based agent environment for Agentic RFT.

class sagemaker.train.custom_agent_lambda.CustomAgentLambda(lambda_arn: str)[source]#

Bases: object

Lambda-based agent environment for Agentic RFT.

Creates and wraps Lambda functions that serve as agent environments or bridges between SageMaker and custom agent environments (e.g., LangSmith, EKS, Fargate).

Parameters:: lambda_arn – ARN of the Lambda function.

classmethod create(source: str, function_name: str | None = None, role: str | None = None, runtime: str = 'python3.12', handler: str = 'lambda_function.handler', timeout: int = 900, memory_size: int = 256, environment: dict | None = None, sagemaker_session=None) → CustomAgentLambda[source]#

Create a new Lambda function and return an CustomAgentLambda.

The source parameter accepts three formats:

S3 URI (s3://bucket/key.zip): deploys from an S3 artifact.
Local file path: reads the file, packages it as a zip, and uploads.
Inline code string: packages the raw code as a zip and uploads.

Detection order: S3 URI → existing local path → inline code.

Parameters:

source – S3 URI, local file path, or inline Python code string.
function_name – Lambda function name. If not provided, a unique name is generated automatically.
role – IAM role ARN for the Lambda execution role.
runtime – Lambda runtime (default: "python3.12").
handler – Lambda handler (default: "lambda_function.handler").
timeout – Lambda timeout in seconds (default: 900).
memory_size – Lambda memory in MB (default: 256).
environment – Dict of environment variables for the Lambda.
sagemaker_session – Optional SageMaker session for role resolution.

Returns:

CustomAgentLambda wrapping the created Lambda ARN.

Raises:

ValueError – If source is empty.

classmethod get(lambda_arn: str) → CustomAgentLambda[source]#

Wrap an existing Lambda ARN.

Validates the Lambda exists by calling GetFunction.

Parameters:: lambda_arn – ARN of an existing Lambda function.
Returns:: CustomAgentLambda wrapping the Lambda ARN.
Raises:: botocore.exceptions.ClientError – If the Lambda does not exist.

Interactive Training Sessions#

SageMaker training session SDK: client surface for interactive training.

Create a ServiceClient, derive a TrainingClient, run the training loop, then derive a SamplingClient for inference:

from sagemaker.train.training_session import (
    AdamParams,
    Datum,
    ModelInput,
    SamplingParams,
    ServiceClient,
    TensorData,
    TensorDtype,
)

svc = ServiceClient(
    region="us-west-2",
    role_arn="arn:aws:iam::123456789012:role/SageMakerTrainingSessionsRole",
    model_package_group_arn="arn:aws:sagemaker:us-west-2:123456789012:model-package-group/my-group",
    s3_output_path="s3://my-bucket/output",
)
tc = svc.create_lora_training_client(base_model=svc.list_base_models()[0])
for tokens in your_token_id_batches:
    # Every datum carries the loss inputs its loss function needs, one value
    # per token of the datum.
    batch = [
        Datum(
            model_input=ModelInput.from_ints(tokens),
            loss_fn_inputs={
                "target_tokens": TensorData(
                    data=tokens[1:] + [0], dtype=TensorDtype.INT64
                ),
                "weights": TensorData(data=[1.0] * len(tokens)),
            },
        )
    ]
    # Each call returns an APIFuture right away, so the optim_step is queued
    # behind the forward_backward without waiting for it to land.
    fwd_bwd = tc.forward_backward(batch, loss_fn="cross_entropy")
    optim = tc.optim_step(AdamParams(learning_rate=1e-4))
    print(fwd_bwd.result().metrics, optim.result().metrics)
sc = tc.save_weights_and_get_sampling_client()
sc.sample(
    "What is the capital of France?",
    num_samples=8,
    sampling_params=SamplingParams(max_tokens=512),
).result()

From async code, await the same futures instead of calling .result():

fwd_bwd = tc.forward_backward(batch, loss_fn="cross_entropy")
optim = tc.optim_step(AdamParams(learning_rate=1e-4))
fwd_bwd_result, optim_result = await asyncio.gather(fwd_bwd, optim)

class sagemaker.train.training_session.APIFuture(operation_id: str, poll_fn: ~typing.Callable[[], ~typing.Dict[str, ~typing.Any]], parse_fn: ~typing.Callable[[~typing.Dict[str, ~typing.Any]], ~sagemaker.train.training_session.futures._T], *, status_key: str = 'status', operation_type: str = 'operation', poll_interval: float = 2, time_fn: ~typing.Callable[[], float] = <built-in function monotonic>, sleep_fn: ~typing.Callable[[float], None] = <built-in function sleep>, async_sleep_fn: ~typing.Callable[[float], ~typing.Awaitable[None]] = <function sleep>, max_transient_poll_errors: int = 150, rand_fn: ~typing.Callable[[float, float], float] = <bound method Random.uniform of <random.Random object>>)[source]#

Bases: Generic[_T]

Future for an async runtime operation (training or sample).

Call result() to block until the operation reaches a terminal state. The same future type serves both training and sample operations; the caller supplies a poll_fn that fetches the latest response and a parse_fn that turns a successful response into the public result object.

Generic in the result type: parse_fn fixes it, so result() on an APIFuture[T] returns a T.

Parameters:

operation_id – Server-assigned id of the operation (for messages).
poll_fn – Zero-arg callable returning the latest operation response dict. Each call should hit Get*OperationResult for this operation.
parse_fn – Callable mapping a completed response dict to the result object returned by result().
status_key – Key under which the response carries its status string.
operation_type – Label used in error messages (e.g. "training").
poll_interval – Seconds between poll attempts.
time_fn – Monotonic clock, injectable for tests. Defaults to time.monotonic().
sleep_fn – Sleep function used by result(), injectable for tests. Defaults to time.sleep().
async_sleep_fn – Awaitable sleep used by result_async(), injectable for tests. Defaults to asyncio.sleep().
max_transient_poll_errors – Consecutive transient poll errors (throttling / 5xx on Get*OperationResult) tolerated before giving up. A successful poll resets the counter.
rand_fn – Jitter draw for the backoff between failed polls, injectable for tests. Defaults to random.uniform().

cancel() → bool[source]#

Always False: a submitted operation cannot be called back.

Present because the rest of this surface follows the standard future protocol, and a caller looping over futures should not have to special case this one. The service has the operation by the time a caller holds the future, so there is nothing to take back; stop the session to end the work it is doing.

cancelled() → bool[source]#: Always False: see cancel().

done() → bool[source]#

Whether this future has resolved, without waiting for it.

True once the operation’s outcome is known, whether that is a result or a failure, so a loop that drains futures by checking this terminates on a failed operation as well as a successful one. A future whose operation has finished on the service but has not been waited on yet still reports False: nothing has polled for it.

A failed request is not an outcome and is not recorded here. A call that could not be made or answered, credentials that expired mid-run being the clearest case, says nothing about the operation, which may still be running and may still succeed, so a later wait asks again rather than replaying an answer that was never about the operation.

exception(timeout: float | None = None) → BaseException | None[source]#

Wait for the operation and return the error it failed with.

Returns None if the operation succeeded, so a caller can inspect several futures without wrapping each in try.

A failed operation hands back the same error object on every wait, the way any settled future does, so identity comparisons hold. Re-raising one object repeatedly does lengthen its traceback each time, which is worth knowing if a retry loop waits on the same future many times over; read __cause__ and the message rather than the traceback depth.

Parameters:: timeout – Max seconds to wait before raising TimeoutError. None waits indefinitely.
Raises:: TimeoutError – The operation did not complete within timeout.

property operation_id: str#: Return the server-assigned operation id.

result(timeout: float | None = None) → _T[source]#

Block until the operation completes and return its parsed result.

Parameters:

timeout – Max seconds to wait before raising TimeoutError. None waits indefinitely.

Returns:

The object produced by parse_fn from the completed response.

Raises:

OperationError – The operation reached a terminal failure state.
TimeoutError – The operation did not complete within timeout.
Exception – The error botocore raised for a failed poll, either a transient one that recurred past max_transient_poll_errors in a row, or one that was never worth retrying.

async result_async(timeout: float | None = None) → _T[source]#

Await the operation result without occupying a thread while it runs.

Equivalent to await future, and the way to await with a timeout of your own: await future cannot take arguments, so it uses the default.

Parameters:

timeout – Max seconds to wait before raising TimeoutError. None waits indefinitely.

Raises:

OperationError – The operation reached a terminal failure state.
TimeoutError – The operation did not complete within timeout.
Exception – The error botocore raised for a failed poll, either a transient one that recurred past max_transient_poll_errors in a row, or one that was never worth retrying.

class sagemaker.train.training_session.AdamParams(*, learning_rate: float = 0.0001, beta1: float = 0.9, beta2: float = 0.999, eps: float = 1e-08, weight_decay: float = 0.01, grad_clip_norm: float | None = None, lora_plus_lr_ratio: float | None = None)[source]#

Bases: _ForbidExtraBase

Optimizer parameters for Adam.

Parameters:

learning_rate – Step size for parameter updates.
beta1 – Exponential decay rate for the first moment estimates.
beta2 – Exponential decay rate for the second moment estimates.
eps – Small constant for numerical stability.
weight_decay – L2 regularization coefficient.
grad_clip_norm – Gradient clipping norm.
lora_plus_lr_ratio – Learning-rate multiplier for LoRA+ B matrices relative to A matrices. Typical values are 8-32. When set, the service applies learning_rate * lora_plus_lr_ratio to B matrices while A matrices use learning_rate.

beta1: float#

beta2: float#

eps: float#

classmethod from_any(value: Any) → AdamParams[source]#

Coerce a params-like value to AdamParams.

Accepts an AdamParams (passthrough), a dict of fields, or any object sharing at least one field name. Raises TypeError for an object that shares none, instead of silently returning all defaults.

grad_clip_norm: float | None#

learning_rate: float#

lora_plus_lr_ratio: float | None#

model_config: ClassVar[ConfigDict] = {'extra': 'forbid', 'validate_assignment': True}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

weight_decay: float#

exception sagemaker.train.training_session.ArtifactDeletionError(model_package_arn: str, s3_uri: str, failed_keys: List[str])[source]#

Bases: TrainingSessionError

Raised when deleting a model package’s S3 artifacts did not fully succeed.

The model package is left in place, so it still points at the objects listed on failed_keys. Re-run ServiceClient.delete_model_package() once the cause is addressed (most often a missing S3 permission). The message truncates failed_keys; the attribute is complete.

class sagemaker.train.training_session.Datum(*, model_input: ModelInput | None = None, prompt: Any | None = None, loss_fn_inputs: Dict[str, Any] | None = None, **extra_data: Any)[source]#

Bases: _AllowExtraBase

A single training datum (model input + optional loss function inputs).

Use as the batch element for TrainingClient.forward_backward() and TrainingClient.forward().

Can be constructed from: - A ModelInput with optional loss_fn_inputs - Simple token IDs (via prompt) for convenience

Parameters:

model_input – Structured input (chunks-based). Mutually exclusive with prompt.
prompt – Convenience shorthand for token IDs or text. Converted to a ModelInput internally.
loss_fn_inputs – Per-datum loss function inputs, keyed by name: target_tokens, weights, advantages, logprobs, ref_log_prob, temperature. Which are required depends on the loss function. Values can be TensorData objects, tensors, or raw lists. Per-token entries must be as long as the token sequence.

classmethod from_any(value: Any) → Datum[source]#

Coerce a datum-like value to a Datum.

Accepts a Datum (passthrough), a ModelInput or a bare chunk, a token ID list or text (the prompt shorthand), a wire dict (modelInput/lossFnInputs or snake_case keys), or any object with model_input / prompt / loss_fn_inputs attributes.

Raises TypeError for a value carrying none of those. Every field on this type is optional, so an unrecognised value would otherwise become an empty Datum serializing to {}, dropping a training example from the batch with nothing raised. A mixed batch reaches this path: the element kind is decided by the first entry, and later entries still come through here one at a time.

loss_fn_inputs: Dict[str, Any] | None#

model_config: ClassVar[ConfigDict] = {'extra': 'allow', 'protected_namespaces': ()}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

model_input: ModelInput | None#

prompt: Any | None#

class sagemaker.train.training_session.EncodedTextChunk(*, tokens: List[int])[source]#

Bases: _ForbidExtraBase

A chunk of token IDs representing encoded text.

This is the fundamental building block for model inputs: a sequence of token IDs. Multiple chunks can be combined into a ModelInputChunk list to represent structured prompts.

Parameters:: tokens – List of token IDs.

classmethod from_any(value: Any) → EncodedTextChunk[source]#

Coerce a chunk-like value to an EncodedTextChunk.

Accepts an EncodedTextChunk (passthrough), a ModelInputChunk wrapper (or any object with the same attributes), a wire dict, a plain token ID list, or any object with tokens.

Raises:: TypeError – If the value carries no token IDs. Rejected instead of coerced to an empty chunk, which would send a prompt the model never sees and report success.

property length: int#: Number of tokens in this chunk.

model_config: ClassVar[ConfigDict] = {'extra': 'forbid'}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

tokens: List[int]#

class sagemaker.train.training_session.ForwardBackwardOutput[source]#

Bases: _OperationResult

Result of a completed forward_backward / forward operation.

property loss_fn_outputs: List[Dict[str, Any]] | None#

Return loss-function tensor outputs from the pass, one entry per datum.

Entries carry the loss outputs under actor/* keys (e.g. actor/pg_loss, actor/ppo_kl, actor/pg_clipfrac, actor/grad_norm).

property metrics: Dict[str, float]#

Return scalar training metrics for the pass.

Contains the throughput metrics (promised_flops, global_token_num, mfu, estimated_flops) plus scalar mirrors of the loss diagnostics (actor/pg_loss, actor/grad_norm, …); the tensor-form diagnostics are in loss_fn_outputs.

model_config: ClassVar[ConfigDict] = {'extra': 'forbid'}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

model_post_init(context: Any, /) → None#

This function is meant to behave like a BaseModel method to initialize private attributes.

It takes context as an argument since that’s what pydantic-core passes when calling it.

Parameters:

self – The BaseModel instance.
context – The context.

class sagemaker.train.training_session.LoadStateResponse[source]#

Bases: _OperationResult

Result of a completed load_state / load_state_with_optimizer operation.

Successful resolution of the future is the signal that the state was loaded; the service returns no payload fields for load operations.

model_config: ClassVar[ConfigDict] = {'extra': 'forbid'}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

model_post_init(context: Any, /) → None#

This function is meant to behave like a BaseModel method to initialize private attributes.

It takes context as an argument since that’s what pydantic-core passes when calling it.

Parameters:

self – The BaseModel instance.
context – The context.

class sagemaker.train.training_session.LossFunction(value)[source]#

Bases: str, Enum

Built-in loss functions for training operations.

Supported values: cross_entropy, importance_sampling, ppo, cispo.

CISPO = 'cispo'#

CROSS_ENTROPY = 'cross_entropy'#

IMPORTANCE_SAMPLING = 'importance_sampling'#

PPO = 'ppo'#

class sagemaker.train.training_session.ModelInput(*, chunks: List[Any] | None = None, prompt: Any | None = None, **extra_data: Any)[source]#

Bases: _AllowExtraBase

Structured model input composed of one or more chunks.

Use as the prompt argument to SamplingClient.sample() or to build typed training datums. Accepts either pre-built chunks or a simple prompt (token IDs or text string) via the convenience constructors from_ints() and from_text().

Parameters:: chunks – List of input chunks. Each chunk is either a ModelInputChunk, an EncodedTextChunk, or a raw dict matching the wire format.

append(chunk: Any) → ModelInput[source]#

Return a new ModelInput with chunk appended (does not mutate).

An input built from a prompt of token IDs carries those tokens into the chunks first, so appending extends it rather than replacing it.

Raises:: ValueError – This input still holds untokenized text. Appending built the result from chunks alone, so the text was dropped and the appended chunk became the whole input.

append_int(token: int) → ModelInput[source]#: Return a new ModelInput with a single-token chunk appended.

chunks: List[Any] | None#

classmethod empty() → ModelInput[source]#: Build a ModelInput with no chunks.

classmethod from_any(value: Any) → ModelInput[source]#

Coerce a prompt-like value to a ModelInput.

Accepts a ModelInput (passthrough), an EncodedTextChunk, a wire dict (chunks, prompt, or a single chunk’s tokens), a token ID list, or any object with chunks.

Raises:: ValueError – A dict naming none of those keys. Such a dict once became an input with no chunks, which the service accepts, so a training example was dropped without anything to see.

classmethod from_ints(tokens: List[int]) → ModelInput[source]#: Build a ModelInput from a flat list of token IDs.

classmethod from_text(text: str) → ModelInput[source]#

Build a ModelInput from text.

The text is tokenized when the input is submitted, using the session’s tokenizer, the same handling a bare string gets. Passing one to a client that cannot resolve a tokenizer raises there.

property length: int#

Total number of tokens across all chunks.

Raises:: ValueError – This input has no token IDs to count, for the reasons to_ints() gives.

model_config: ClassVar[ConfigDict] = {'extra': 'allow'}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

prompt: Any | None#

to_ints() → List[int][source]#

Flatten all chunks to a single token ID list.

Raises:: ValueError – This input carries a prompt that is not token IDs, so it has none to return yet. Returning an empty list instead would read a full prompt as empty wherever a token budget is counted.

class sagemaker.train.training_session.ModelInputChunk(*, chunk: EncodedTextChunk, **extra_data: Any)[source]#

Bases: _AllowExtraBase

A single chunk within a structured model input.

Wraps an EncodedTextChunk (the only chunk type today). This layer exists so the type system can accommodate future chunk types (e.g. image embeddings) without breaking existing code.

Parameters:: chunk – The encoded text chunk.

chunk: EncodedTextChunk#

model_config: ClassVar[ConfigDict] = {'extra': 'allow'}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class sagemaker.train.training_session.ModelPackageDeletion(*, model_package_arn: str, model_package_deleted: bool, s3_uri: str | None = None, objects_deleted: int = 0, artifacts_retained: bool = True, retained_reason: str | None = None)[source]#

Bases: _ForbidExtraBase

Outcome of ServiceClient.delete_model_package().

Parameters:

model_package_arn – ARN of the model package that was deleted.
model_package_deleted – True when this call deleted the package. False means it was already gone, so there was nothing left to delete.
s3_uri – The S3 location the model package pointed at, or None when the package exposes none. When the artifacts were kept, this is the only remaining record of where they are.
objects_deleted – Number of S3 objects deleted under s3_uri. Zero when the artifacts were kept, or when the location was already empty. On a versioned bucket a delete adds a marker instead of reclaiming storage, so this counts objects, not bytes freed.
artifacts_retained – True when this call did not attempt to delete the artifacts, either because delete_artifacts was False or because they are not in a location you manage. False with objects_deleted at zero means the location was already empty.
retained_reason – Human-readable reason the artifacts were kept, or None when they were deleted. Diagnostic text for a log; read artifacts_retained to branch on.

artifacts_retained: bool#

model_config: ClassVar[ConfigDict] = {'extra': 'forbid'}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

model_package_arn: str#

model_package_deleted: bool#

objects_deleted: int#

retained_reason: str | None#

s3_uri: str | None#

class sagemaker.train.training_session.ModelPackageInfo(*, model_package_arn: str, version: int | None = None, creation_time: datetime | None = None, description: str | None = None, kind: ModelPackageKind | None = None)[source]#

Bases: _ForbidExtraBase

Summary of one model package version in the session’s group.

A lightweight view (ARN, version, creation time, kind), distinct from the full sagemaker.core.resources.ModelPackage resource. Returned by ServiceClient.list_model_packages() and ServiceClient.get_latest_model_package(). Pass model_package_arn straight to TrainingClient.load_state() / TrainingClient.load_state_with_optimizer() (resume training), or to ServiceClient.create_sampling_client() (inference).

Parameters:

model_package_arn – ARN of the model package version.
version – Model package version number (monotonic within the group).
creation_time – When the model package was registered.
description – The model package description set when it was saved (a customer-provided string, or a platform default).
kind – Whether this is a CHECKPOINT (save_state) or MODEL (save_weights_for_sampler) package. None when not resolved: the field is only available from DescribeModelPackage, so ServiceClient.list_model_packages() leaves it unset (it lists without describing); the get_latest_* helpers always populate it.

creation_time: datetime | None#

description: str | None#

kind: ModelPackageKind | None#

model_config: ClassVar[ConfigDict] = {'extra': 'forbid'}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

model_package_arn: str#

version: int | None#

class sagemaker.train.training_session.ModelPackageKind(value)[source]#

Bases: str, Enum

Kind of model package a save operation registers.

A session registers one of two kinds of model package:

CHECKPOINT: a save_state checkpoint: sharded LoRA weights plus optimizer/FSDP state. Use it to resume an interrupted training run via TrainingClient.load_state_with_optimizer() (weights + optimizer) or TrainingClient.load_state() (weights only, optimizer resets).
MODEL: a save_weights_for_sampler model: the full (unsharded) LoRA weights. This is the deployable artifact (no optimizer state), so use it for downstream tasks like inference/sampling or to start a new training run.

CHECKPOINT = 'checkpoint'#

MODEL = 'model'#

exception sagemaker.train.training_session.OperationError(operation_type: str, message: str, *, operation_id: str | None = None)[source]#

Bases: TrainingSessionError

Raised when an asynchronous operation completed in a failed state.

The operation was accepted, so nothing was raised when it was submitted; it then failed while running, and that surfaces here when the future for it is resolved. A request the service rejects outright raises ClientError instead, at submit time.

operation_id identifies the failed operation server-side and is the one detail to quote when asking for help with it.

class sagemaker.train.training_session.OptimStepResponse[source]#

Bases: _OperationResult

Result of a completed optim_step operation.

property metrics: Dict[str, float]#: Return the optimizer-step metrics, such as the learning rate.

model_config: ClassVar[ConfigDict] = {'extra': 'forbid'}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

model_post_init(context: Any, /) → None#

This function is meant to behave like a BaseModel method to initialize private attributes.

It takes context as an argument since that’s what pydantic-core passes when calling it.

Parameters:

self – The BaseModel instance.
context – The context.

class sagemaker.train.training_session.RequestStatus(value)[source]#

Bases: str, Enum

Lifecycle states of an asynchronous training or sampling operation.

An operation moves from Pending to InProgress to a terminal state: Completed (success), Failed (the operation errored), or Skipped (the session failed and abandoned its remaining queued operations).

COMPLETED = 'Completed'#

FAILED = 'Failed'#

IN_PROGRESS = 'InProgress'#

PENDING = 'Pending'#

SKIPPED = 'Skipped'#

exception sagemaker.train.training_session.ResultMismatchError(operation_type: str, expected: int, received: int)[source]#

Bases: TrainingSessionError

Raised when a completed operation’s result does not match its request.

A per-token result must be as long as the sequence sent. A different length means the result belongs to some other request, so it is raised instead of returned: the values would otherwise be read as this request’s own.

class sagemaker.train.training_session.SampleResponse(*, sequences: List[SampledSequence], prompt_logprobs: List[float] | None = None, topk_prompt_logprobs: List[List[TokenLogprob]] | None = None)[source]#

Bases: _ForbidExtraBase

Result of a completed sampling operation.

Parameters:

sequences – Generated sequences.
prompt_logprobs – Log probabilities for prompt tokens (if requested).
topk_prompt_logprobs – Top-k log probabilities per prompt position (if requested).

model_config: ClassVar[ConfigDict] = {'extra': 'forbid'}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

prompt_logprobs: List[float] | None#

sequences: List[SampledSequence]#

topk_prompt_logprobs: List[List[TokenLogprob]] | None#

class sagemaker.train.training_session.SampledSequence(*, stop_reason: StopReason, tokens: List[int], logprobs: List[float] | None = None)[source]#

Bases: _ForbidExtraBase

A single generated sequence from a sampling operation.

Parameters:

stop_reason – Why generation stopped.
tokens – Generated token IDs.
logprobs – Per-token log probabilities (if requested).

logprobs: List[float] | None#

model_config: ClassVar[ConfigDict] = {'extra': 'forbid'}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

stop_reason: StopReason#

tokens: List[int]#

class sagemaker.train.training_session.SamplingClient(training_session_name: str, runtime: Any, *, model_package_arn: str = '', base_model: str = '', tokenizer: Any | None = None, hub_client: Any | None = None, hf_model_id: str | None = None, service_client: Any | None = None, session_handle: _SessionHandle | None = None, submitter: BackgroundSubmitter | None = None)[source]#

Bases: object

Samples from a checkpoint; model_package_arn="" = base model.

Parameters:

training_session_name – Name of the backing TrainingSession resource.
runtime – Client for the training-session runtime.
model_package_arn – Model package ARN of the checkpoint to sample from; "" = base model.
base_model – Base model (HF model ID or SageMaker Hub ARN), used to resolve a tokenizer via get_tokenizer().
tokenizer – Optional pre-loaded tokenizer for string prompts.
hub_client – Optional SageMaker Hub client, used to look up the HuggingFace model id behind a Hub ARN when a tokenizer is needed. Only consulted when neither tokenizer nor hf_model_id is set, and borrowed from service_client at that point if omitted here.
hf_model_id – Optional HuggingFace model id to load the tokenizer from. Set it when base_model is an ARN whose id cannot be looked up, or to pin the tokenizer to a specific repo.
service_client – The ServiceClient that created this client, used to stop the session when this client owns it.
session_handle – Existing _SessionHandle to share, passed by TrainingClient.create_sampling_client() so this client and its TrainingClient refcount one session. When omitted a fresh handle is created for training_session_name.
submitter – Existing BackgroundSubmitter to share, passed by TrainingClient.create_sampling_client() so refreshing a sampler does not add an event loop and thread each time. When omitted this client builds and owns its own.

Whether this client owns its session is read off session_handle rather than passed separately: sharing a handle means a TrainingClient is driving the session, so this client must not stop it; a fresh handle means the session was created implicitly for this client alone and has no other owner. Ownership governs stop() and __exit__ only; the automatic stop is the handle’s job either way. Submitter ownership is read the same way but tracked separately: a client can share a handle and still own its submitter.

compute_logprobs(prompt: Any) → APIFuture[List[float | None]][source]#

Submit a compute-logprobs operation; returns an APIFuture.

Unlike sample(), this generates no new tokens; it scores the given prompt and returns one log probability per prompt token. Useful for evaluation, reranking, or importance weighting. Non-blocking, like sample().

Parameters:: prompt – A list of token IDs, a string (requires a tokenizer), or a pre-built ModelInput dict.
Returns:: An APIFuture resolving to a list with one logprob per prompt token (the first entry is None, since it has no preceding context).
Raises:: ResultMismatchError – On resolution, when the result’s length does not match the sequence sent.

async compute_logprobs_async(prompt: Any) → List[float | None][source]#: Coroutine twin of compute_logprobs(); awaits the logprobs directly.

get_tokenizer() → Any[source]#

Return a HuggingFace tokenizer for the session’s base model.

The runtime accepts only token IDs, so callers with text must tokenize client-side before sample(). The tokenizer is loaded once and cached (and reused for string prompts passed to sample()).

Returns:

A transformers.PreTrainedTokenizer.

Raises:

ValueError – If the base model has no tokenizer available (e.g. a model whose tokenizer is not published to HuggingFace); pass pre-tokenized token IDs in that case.
ImportError – If transformers is not installed (install the optional extra: pip install sagemaker-train[tokenizer]).

sample(prompt: Any, num_samples: int = 1, sampling_params: SamplingParams | Dict[str, Any] | Any | None = None, *, include_prompt_logprobs: bool = False, topk_prompt_logprobs: int | None = None) → APIFuture[SampleResponse][source]#

Submit a sampling operation; returns an APIFuture.

Parameters:

prompt – A list of token IDs, a string (requires a tokenizer), or a pre-built ModelInput dict.
num_samples – Number of sequences to generate.
sampling_params – Generation parameters. Uses defaults if omitted. Accepts a SamplingParams, a dict of the same fields, or any object carrying at least one of them, so an existing config object can be passed as-is.
include_prompt_logprobs – Request per-token logprobs over the prompt.
topk_prompt_logprobs – If set, request the top-k logprobs per prompt position.

Returns:

An APIFuture resolving to a SampleResponse. Call .result() on it (or await it) to get the response. The submit itself runs on a background thread: this returns immediately, many requests can be queued concurrently, and a submit error surfaces when the future is resolved, not here.

async sample_async(prompt: Any, num_samples: int = 1, sampling_params: SamplingParams | Dict[str, Any] | Any | None = None, *, include_prompt_logprobs: bool = False, topk_prompt_logprobs: int | None = None) → SampleResponse[source]#

Coroutine twin of sample(); awaits the response directly.

resp = await sc.sample_async(...) resolves to the SampleResponse. Sample ops are order-independent, so unlike the training *_async methods there is no sequence position to hold and therefore no intermediate future to return.

stop() → None[source]#

Stop this client’s training session.

Only meaningful for a client created by ServiceClient.create_sampling_client(), which creates a session implicitly and therefore owns it. For a client created by TrainingClient.create_sampling_client() the session belongs to the TrainingClient, so this retires only this client; stop that session through its TrainingClient instead.

Drains the background submitter first (best effort) when this client owns it, so every operation already queued reaches the service before the stop request. The client accepts no new operations afterwards.

class sagemaker.train.training_session.SamplingParams(*, max_tokens: int = 512, temperature: Annotated[float | None, Ge(ge=0)] = None, top_p: Annotated[float | None, Ge(ge=0), Le(le=1)] = None, top_k: int | None = None, seed: int | None = None, stop: List[str] | None = None)[source]#

Bases: _ForbidExtraBase

Parameters for sampling / inference operations.

Parameters:

max_tokens – Maximum number of tokens to generate. The service requires a minimum of 512.
temperature – Sampling temperature. Must not be negative; 0 is greedy decoding. The service applies its own upper limit.
top_p – Nucleus sampling probability threshold, between 0 and 1. It is a cumulative probability, so a percentage such as 95 is rejected here rather than at the service.
top_k – Top-k sampling parameter. -1 disables it.
seed – Random seed for reproducibility.
stop – Stop sequences.

classmethod from_any(value: Any) → SamplingParams[source]#

Coerce a params-like value to SamplingParams.

Accepts a SamplingParams (passthrough), a dict of fields, or any object sharing at least one field name. Raises TypeError for an object that shares none, instead of silently returning all defaults.

max_tokens: int#

model_config: ClassVar[ConfigDict] = {'extra': 'forbid', 'validate_assignment': True}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

seed: int | None#

stop: List[str] | None#

temperature: float | None#

top_k: int | None#

top_p: float | None#

class sagemaker.train.training_session.SaveStateResponse[source]#

Bases: _OperationResult

Result of a completed save_state operation.

model_config: ClassVar[ConfigDict] = {'extra': 'forbid'}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

property model_package_arn: str | None#: Return the ARN of the model package the checkpoint was registered as.

model_post_init(context: Any, /) → None#

This function is meant to behave like a BaseModel method to initialize private attributes.

It takes context as an argument since that’s what pydantic-core passes when calling it.

Parameters:

self – The BaseModel instance.
context – The context.

class sagemaker.train.training_session.SaveWeightsResponse[source]#

Bases: _OperationResult

Result of a completed save_weights_for_sampler operation.

model_config: ClassVar[ConfigDict] = {'extra': 'forbid'}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

property model_package_arn: str | None#: Return the ARN of the model package the weights were registered as.

model_post_init(context: Any, /) → None#

This function is meant to behave like a BaseModel method to initialize private attributes.

It takes context as an argument since that’s what pydantic-core passes when calling it.

Parameters:

self – The BaseModel instance.
context – The context.

Bases: object

Entry point for creating and managing training sessions.

ServiceClient holds session-level defaults (role_arn, optional model_package_group_arn, s3_output_path, kms_key_arn, vpc_config) so that child-client creation needs only the model identity.

Anything left unset gets a SageMaker-conventional default:

Session name: {session_name_prefix}-{8 hex}, prefix defaulting to DEFAULT_SESSION_NAME_PREFIX. session_name_prefix names sessions only; it has no effect on the model package group or the S3 path.
Model package group: the account’s shared DEFAULT_MODEL_PACKAGE_GROUP_NAME, created on first use. Pass model_package_group_arn to keep one experiment’s checkpoints in their own group.
S3 output: s3://{sagemaker default bucket}/{session name}/output, resolved through the same Session.default_bucket() the broader SageMaker PySDK uses, so an existing sagemaker_config override applies.

Every client this ServiceClient creates cleans up after itself: a session is stopped once the last client referencing it becomes unreachable, so one is not left Running by a script that exits or a notebook cell that rebinds its client. The automatic path stops the session and leaves deletion to an explicit call, so a stopped session stays describable and a run that went wrong can still be inspected. Prefer a with block or an explicit .stop() where you can: the automatic stop is a backstop, not a substitute for scoping the session. Durability comes from a saved checkpoint, which create_sampling_client() takes directly.

Example:

from sagemaker.train.training_session import ServiceClient

svc = ServiceClient(
    region="us-west-2",
    role_arn="arn:aws:iam::123456789012:role/SageMakerTrainingSessionsRole",
    model_package_group_arn="arn:aws:sagemaker:us-west-2:123456789012:model-package-group/my-group",
    s3_output_path="s3://my-bucket/output",
)
# rank and alpha are keyword-only, defaulting to DEFAULT_LORA_RANK
# and DEFAULT_LORA_ALPHA.
tc = svc.create_lora_training_client(
    base_model="<id from list_base_models()>",
    rank=8,
    alpha=16,
)
sp = svc.create_sampling_client(base_model="<id from list_base_models()>")

Create a TrainingSession resource and return a TrainingClient.

Blocks until the session reaches Running before returning.

Every parameter after base_model is keyword-only, deliberately: create_lora_training_client(model, 16) does not say what 16 is, and a bare number at a call site is exactly the kind of thing that gets mis-read as alpha. Requiring rank=16 costs a caller nothing.

seed, train_mlp, train_attn, train_unembed and user_metadata are accepted for compatibility with existing training scripts but are not honoured; each one that is set logs a warning naming it.

Session-level parameters (model_package_group_arn, s3_output_path) default to the values set on this ServiceClient. Pass them here to override per-session.

Parameters:

base_model – A model id from list_base_models() ("provider/model"), or a full Hub model ARN to pin a specific model version. An id is resolved through the Hub to that model’s current version.
rank – LoRA rank (default 32).
alpha – LoRA alpha / scaling factor (default 64).
training_session_name – Name for the new session. When omitted, the SDK auto-generates {session_name_prefix}-{8 hex}.
model_package_group_arn – Model package group ARN for checkpoints. Falls back to ServiceClient’s default. When neither is set, the SDK uses the account’s shared default group (DEFAULT_MODEL_PACKAGE_GROUP_NAME), creating it on first use.
s3_output_path – S3 path for outputs. Falls back to ServiceClient’s default. When neither is set, the SDK uses the SageMaker default bucket convention.
kms_key_arn – ARN of an AWS KMS key used to encrypt the session’s outputs (saved weights and checkpoints) in Amazon S3. Falls back to ServiceClient’s default. When neither is set, outputs use Amazon S3 managed encryption.
vpc_config – VpcConfig (or a dict with subnets and security_group_ids) to run the session inside your VPC. Falls back to ServiceClient’s default.
seed – Accepted but not used; setting it logs a warning. LoRA initialisation is seeded by the service and is not caller-settable.
train_mlp – Accepted but not used; setting it logs a warning. LoRA is applied to the layer set the service selects.
train_attn – Accepted but not used; setting it logs a warning.
train_unembed – Accepted but not used; setting it logs a warning.
user_metadata – Accepted but not used; setting it logs a warning. Nothing is recorded against the session. Tag the model package group instead if you need run metadata to persist.

Returns:

A TrainingClient bound to the running session.

async create_lora_training_client_async(base_model: str, **kwargs: Any) → TrainingClient[source]#

Coroutine twin of create_lora_training_client().

Session creation waits for the session to reach Running, which can take minutes on a cold model. Running it on a worker thread keeps that wait off the event loop, so an asyncio caller can provision and prepare data concurrently.

base_model is the only positional parameter here too, and forwarding **kwargs is what keeps rank keyword-only through this path rather than quietly reintroducing a positional form the sync method rejects.

create_sampling_client(*, base_model: str | None = None, model_package_arn: str | None = None, retry_config: Any | None = None) → SamplingClient[source]#

Create a SamplingClient, implicitly creating a training session.

The SDK creates a training session under the hood and returns a SamplingClient bound to it. The session uses SDK-generated defaults for MPG and S3 output (or the values on this ServiceClient).

Because the session is implicit, the returned client owns it: use the client as a context manager (or call SamplingClient.stop()) so the session is stopped at a point you chose:

with svc.create_sampling_client(base_model="<id from list_base_models()>") as sp:
    out = sp.sample("hello").result()
# session stopped here

If you neither scope nor stop it, the SDK stops the session once the client becomes unreachable (on garbage collection or at interpreter exit) instead of leaving it Running until its inactivity timeout stops it. That is a backstop for interactive work, not a replacement for the with block above: it fires at a time the interpreter picks, and never at all if the process is killed with os._exit or SIGKILL.

Specify exactly one of base_model or model_package_arn:

base_model: sample from the base model directly.
model_package_arn: sample from saved weights (the SDK resolves the base model from the package metadata).

Parameters:

base_model – A model id from list_base_models(), or a full Hub model ARN.
model_package_arn – ARN of a model package (saved weights) to sample from.
retry_config – Accepted but not used; setting it logs a warning. Throttling and 5xx failures are retried on a fixed budget; any other error surfaces immediately.

Returns:

A SamplingClient bound to the implicit session.

Raises:

ValueError – If neither or both of base_model / model_package_arn are supplied.

async create_sampling_client_async(**kwargs: Any) → SamplingClient[source]#

Coroutine twin of create_sampling_client().

Creates the implicit session on a worker thread, so the wait for it to reach Running does not block the event loop.

create_training_client_from_state(model_package_arn: str, **kwargs: Any) → TrainingClient[source]#

Create a session from saved model weights.

Convenience over create_lora_training_client() + load_state: creates a fresh session (forwarding all the same keyword arguments) and loads model_package_arn into it, so the returned client resumes from that checkpoint.

This restores only the model weights, not optimizer state; the optimizer resets. To also restore optimizer state (e.g. Adam momentum), use create_training_client_from_state_with_optimizer(). The checkpoint carries both; which factory you call decides what is restored from it.

Parameters:

model_package_arn – ARN of the model package to load on startup.
**kwargs – Forwarded to create_lora_training_client(). When base_model is omitted, it is read from the checkpoint’s model package metadata (pinned to the version the weights were trained on).

async create_training_client_from_state_async(model_package_arn: str, **kwargs: Any) → TrainingClient[source]#: Coroutine twin of create_training_client_from_state().

create_training_client_from_state_with_optimizer(model_package_arn: str, **kwargs: Any) → TrainingClient[source]#: Like create_training_client_from_state(), restoring the optimizer.

async create_training_client_from_state_with_optimizer_async(model_package_arn: str, **kwargs: Any) → TrainingClient[source]#: Coroutine twin of create_training_client_from_state_with_optimizer().

delete_model_package(model_package_arn: str, *, delete_artifacts: bool = False) → ModelPackageDeletion[source]#

Delete a model package, optionally deleting its S3 artifacts too.

The model package carries the only reference to the S3 location its weights were written to, so by default this deletes the package and leaves the objects: they keep costing storage with nothing pointing at them. Keep the returned s3_uri if you intend to clean up later.

With delete_artifacts=True the objects go first and the package second, so a failure part-way leaves the package still pointing at whatever survived and the call can be retried.

Idempotent: a package that is already gone returns a result with model_package_deleted=False instead of raising.

Note

ServiceClient.create_sampling_client() and TrainingClient.load_state() read a package’s artifacts, so once they are deleted do not rely on this package for either.

Parameters:

model_package_arn – ARN of the model package version to delete.
delete_artifacts – When True, delete the S3 objects under the package’s location before deleting the package. Objects whose storage SageMaker manages on your behalf are left to it; deleting the package is what releases those.

Returns:

A ModelPackageDeletion describing what was deleted and, when the artifacts were kept, where they still are.

Raises:

ArtifactDeletionError – If some object could not be deleted. The model package is left in place, so a retry can find them again.
ValueError – If the package’s location names a bucket root, where deleting “its” artifacts would empty the whole bucket.

Example:

# Delete the package, keep the weights
result = svc.delete_model_package(ckpt.model_package_arn)
print("weights still at", result.s3_uri)

# Delete both
result = svc.delete_model_package(
    ckpt.model_package_arn, delete_artifacts=True
)
print("removed", result.objects_deleted, "objects")

delete_session(training_session: str, *, stop_if_active: bool = False) → None[source]#

Delete a training session.

The service only deletes sessions in a terminal state (Stopped, Failed, or DeletionFailed); deletion is asynchronous: the session flips to Deleting and a background worker removes it shortly after. Pass stop_if_active=True to stop a still-running session before deleting it.

Parameters:

training_session – Session name or ARN. A TrainingSessionArn from list_sessions() works as-is; the name is taken from the ARN’s last segment.
stop_if_active – When True, issue a stop and wait for a terminal state before deleting.

get_latest_model_package(kind: ModelPackageKind, *, model_package_group_arn: str | None = None, scan_limit: int = 20) → ModelPackageInfo | None[source]#

Return the latest model package of a given kind, or None.

Lists newest first and describes candidates in order, returning the first whose kind matches. Use ModelPackageKind.CHECKPOINT for resumable state (save_state) or ModelPackageKind.MODEL for deployable weights (save_weights_for_sampler).

Parameters:

kind – The kind of package to search for.
model_package_group_arn – The model package group. When omitted, uses the value set on this ServiceClient.
scan_limit – Most-recent packages to describe while searching (default DEFAULT_LATEST_SCAN_LIMIT). Exhausting it without a match logs a warning: None from a truncated search does not mean no such package exists.

list_base_models() → List[str][source]#

List the base models available for fine-tuning in this client’s region.

Queries the SageMaker Public Hub and returns the models this client can fine-tune end to end, which is narrower than every model the Hub tags as fine-tunable. Results are region-specific.

Each model is named by its own model id ("provider/model"), spelled as its provider publishes it wherever the Hub carries one. Pass one straight to create_lora_training_client() as base_model. The set changes as models are added, so call this rather than hard-coding an id.

list_model_packages(*, kind: ModelPackageKind | None = None, model_package_group_arn: str | None = None, max_results: int | None = 10) → List[ModelPackageInfo][source]#

List model packages, newest first.

Each ModelPackageInfo carries the ARN, version, creation time, and description. Pass its model_package_arn to TrainingClient.load_state() or ServiceClient.create_sampling_client().

When kind is supplied, the method describes each candidate to resolve its kind and returns only those matching the filter. Without kind, the kind field on each result is left unset (stays at one API call regardless of group size).

Parameters:

kind – Optional filter; only return packages of this kind.
model_package_group_arn – The model package group (name or ARN). When omitted, uses the value set on this ServiceClient.
max_results – Cap on packages returned (default 10, newest first). Pass None to return every version in the group.

list_sessions(*, name_contains: str | None = None, status_equals: str | SessionStatus | None = None, max_results: int | None = None) → list[TrainingSessionSummary][source]#

List training sessions, auto-paginating.

Parameters:

name_contains – Only return sessions whose name contains this string.
status_equals – Only return sessions in this status. Accepts a SessionStatus enum or a raw string (e.g. "Running").
max_results – Cap the total number of summaries returned. When None, every matching session is returned. Larger than one page, this pages until the cap is reached.

Returns:

A list of TrainingSessionSummary.

stop_session(training_session: str) → None[source]#

Stop a running training session.

Transitions the session toward a terminal Stopped state.

Idempotent: a session that is already stopping, already terminal, or being deleted is left as it is and the call returns normally. That makes stop order-independent, which is what the rest of the lifecycle relies on: an explicit stop_session here and the automatic stop that follows the last client’s collection both land on the same session, and neither has to know whether the other went first.

Parameters:: training_session – Session name or ARN. A TrainingSessionArn from list_sessions() works as-is; the name is taken from the ARN’s last segment.

class sagemaker.train.training_session.SessionStatus(value)[source]#

Bases: str, Enum

Status of a training session.

Use with ServiceClient.list_sessions() to filter by status without memorising the service’s raw PascalCase strings.

BLOCKED = 'Blocked'#

CREATED = 'Created'#

DELETING = 'Deleting'#

DELETION_FAILED = 'DeletionFailed'#

EXPIRED = 'Expired'#

FAILED = 'Failed'#

RUNNING = 'Running'#

SCHEDULED = 'Scheduled'#

STOPPED = 'Stopped'#

STOPPING = 'Stopping'#

class sagemaker.train.training_session.StopReason(value)[source]#

Bases: str, Enum

Reason why text generation stopped.

LENGTH = 'length'#

STOP = 'stop'#

class sagemaker.train.training_session.TensorData(data: Any = None, *, dtype: TensorDtype = TensorDtype.FLOAT32, shape: List[int] | None = None)[source]#

Bases: _ForbidExtraBase

A named tensor with dtype and shape metadata.

Use in loss_fn_inputs to pass typed tensor data (advantages, weights, logprobs, etc.) alongside training datums.

Tensor-like values (torch tensors, numpy arrays, anything with tolist()) are held as-is and only flattened to a plain list when the request is serialized or data is first read, so building datums from large tensors does not copy them up front.

Parameters:

data – Flat list of numeric values, or a tensor-like object converted lazily.
dtype – Tensor element type.
shape – Tensor dimensions (e.g. [batch_size, seq_len]).

property data: List[Any]#: Tensor values as a flat list (materialized from the source on first access).

dtype: TensorDtype#

classmethod from_any(value: Any) → TensorData[source]#

Coerce a tensor-like value to a TensorData.

Accepts a TensorData (passthrough), a wire-format dict, a torch tensor / numpy array (anything with tolist()), a plain list, or any object with data/dtype attributes.

model_config: ClassVar[ConfigDict] = {'extra': 'forbid'}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

model_post_init(context: Any, /) → None#

This function is meant to behave like a BaseModel method to initialize private attributes.

It takes context as an argument since that’s what pydantic-core passes when calling it.

Parameters:

self – The BaseModel instance.
context – The context.

shape: List[int] | None#

to_numpy() → Any[source]#: Convert to a numpy.ndarray (requires numpy).

to_torch() → Any[source]#: Convert to a torch.Tensor (requires torch).

tolist() → List[Any][source]#: Return the tensor values as a plain list.

class sagemaker.train.training_session.TensorDtype(value)[source]#

Bases: str, Enum

Data type for tensor values in loss function inputs.

Maps to the dtypes the runtime accepts in TensorData.dtype.

BFLOAT16 = 'bfloat16'#

FLOAT16 = 'float16'#

FLOAT32 = 'float32'#

INT32 = 'int32'#

INT64 = 'int64'#

class sagemaker.train.training_session.TokenLogprob(*, token_id: int, logprob: float)[source]#

Bases: _ForbidExtraBase

A token ID and its log probability.

logprob: float#

model_config: ClassVar[ConfigDict] = {'extra': 'forbid'}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

token_id: int#

class sagemaker.train.training_session.TrainingClient(training_session_name: str, runtime: Any, *, base_model: str = '', tokenizer: Any | None = None, hub_client: Any | None = None, hf_model_id: str | None = None, model_package_group_arn: str = '', service_client: Any | None = None)[source]#

Bases: object

Drives forward_backward / optim_step / save_weights_for_sampler over a session.

Instances are created by ServiceClient.create_lora_training_client() and its *_from_state variants, not directly by callers.

Parameters:

training_session_name – Name of the backing TrainingSession resource.
runtime – Client for the training-session runtime.
base_model – Base model (HF model ID or SageMaker Hub ARN), used to resolve a tokenizer via get_tokenizer().
tokenizer – Optional pre-loaded tokenizer for string batches.
hub_client – Optional SageMaker Hub client, used to look up the HuggingFace model id behind a Hub ARN when a tokenizer is needed. Only consulted when neither tokenizer nor hf_model_id is set.
hf_model_id – Optional HuggingFace model id to load the tokenizer from. Set it when base_model is an ARN whose id cannot be looked up, or to pin the tokenizer to a specific repo.
model_package_group_arn – The session’s model package group, where save_state / save_weights_for_sampler register checkpoints. Lets list_model_packages() resolve the group without the caller re-supplying it.
service_client – The ServiceClient that created this client, used to delegate checkpoint listing and to stop the session.

The backing session is stopped once the last client referencing it becomes unreachable, including any SamplingClient derived from this one, which shares this client’s session handle. The automatic teardown stops the session, leaving it describable for later inspection; deleting one stays an explicit delete() call.

Releasing the session is handled for you, so a caller can focus on the training loop. Anything worth keeping goes in a checkpoint: model packages are durable on their own, so trained weights outlive both the session and the process, and a new session can be created from one at any time.

create_sampling_client(model_package_arn: str = '', *, retry_config: Any | None = None) → SamplingClient[source]#

Create a SamplingClient; model_package_arn="" = base model.

Parameters:

model_package_arn – Model package ARN of the checkpoint to sample from. Empty string samples from the base model.
retry_config – Accepted but not used; setting it logs a warning. Throttling and 5xx failures are retried on a fixed budget; any other error surfaces immediately.

Returns:

A SamplingClient bound to this session and checkpoint.

async create_sampling_client_async(model_package_arn: str = '', *, retry_config: Any | None = None) → SamplingClient[source]#

Coroutine twin of create_sampling_client().

Binding to an existing session issues no control-plane call, so this is cheap; it exists so an asyncio caller can use one calling convention throughout without remembering which factories block.

property current_sequence_number: int#

The last sequence number assigned to an operation (read-only).

Safe to read for diagnostics: unlike the internal sequence allocator, this does NOT advance the counter, so inspecting it never consumes a sequence the service expects to see dispatched.

delete(*, stop_if_active: bool = False) → None[source]#

Delete this client’s training session.

Convenience over ServiceClient.delete_session() for this session. The service only deletes terminal sessions; pass stop_if_active=True to stop and wait for a terminal state first. Drains the background submitter first; the client accepts no new operations afterwards.

forward(batch: Any, loss_fn: LossFunction | str = 'cross_entropy', loss_fn_inputs: Dict[str, Any] | None = None, loss_fn_config: Dict[str, float] | None = None) → APIFuture[ForwardBackwardOutput][source]#

Submit a forward-only pass (no gradient); returns an APIFuture.

Same input shape as forward_backward(), but runs only the forward pass, useful for computing loss-fn outputs / logprobs without updating optimizer state. Non-blocking; the future resolves to a ForwardBackwardOutput with metrics and loss_fn_outputs.

Parameters:

batch – Training data.
loss_fn – Loss function name or LossFunction enum.
loss_fn_inputs – Optional loss inputs applied to each datum.
loss_fn_config – Optional loss function configuration (e.g. {"clip_low_threshold": 0.2, "clip_high_threshold": 5.0}).

Raises:

ValueError – A per-token lossFnInputs entry is not the same length as its datum’s modelInput, as for forward_backward().

A batch too large to send in one request is chunked as it is for forward_backward(), and the returned future resolves to a single combined ForwardBackwardOutput.

async forward_async(batch: Any, loss_fn: LossFunction | str = 'cross_entropy', loss_fn_inputs: Dict[str, Any] | None = None, loss_fn_config: Dict[str, float] | None = None) → APIFuture[ForwardBackwardOutput][source]#: Coroutine twin of forward() for async code.

forward_backward(batch: Any, loss_fn: LossFunction | str = 'cross_entropy', loss_fn_inputs: Dict[str, Any] | None = None, loss_fn_config: Dict[str, float] | None = None) → APIFuture[ForwardBackwardOutput][source]#

Submit a forward-backward pass; returns an APIFuture.

Non-blocking: the submit runs on a background thread, so this returns immediately. Resolve the future with .result() (or await it) for the ForwardBackwardOutput; a submit error surfaces there, not at the call site. Ops pipeline naturally: submit the next op before this future resolves; the sequence numbers stamped at call time preserve ordering on the service.

Parameters:

batch – Training data: a list of token ID lists, a single token ID list, a list of strings (requires a tokenizer), or raw datum dicts (advanced).
loss_fn – Loss function name or LossFunction enum.
loss_fn_inputs – Loss inputs applied to each datum that does not already carry its own. Which keys loss_fn requires is checked before the request leaves.
loss_fn_config – Optional loss function configuration (e.g. {"clip_low_threshold": 0.2, "clip_high_threshold": 5.0} for CISPO/PPO clipping).

Raises:

ValueError – A per-token lossFnInputs entry is not the same length as its datum’s modelInput. Raised before the request is sent, naming the datum and the key.

A batch too large to send in one request is split into chunks (bounded by datum count and estimated size) submitted as consecutive forwardBackward ops. The service accumulates gradients across consecutive forward-backward ops until the next optim_step, so this is equivalent to sending the whole batch at once. When the batch fits in one chunk (the common case) the request is unchanged. The returned future resolves to a single combined ForwardBackwardOutput.

async forward_backward_async(batch: Any, loss_fn: LossFunction | str = 'cross_entropy', loss_fn_inputs: Dict[str, Any] | None = None, loss_fn_config: Dict[str, float] | None = None) → APIFuture[ForwardBackwardOutput][source]#

Coroutine twin of forward_backward() for async code.

fut = await tc.forward_backward_async(...) completes the enqueue, so the op is sequenced before the next line runs, and await fut yields the ForwardBackwardOutput.

get_latest_model_package(kind: ModelPackageKind) → ModelPackageInfo | None[source]#

Return the latest model package of the given kind THIS client saved.

Answered purely from this client’s in-memory save cache, so no API call, and immune to the brief lag before a just-saved model package is describable. Returns None if this client has not saved a package of that kind in this process.

A fresh client starts empty: a package from a session this client did not drive is not implied here. Use ServiceClient.get_latest_model_package() or ServiceClient.list_model_packages() to discover packages in a model package group.

Parameters:: kind – The kind of model package to look up (e.g. ModelPackageKind.CHECKPOINT or ModelPackageKind.MODEL).
Returns:: A ModelPackageInfo with model_package_arn and kind, or None if nothing of that kind was saved by this client.

get_tokenizer() → Any[source]#

Return a HuggingFace tokenizer for the session’s base model.

The runtime accepts only token IDs, so callers with text must tokenize client-side: tok = tc.get_tokenizer(); ids = tok.encode("hello"). The tokenizer is loaded once and cached (and reused for string batches passed to forward_backward()).

Returns:

A transformers.PreTrainedTokenizer.

Raises:

ValueError – If the base model has no tokenizer available (e.g. a model whose tokenizer is not published to HuggingFace); pass pre-tokenized token IDs in that case.
ImportError – If transformers is not installed (install the optional extra: pip install sagemaker-train[tokenizer]).

list_model_packages(*, kind: ModelPackageKind | None = None, max_results: int | None = None) → List[ModelPackageInfo][source]#

List this session’s model packages, newest first.

Convenience over ServiceClient.list_model_packages() that supplies the session’s own model package group.

Parameters:

kind – Optional filter; only return packages of this kind (e.g. ModelPackageKind.CHECKPOINT or ModelPackageKind.MODEL).
max_results – Cap on packages returned. None uses the ServiceClient default (10 newest).

load_state(model_package_arn: str) → APIFuture[LoadStateResponse][source]#

Submit a load of a previously saved training state, weights only.

Restores the LoRA weights from a save_state() checkpoint and leaves the optimizer freshly initialized, so training continues from the saved weights with the optimizer reset. To also restore optimizer state (e.g. Adam momentum), use load_state_with_optimizer().

The checkpoint itself carries both weights and optimizer state; which method you call decides what is restored from it.

Non-blocking; the future resolves to a LoadStateResponse.

Parameters:: model_package_arn – ARN of the model package to load.

async load_state_async(model_package_arn: str) → APIFuture[LoadStateResponse][source]#: Coroutine twin of load_state() for async code.

load_state_with_optimizer(model_package_arn: str) → APIFuture[LoadStateResponse][source]#

Submit a load of saved training state, weights and optimizer.

Like load_state(), but also restores the optimizer state saved by save_state(), so training resumes exactly where it left off instead of restarting the optimizer from scratch.

Non-blocking; the future resolves to a LoadStateResponse.

Parameters:: model_package_arn – ARN of the model package to load.

async load_state_with_optimizer_async(model_package_arn: str) → APIFuture[LoadStateResponse][source]#: Coroutine twin of load_state_with_optimizer().

optim_step(adam_params: AdamParams | Dict[str, Any] | Any | None = None) → APIFuture[OptimStepResponse][source]#

Submit an optimizer step; returns an APIFuture.

Non-blocking; the future resolves to an OptimStepResponse with metrics. Pipelines with forward_backward(); the sequence numbers stamped at call time preserve ordering on the service.

Parameters:: adam_params – Optimizer parameters. Uses defaults if omitted. Accepts an AdamParams, a dict of the same fields, or any object carrying at least one of them, so an existing optimizer config object can be passed as-is.

async optim_step_async(adam_params: AdamParams | Dict[str, Any] | Any | None = None) → APIFuture[OptimStepResponse][source]#

Coroutine twin of optim_step() for async code.

fut = await tc.optim_step_async(...) completes the enqueue, guaranteeing the step is sequenced after any forward-backward already submitted, and await fut yields the OptimStepResponse.

save_state(description: str = '', *, ttl_seconds: int | None = None, overwrite: bool = False, name: str | None = None) → APIFuture[SaveStateResponse][source]#

Submit a save of full training state (weights + optimizer).

Non-blocking; the future resolves to a SaveStateResponse carrying model_package_arn, always a model package version that load_state() can resume from.

Parameters:

description – Optional human-readable description stored on the model package (e.g. "step 100 checkpoint").
ttl_seconds – Accepted but not used; setting it logs a warning. Checkpoint lifetime is governed by the platform’s retention policy, not per-save. Delete the model package version when you no longer need it.
overwrite – Accepted but not used; setting it logs a warning. Every save registers a new model package version, so a save never overwrites an earlier one, which is what False (the default) asks for, so it does not warn. True warns and is ignored.
name – Accepted but not used; setting it logs a warning. Checkpoints are model package versions addressed by ARN, so there is no caller-chosen name to set. Use description.

async save_state_async(description: str = '', *, ttl_seconds: int | None = None, overwrite: bool = False, name: str | None = None) → APIFuture[SaveStateResponse][source]#: Coroutine twin of save_state() for async code.

save_weights_and_get_sampling_client(description: str = '', merge_weights: bool = False, *, name: str | None = None, retry_config: Any | None = None) → SamplingClient[source]#

Save weights, load them for sampling, and return a SamplingClient.

Blocks until save_weights_for_sampler completes, points the sampling path at the freshly saved checkpoint, then returns a SamplingClient bound to it.

Parameters:

description – Optional human-readable description stored on the model package (passed through to save_weights_for_sampler()).
merge_weights – Whether to also write full adapter-free weights for standalone deployment (passed through to save_weights_for_sampler()).
name – Accepted but not used; setting it logs a warning. Checkpoints are model package versions addressed by ARN, so there is no caller-chosen name to set. Use description.
retry_config – Accepted but not used; setting it logs a warning. Throttling and 5xx failures are retried on a fixed budget; any other error surfaces immediately.

Returns:

A SamplingClient for sampling from the saved checkpoint.

Raises:

TrainingSessionError – The save completed without a checkpoint ARN. Raised rather than falling back: an empty ARN means the base model to SamplingClient, so the caller would have sampled untrained weights believing they were the trained ones.

async save_weights_and_get_sampling_client_async(description: str = '', merge_weights: bool = False, *, name: str | None = None, retry_config: Any | None = None) → SamplingClient[source]#

Coroutine twin of save_weights_and_get_sampling_client().

The save has to complete before the checkpoint can be sampled, so this blocks on a worker thread and not on the event loop, letting an RL loop keep servicing rollouts while the weights are written.

save_weights_for_sampler(description: str = '', merge_weights: bool = False, *, ttl_seconds: int | None = None, name: str | None = None) → APIFuture[SaveWeightsResponse][source]#

Submit a save-weights op; returns an APIFuture.

Non-blocking; the future resolves to a SaveWeightsResponse carrying model_package_arn. The op registers a model package version in the session’s model package group.

When merge_weights=True, the service merges the LoRA adapter into the base model and writes full, adapter-free weights (checkpoints/hf_merged/) alongside the adapter. Use this when the checkpoint will be deployed as a single standalone model, e.g. on GPUs without a working LoRA serving path.

Parameters:

description – Optional human-readable description stored on the model package (e.g. "weights at step 100").
merge_weights – Whether to also write full adapter-free weights for standalone deployment.
ttl_seconds – Accepted but not used; setting it logs a warning. Checkpoint lifetime is governed by the platform’s retention policy, not per-save. Delete the model package version when you no longer need it.
name – Accepted but not used; setting it logs a warning. Checkpoints are model package versions addressed by ARN, so there is no caller-chosen name to set. Use description.

async save_weights_for_sampler_async(description: str = '', merge_weights: bool = False, *, ttl_seconds: int | None = None, name: str | None = None) → APIFuture[SaveWeightsResponse][source]#: Coroutine twin of save_weights_for_sampler() for async code.

stop() → None[source]#

Stop this client’s training session.

Convenience over ServiceClient.stop_session() for this session, with no need to repeat the session name. Safe to call on an already-stopped session.

Drains the background submitter first (best effort), so every operation already queued reaches the service before the stop request; the client accepts no new operations afterwards.

exception sagemaker.train.training_session.TrainingSessionError[source]#

Bases: Exception

Base exception for training session errors.

Every error class in this module derives from this one, and it is also raised directly where the message says all there is to say: asking a closed client for another operation, for instance. A wait that runs out of time raises the built-in TimeoutError instead, so bound waits are caught the same way as with any other future.

class sagemaker.train.training_session.TrainingSessionSummary(*, training_session_arn: str, status: SessionStatus | str | None = None, creation_time: datetime | None = None, last_modified_time: datetime | None = None)[source]#

Bases: _ForbidExtraBase

Summary of one training session, as listed.

Returned by ServiceClient.list_sessions(). Carries the fields the list reports; anything else about a session comes from describing it.

Parameters:

training_session_arn – ARN of the session.
status – Current session status.
creation_time – When the session was created.
last_modified_time – When the session last changed state.

creation_time: datetime | None#

last_modified_time: datetime | None#

model_config: ClassVar[ConfigDict] = {'extra': 'forbid'}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

status: SessionStatus | str | None#: str is part of the type so a status this version does not know about still reaches the caller rather than failing the whole list call.

training_session_arn: str#

property training_session_name: str#: The session name, taken from the tail of training_session_arn.

class sagemaker.train.training_session.VpcConfig(*, subnets: List[str], security_group_ids: List[str])[source]#

Bases: _ForbidExtraBase

VPC configuration for a training session.

When set, the session’s compute runs inside your VPC: traffic to your Amazon S3 output location and to other resources flows through the subnets and security groups you specify, so make sure they allow the required access (for example an S3 gateway endpoint).

Parameters:

subnets – IDs of the subnets to run the session in.
security_group_ids – IDs of the security groups to apply.

classmethod from_any(value: Any) → VpcConfig[source]#

Coerce a VPC-config-like value to VpcConfig.

Accepts a VpcConfig (passthrough), a dict of fields, or any object with subnets / security_group_ids attributes.

model_config: ClassVar[ConfigDict] = {'extra': 'forbid'}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

security_group_ids: List[str]#

subnets: List[str]#

sagemaker.train.training_session.is_retryable(error: BaseException) → bool[source]#

Return whether sending the failed call again may succeed.

True for throttling, an HTTP 5xx, and connection-level failures such as a read or connect timeout: the request did not arrive, or arrived and could not be served for a reason unrelated to what it asked for. False for everything else, including a rejected request, since resending that one gets the same answer.

The SDK already retries its own internal calls on this basis. Use it to decide the same thing for a call of your own. Resolve the future inside the try: submitting is non-blocking, so the operation’s own failure surfaces where the result is awaited, not where it was queued:

future = training_client.forward_backward(batch, loss_fn="cross_entropy")
try:
    result = future.result()
except Exception as error:
    if not is_retryable(error):
        raise
    ...

Parameters:: error – The exception raised by the failed call.
Returns:: True if the failure is transient, False otherwise.

ServiceClient: entry point for creating and managing training sessions.

Holds session-level defaults (role, MPG, S3 output path) so that child-client creation is concise. Use ServiceClient.create_lora_training_client() to start a training session and get back a TrainingClient, or ServiceClient.create_sampling_client() to sample without training.

sagemaker.train.training_session.service_client.DEFAULT_MODEL_PACKAGE_GROUP_NAME = 'sagemaker-default'#: Name of the model package group used when the caller supplies none. One shared group per account/region, so kind-filtered discovery has to look past packages it did not write; see DEFAULT_LATEST_SCAN_LIMIT.

sagemaker.train.training_session.service_client.DEFAULT_SESSION_NAME_PREFIX = 'sagemaker-training-session'#: Prefix for auto-generated session names, when no session_name_prefix is set.

Bases: object

Entry point for creating and managing training sessions.

ServiceClient holds session-level defaults (role_arn, optional model_package_group_arn, s3_output_path, kms_key_arn, vpc_config) so that child-client creation needs only the model identity.

Anything left unset gets a SageMaker-conventional default:

Session name: {session_name_prefix}-{8 hex}, prefix defaulting to DEFAULT_SESSION_NAME_PREFIX. session_name_prefix names sessions only; it has no effect on the model package group or the S3 path.
Model package group: the account’s shared DEFAULT_MODEL_PACKAGE_GROUP_NAME, created on first use. Pass model_package_group_arn to keep one experiment’s checkpoints in their own group.
S3 output: s3://{sagemaker default bucket}/{session name}/output, resolved through the same Session.default_bucket() the broader SageMaker PySDK uses, so an existing sagemaker_config override applies.

Every client this ServiceClient creates cleans up after itself: a session is stopped once the last client referencing it becomes unreachable, so one is not left Running by a script that exits or a notebook cell that rebinds its client. The automatic path stops the session and leaves deletion to an explicit call, so a stopped session stays describable and a run that went wrong can still be inspected. Prefer a with block or an explicit .stop() where you can: the automatic stop is a backstop, not a substitute for scoping the session. Durability comes from a saved checkpoint, which create_sampling_client() takes directly.

Example:

from sagemaker.train.training_session import ServiceClient

svc = ServiceClient(
    region="us-west-2",
    role_arn="arn:aws:iam::123456789012:role/SageMakerTrainingSessionsRole",
    model_package_group_arn="arn:aws:sagemaker:us-west-2:123456789012:model-package-group/my-group",
    s3_output_path="s3://my-bucket/output",
)
# rank and alpha are keyword-only, defaulting to DEFAULT_LORA_RANK
# and DEFAULT_LORA_ALPHA.
tc = svc.create_lora_training_client(
    base_model="<id from list_base_models()>",
    rank=8,
    alpha=16,
)
sp = svc.create_sampling_client(base_model="<id from list_base_models()>")

Create a TrainingSession resource and return a TrainingClient.

Blocks until the session reaches Running before returning.

Every parameter after base_model is keyword-only, deliberately: create_lora_training_client(model, 16) does not say what 16 is, and a bare number at a call site is exactly the kind of thing that gets mis-read as alpha. Requiring rank=16 costs a caller nothing.

seed, train_mlp, train_attn, train_unembed and user_metadata are accepted for compatibility with existing training scripts but are not honoured; each one that is set logs a warning naming it.

Session-level parameters (model_package_group_arn, s3_output_path) default to the values set on this ServiceClient. Pass them here to override per-session.

Parameters:

base_model – A model id from list_base_models() ("provider/model"), or a full Hub model ARN to pin a specific model version. An id is resolved through the Hub to that model’s current version.
rank – LoRA rank (default 32).
alpha – LoRA alpha / scaling factor (default 64).
training_session_name – Name for the new session. When omitted, the SDK auto-generates {session_name_prefix}-{8 hex}.
model_package_group_arn – Model package group ARN for checkpoints. Falls back to ServiceClient’s default. When neither is set, the SDK uses the account’s shared default group (DEFAULT_MODEL_PACKAGE_GROUP_NAME), creating it on first use.
s3_output_path – S3 path for outputs. Falls back to ServiceClient’s default. When neither is set, the SDK uses the SageMaker default bucket convention.
kms_key_arn – ARN of an AWS KMS key used to encrypt the session’s outputs (saved weights and checkpoints) in Amazon S3. Falls back to ServiceClient’s default. When neither is set, outputs use Amazon S3 managed encryption.
vpc_config – VpcConfig (or a dict with subnets and security_group_ids) to run the session inside your VPC. Falls back to ServiceClient’s default.
seed – Accepted but not used; setting it logs a warning. LoRA initialisation is seeded by the service and is not caller-settable.
train_mlp – Accepted but not used; setting it logs a warning. LoRA is applied to the layer set the service selects.
train_attn – Accepted but not used; setting it logs a warning.
train_unembed – Accepted but not used; setting it logs a warning.
user_metadata – Accepted but not used; setting it logs a warning. Nothing is recorded against the session. Tag the model package group instead if you need run metadata to persist.

Returns:

A TrainingClient bound to the running session.

async create_lora_training_client_async(base_model: str, **kwargs: Any) → TrainingClient[source]#

Coroutine twin of create_lora_training_client().

Session creation waits for the session to reach Running, which can take minutes on a cold model. Running it on a worker thread keeps that wait off the event loop, so an asyncio caller can provision and prepare data concurrently.

base_model is the only positional parameter here too, and forwarding **kwargs is what keeps rank keyword-only through this path rather than quietly reintroducing a positional form the sync method rejects.

create_sampling_client(*, base_model: str | None = None, model_package_arn: str | None = None, retry_config: Any | None = None) → SamplingClient[source]#

Create a SamplingClient, implicitly creating a training session.

The SDK creates a training session under the hood and returns a SamplingClient bound to it. The session uses SDK-generated defaults for MPG and S3 output (or the values on this ServiceClient).

Because the session is implicit, the returned client owns it: use the client as a context manager (or call SamplingClient.stop()) so the session is stopped at a point you chose:

with svc.create_sampling_client(base_model="<id from list_base_models()>") as sp:
    out = sp.sample("hello").result()
# session stopped here

If you neither scope nor stop it, the SDK stops the session once the client becomes unreachable (on garbage collection or at interpreter exit) instead of leaving it Running until its inactivity timeout stops it. That is a backstop for interactive work, not a replacement for the with block above: it fires at a time the interpreter picks, and never at all if the process is killed with os._exit or SIGKILL.

Specify exactly one of base_model or model_package_arn:

base_model: sample from the base model directly.
model_package_arn: sample from saved weights (the SDK resolves the base model from the package metadata).

Parameters:

base_model – A model id from list_base_models(), or a full Hub model ARN.
model_package_arn – ARN of a model package (saved weights) to sample from.
retry_config – Accepted but not used; setting it logs a warning. Throttling and 5xx failures are retried on a fixed budget; any other error surfaces immediately.

Returns:

A SamplingClient bound to the implicit session.

Raises:

ValueError – If neither or both of base_model / model_package_arn are supplied.

async create_sampling_client_async(**kwargs: Any) → SamplingClient[source]#

Coroutine twin of create_sampling_client().

Creates the implicit session on a worker thread, so the wait for it to reach Running does not block the event loop.

create_training_client_from_state(model_package_arn: str, **kwargs: Any) → TrainingClient[source]#

Create a session from saved model weights.

Convenience over create_lora_training_client() + load_state: creates a fresh session (forwarding all the same keyword arguments) and loads model_package_arn into it, so the returned client resumes from that checkpoint.

This restores only the model weights, not optimizer state; the optimizer resets. To also restore optimizer state (e.g. Adam momentum), use create_training_client_from_state_with_optimizer(). The checkpoint carries both; which factory you call decides what is restored from it.

Parameters:

model_package_arn – ARN of the model package to load on startup.
**kwargs – Forwarded to create_lora_training_client(). When base_model is omitted, it is read from the checkpoint’s model package metadata (pinned to the version the weights were trained on).

async create_training_client_from_state_async(model_package_arn: str, **kwargs: Any) → TrainingClient[source]#: Coroutine twin of create_training_client_from_state().

create_training_client_from_state_with_optimizer(model_package_arn: str, **kwargs: Any) → TrainingClient[source]#: Like create_training_client_from_state(), restoring the optimizer.

async create_training_client_from_state_with_optimizer_async(model_package_arn: str, **kwargs: Any) → TrainingClient[source]#: Coroutine twin of create_training_client_from_state_with_optimizer().

delete_model_package(model_package_arn: str, *, delete_artifacts: bool = False) → ModelPackageDeletion[source]#

Delete a model package, optionally deleting its S3 artifacts too.

The model package carries the only reference to the S3 location its weights were written to, so by default this deletes the package and leaves the objects: they keep costing storage with nothing pointing at them. Keep the returned s3_uri if you intend to clean up later.

With delete_artifacts=True the objects go first and the package second, so a failure part-way leaves the package still pointing at whatever survived and the call can be retried.

Idempotent: a package that is already gone returns a result with model_package_deleted=False instead of raising.

Note

ServiceClient.create_sampling_client() and TrainingClient.load_state() read a package’s artifacts, so once they are deleted do not rely on this package for either.

Parameters:

model_package_arn – ARN of the model package version to delete.
delete_artifacts – When True, delete the S3 objects under the package’s location before deleting the package. Objects whose storage SageMaker manages on your behalf are left to it; deleting the package is what releases those.

Returns:

A ModelPackageDeletion describing what was deleted and, when the artifacts were kept, where they still are.

Raises:

ArtifactDeletionError – If some object could not be deleted. The model package is left in place, so a retry can find them again.
ValueError – If the package’s location names a bucket root, where deleting “its” artifacts would empty the whole bucket.

Example:

# Delete the package, keep the weights
result = svc.delete_model_package(ckpt.model_package_arn)
print("weights still at", result.s3_uri)

# Delete both
result = svc.delete_model_package(
    ckpt.model_package_arn, delete_artifacts=True
)
print("removed", result.objects_deleted, "objects")

delete_session(training_session: str, *, stop_if_active: bool = False) → None[source]#

Delete a training session.

The service only deletes sessions in a terminal state (Stopped, Failed, or DeletionFailed); deletion is asynchronous: the session flips to Deleting and a background worker removes it shortly after. Pass stop_if_active=True to stop a still-running session before deleting it.

Parameters:

training_session – Session name or ARN. A TrainingSessionArn from list_sessions() works as-is; the name is taken from the ARN’s last segment.
stop_if_active – When True, issue a stop and wait for a terminal state before deleting.

get_latest_model_package(kind: ModelPackageKind, *, model_package_group_arn: str | None = None, scan_limit: int = 20) → ModelPackageInfo | None[source]#

Return the latest model package of a given kind, or None.

Lists newest first and describes candidates in order, returning the first whose kind matches. Use ModelPackageKind.CHECKPOINT for resumable state (save_state) or ModelPackageKind.MODEL for deployable weights (save_weights_for_sampler).

Parameters:

kind – The kind of package to search for.
model_package_group_arn – The model package group. When omitted, uses the value set on this ServiceClient.
scan_limit – Most-recent packages to describe while searching (default DEFAULT_LATEST_SCAN_LIMIT). Exhausting it without a match logs a warning: None from a truncated search does not mean no such package exists.

list_base_models() → List[str][source]#

List the base models available for fine-tuning in this client’s region.

Queries the SageMaker Public Hub and returns the models this client can fine-tune end to end, which is narrower than every model the Hub tags as fine-tunable. Results are region-specific.

Each model is named by its own model id ("provider/model"), spelled as its provider publishes it wherever the Hub carries one. Pass one straight to create_lora_training_client() as base_model. The set changes as models are added, so call this rather than hard-coding an id.

list_model_packages(*, kind: ModelPackageKind | None = None, model_package_group_arn: str | None = None, max_results: int | None = 10) → List[ModelPackageInfo][source]#

List model packages, newest first.

Each ModelPackageInfo carries the ARN, version, creation time, and description. Pass its model_package_arn to TrainingClient.load_state() or ServiceClient.create_sampling_client().

When kind is supplied, the method describes each candidate to resolve its kind and returns only those matching the filter. Without kind, the kind field on each result is left unset (stays at one API call regardless of group size).

Parameters:

kind – Optional filter; only return packages of this kind.
model_package_group_arn – The model package group (name or ARN). When omitted, uses the value set on this ServiceClient.
max_results – Cap on packages returned (default 10, newest first). Pass None to return every version in the group.

list_sessions(*, name_contains: str | None = None, status_equals: str | SessionStatus | None = None, max_results: int | None = None) → list[TrainingSessionSummary][source]#

List training sessions, auto-paginating.

Parameters:

name_contains – Only return sessions whose name contains this string.
status_equals – Only return sessions in this status. Accepts a SessionStatus enum or a raw string (e.g. "Running").
max_results – Cap the total number of summaries returned. When None, every matching session is returned. Larger than one page, this pages until the cap is reached.

Returns:

A list of TrainingSessionSummary.

stop_session(training_session: str) → None[source]#

Stop a running training session.

Transitions the session toward a terminal Stopped state.

Idempotent: a session that is already stopping, already terminal, or being deleted is left as it is and the call returns normally. That makes stop order-independent, which is what the rest of the lifecycle relies on: an explicit stop_session here and the automatic stop that follows the last client’s collection both land on the same session, and neither has to know whether the other went first.

Parameters:: training_session – Session name or ARN. A TrainingSessionArn from list_sessions() works as-is; the name is taken from the ARN’s last segment.

TrainingClient: drives the training loop over a TrainingSession resource.

Owns the monotonic sequenceNumber for the session and turns each primitive (forward_backward / optim_step / save_weights_for_sampler / save_state / load_state) into an ExecuteTrainingOperation call, returning a unified APIFuture. Submits are non-blocking: each method stamps its sequence slot at call time, queues the wire call on a background submitter thread, and returns its future immediately, so a training loop can queue a large batch of operations in seconds and resolve results later. Operations can be pipelined: submit optim_step before the preceding forward_backward future resolves; the sequence numbers stamped at call time preserve order on the service side.

class sagemaker.train.training_session.training_client.TrainingClient(training_session_name: str, runtime: Any, *, base_model: str = '', tokenizer: Any | None = None, hub_client: Any | None = None, hf_model_id: str | None = None, model_package_group_arn: str = '', service_client: Any | None = None)[source]#

Bases: object

Drives forward_backward / optim_step / save_weights_for_sampler over a session.

Instances are created by ServiceClient.create_lora_training_client() and its *_from_state variants, not directly by callers.

Parameters:

training_session_name – Name of the backing TrainingSession resource.
runtime – Client for the training-session runtime.
base_model – Base model (HF model ID or SageMaker Hub ARN), used to resolve a tokenizer via get_tokenizer().
tokenizer – Optional pre-loaded tokenizer for string batches.
hub_client – Optional SageMaker Hub client, used to look up the HuggingFace model id behind a Hub ARN when a tokenizer is needed. Only consulted when neither tokenizer nor hf_model_id is set.
hf_model_id – Optional HuggingFace model id to load the tokenizer from. Set it when base_model is an ARN whose id cannot be looked up, or to pin the tokenizer to a specific repo.
model_package_group_arn – The session’s model package group, where save_state / save_weights_for_sampler register checkpoints. Lets list_model_packages() resolve the group without the caller re-supplying it.
service_client – The ServiceClient that created this client, used to delegate checkpoint listing and to stop the session.

The backing session is stopped once the last client referencing it becomes unreachable, including any SamplingClient derived from this one, which shares this client’s session handle. The automatic teardown stops the session, leaving it describable for later inspection; deleting one stays an explicit delete() call.

Releasing the session is handled for you, so a caller can focus on the training loop. Anything worth keeping goes in a checkpoint: model packages are durable on their own, so trained weights outlive both the session and the process, and a new session can be created from one at any time.

create_sampling_client(model_package_arn: str = '', *, retry_config: Any | None = None) → SamplingClient[source]#

Create a SamplingClient; model_package_arn="" = base model.

Parameters:

model_package_arn – Model package ARN of the checkpoint to sample from. Empty string samples from the base model.
retry_config – Accepted but not used; setting it logs a warning. Throttling and 5xx failures are retried on a fixed budget; any other error surfaces immediately.

Returns:

A SamplingClient bound to this session and checkpoint.

async create_sampling_client_async(model_package_arn: str = '', *, retry_config: Any | None = None) → SamplingClient[source]#

Coroutine twin of create_sampling_client().

Binding to an existing session issues no control-plane call, so this is cheap; it exists so an asyncio caller can use one calling convention throughout without remembering which factories block.

property current_sequence_number: int#

The last sequence number assigned to an operation (read-only).

Safe to read for diagnostics: unlike the internal sequence allocator, this does NOT advance the counter, so inspecting it never consumes a sequence the service expects to see dispatched.

delete(*, stop_if_active: bool = False) → None[source]#

Delete this client’s training session.

Convenience over ServiceClient.delete_session() for this session. The service only deletes terminal sessions; pass stop_if_active=True to stop and wait for a terminal state first. Drains the background submitter first; the client accepts no new operations afterwards.

forward(batch: Any, loss_fn: LossFunction | str = 'cross_entropy', loss_fn_inputs: Dict[str, Any] | None = None, loss_fn_config: Dict[str, float] | None = None) → APIFuture[ForwardBackwardOutput][source]#

Submit a forward-only pass (no gradient); returns an APIFuture.

Same input shape as forward_backward(), but runs only the forward pass, useful for computing loss-fn outputs / logprobs without updating optimizer state. Non-blocking; the future resolves to a ForwardBackwardOutput with metrics and loss_fn_outputs.

Parameters:

batch – Training data.
loss_fn – Loss function name or LossFunction enum.
loss_fn_inputs – Optional loss inputs applied to each datum.
loss_fn_config – Optional loss function configuration (e.g. {"clip_low_threshold": 0.2, "clip_high_threshold": 5.0}).

Raises:

ValueError – A per-token lossFnInputs entry is not the same length as its datum’s modelInput, as for forward_backward().

A batch too large to send in one request is chunked as it is for forward_backward(), and the returned future resolves to a single combined ForwardBackwardOutput.

async forward_async(batch: Any, loss_fn: LossFunction | str = 'cross_entropy', loss_fn_inputs: Dict[str, Any] | None = None, loss_fn_config: Dict[str, float] | None = None) → APIFuture[ForwardBackwardOutput][source]#: Coroutine twin of forward() for async code.

forward_backward(batch: Any, loss_fn: LossFunction | str = 'cross_entropy', loss_fn_inputs: Dict[str, Any] | None = None, loss_fn_config: Dict[str, float] | None = None) → APIFuture[ForwardBackwardOutput][source]#

Submit a forward-backward pass; returns an APIFuture.

Non-blocking: the submit runs on a background thread, so this returns immediately. Resolve the future with .result() (or await it) for the ForwardBackwardOutput; a submit error surfaces there, not at the call site. Ops pipeline naturally: submit the next op before this future resolves; the sequence numbers stamped at call time preserve ordering on the service.

Parameters:

batch – Training data: a list of token ID lists, a single token ID list, a list of strings (requires a tokenizer), or raw datum dicts (advanced).
loss_fn – Loss function name or LossFunction enum.
loss_fn_inputs – Loss inputs applied to each datum that does not already carry its own. Which keys loss_fn requires is checked before the request leaves.
loss_fn_config – Optional loss function configuration (e.g. {"clip_low_threshold": 0.2, "clip_high_threshold": 5.0} for CISPO/PPO clipping).

Raises:

ValueError – A per-token lossFnInputs entry is not the same length as its datum’s modelInput. Raised before the request is sent, naming the datum and the key.

A batch too large to send in one request is split into chunks (bounded by datum count and estimated size) submitted as consecutive forwardBackward ops. The service accumulates gradients across consecutive forward-backward ops until the next optim_step, so this is equivalent to sending the whole batch at once. When the batch fits in one chunk (the common case) the request is unchanged. The returned future resolves to a single combined ForwardBackwardOutput.

async forward_backward_async(batch: Any, loss_fn: LossFunction | str = 'cross_entropy', loss_fn_inputs: Dict[str, Any] | None = None, loss_fn_config: Dict[str, float] | None = None) → APIFuture[ForwardBackwardOutput][source]#

Coroutine twin of forward_backward() for async code.

fut = await tc.forward_backward_async(...) completes the enqueue, so the op is sequenced before the next line runs, and await fut yields the ForwardBackwardOutput.

get_latest_model_package(kind: ModelPackageKind) → ModelPackageInfo | None[source]#

Return the latest model package of the given kind THIS client saved.

Answered purely from this client’s in-memory save cache, so no API call, and immune to the brief lag before a just-saved model package is describable. Returns None if this client has not saved a package of that kind in this process.

A fresh client starts empty: a package from a session this client did not drive is not implied here. Use ServiceClient.get_latest_model_package() or ServiceClient.list_model_packages() to discover packages in a model package group.

Parameters:: kind – The kind of model package to look up (e.g. ModelPackageKind.CHECKPOINT or ModelPackageKind.MODEL).
Returns:: A ModelPackageInfo with model_package_arn and kind, or None if nothing of that kind was saved by this client.

get_tokenizer() → Any[source]#

Return a HuggingFace tokenizer for the session’s base model.

The runtime accepts only token IDs, so callers with text must tokenize client-side: tok = tc.get_tokenizer(); ids = tok.encode("hello"). The tokenizer is loaded once and cached (and reused for string batches passed to forward_backward()).

Returns:

A transformers.PreTrainedTokenizer.

Raises:

ValueError – If the base model has no tokenizer available (e.g. a model whose tokenizer is not published to HuggingFace); pass pre-tokenized token IDs in that case.
ImportError – If transformers is not installed (install the optional extra: pip install sagemaker-train[tokenizer]).

list_model_packages(*, kind: ModelPackageKind | None = None, max_results: int | None = None) → List[ModelPackageInfo][source]#

List this session’s model packages, newest first.

Convenience over ServiceClient.list_model_packages() that supplies the session’s own model package group.

Parameters:

kind – Optional filter; only return packages of this kind (e.g. ModelPackageKind.CHECKPOINT or ModelPackageKind.MODEL).
max_results – Cap on packages returned. None uses the ServiceClient default (10 newest).

load_state(model_package_arn: str) → APIFuture[LoadStateResponse][source]#

Submit a load of a previously saved training state, weights only.

Restores the LoRA weights from a save_state() checkpoint and leaves the optimizer freshly initialized, so training continues from the saved weights with the optimizer reset. To also restore optimizer state (e.g. Adam momentum), use load_state_with_optimizer().

The checkpoint itself carries both weights and optimizer state; which method you call decides what is restored from it.

Non-blocking; the future resolves to a LoadStateResponse.

Parameters:: model_package_arn – ARN of the model package to load.

async load_state_async(model_package_arn: str) → APIFuture[LoadStateResponse][source]#: Coroutine twin of load_state() for async code.

load_state_with_optimizer(model_package_arn: str) → APIFuture[LoadStateResponse][source]#

Submit a load of saved training state, weights and optimizer.

Like load_state(), but also restores the optimizer state saved by save_state(), so training resumes exactly where it left off instead of restarting the optimizer from scratch.

Non-blocking; the future resolves to a LoadStateResponse.

Parameters:: model_package_arn – ARN of the model package to load.

async load_state_with_optimizer_async(model_package_arn: str) → APIFuture[LoadStateResponse][source]#: Coroutine twin of load_state_with_optimizer().

optim_step(adam_params: AdamParams | Dict[str, Any] | Any | None = None) → APIFuture[OptimStepResponse][source]#

Submit an optimizer step; returns an APIFuture.

Non-blocking; the future resolves to an OptimStepResponse with metrics. Pipelines with forward_backward(); the sequence numbers stamped at call time preserve ordering on the service.

Parameters:: adam_params – Optimizer parameters. Uses defaults if omitted. Accepts an AdamParams, a dict of the same fields, or any object carrying at least one of them, so an existing optimizer config object can be passed as-is.

async optim_step_async(adam_params: AdamParams | Dict[str, Any] | Any | None = None) → APIFuture[OptimStepResponse][source]#

Coroutine twin of optim_step() for async code.

fut = await tc.optim_step_async(...) completes the enqueue, guaranteeing the step is sequenced after any forward-backward already submitted, and await fut yields the OptimStepResponse.

save_state(description: str = '', *, ttl_seconds: int | None = None, overwrite: bool = False, name: str | None = None) → APIFuture[SaveStateResponse][source]#

Submit a save of full training state (weights + optimizer).

Non-blocking; the future resolves to a SaveStateResponse carrying model_package_arn, always a model package version that load_state() can resume from.

Parameters:

description – Optional human-readable description stored on the model package (e.g. "step 100 checkpoint").
ttl_seconds – Accepted but not used; setting it logs a warning. Checkpoint lifetime is governed by the platform’s retention policy, not per-save. Delete the model package version when you no longer need it.
overwrite – Accepted but not used; setting it logs a warning. Every save registers a new model package version, so a save never overwrites an earlier one, which is what False (the default) asks for, so it does not warn. True warns and is ignored.
name – Accepted but not used; setting it logs a warning. Checkpoints are model package versions addressed by ARN, so there is no caller-chosen name to set. Use description.

async save_state_async(description: str = '', *, ttl_seconds: int | None = None, overwrite: bool = False, name: str | None = None) → APIFuture[SaveStateResponse][source]#: Coroutine twin of save_state() for async code.

save_weights_and_get_sampling_client(description: str = '', merge_weights: bool = False, *, name: str | None = None, retry_config: Any | None = None) → SamplingClient[source]#

Save weights, load them for sampling, and return a SamplingClient.

Blocks until save_weights_for_sampler completes, points the sampling path at the freshly saved checkpoint, then returns a SamplingClient bound to it.

Parameters:

description – Optional human-readable description stored on the model package (passed through to save_weights_for_sampler()).
merge_weights – Whether to also write full adapter-free weights for standalone deployment (passed through to save_weights_for_sampler()).
name – Accepted but not used; setting it logs a warning. Checkpoints are model package versions addressed by ARN, so there is no caller-chosen name to set. Use description.
retry_config – Accepted but not used; setting it logs a warning. Throttling and 5xx failures are retried on a fixed budget; any other error surfaces immediately.

Returns:

A SamplingClient for sampling from the saved checkpoint.

Raises:

TrainingSessionError – The save completed without a checkpoint ARN. Raised rather than falling back: an empty ARN means the base model to SamplingClient, so the caller would have sampled untrained weights believing they were the trained ones.

async save_weights_and_get_sampling_client_async(description: str = '', merge_weights: bool = False, *, name: str | None = None, retry_config: Any | None = None) → SamplingClient[source]#

Coroutine twin of save_weights_and_get_sampling_client().

The save has to complete before the checkpoint can be sampled, so this blocks on a worker thread and not on the event loop, letting an RL loop keep servicing rollouts while the weights are written.

save_weights_for_sampler(description: str = '', merge_weights: bool = False, *, ttl_seconds: int | None = None, name: str | None = None) → APIFuture[SaveWeightsResponse][source]#

Submit a save-weights op; returns an APIFuture.

Non-blocking; the future resolves to a SaveWeightsResponse carrying model_package_arn. The op registers a model package version in the session’s model package group.

When merge_weights=True, the service merges the LoRA adapter into the base model and writes full, adapter-free weights (checkpoints/hf_merged/) alongside the adapter. Use this when the checkpoint will be deployed as a single standalone model, e.g. on GPUs without a working LoRA serving path.

Parameters:

description – Optional human-readable description stored on the model package (e.g. "weights at step 100").
merge_weights – Whether to also write full adapter-free weights for standalone deployment.
ttl_seconds – Accepted but not used; setting it logs a warning. Checkpoint lifetime is governed by the platform’s retention policy, not per-save. Delete the model package version when you no longer need it.
name – Accepted but not used; setting it logs a warning. Checkpoints are model package versions addressed by ARN, so there is no caller-chosen name to set. Use description.

async save_weights_for_sampler_async(description: str = '', merge_weights: bool = False, *, ttl_seconds: int | None = None, name: str | None = None) → APIFuture[SaveWeightsResponse][source]#: Coroutine twin of save_weights_for_sampler() for async code.

stop() → None[source]#

Stop this client’s training session.

Convenience over ServiceClient.stop_session() for this session, with no need to repeat the session name. Safe to call on an already-stopped session.

Drains the background submitter first (best effort), so every operation already queued reaches the service before the stop request; the client accepts no new operations afterwards.

SamplingClient: sample from a saved checkpoint (or base model).

Created via TrainingClient.create_sampling_client(), TrainingClient.save_weights_and_get_sampling_client(), or ServiceClient.create_sampling_client() (sampling without training), not directly by callers. Only the last of those creates the backing session, and so only it stops that session when used as a context manager; the other two share the TrainingClient’s session handle, which is what marks them non-owning. Turns a prompt + SamplingParams into a sampling operation and returns an APIFuture over the SampleResponse. Submits are non-blocking: the wire call runs on a background submitter thread, and sample ops are order-independent, so many requests can be queued concurrently in seconds.

model_package_arn selects which checkpoint to sample from; "" means the base model. It is sent on the wire as modelPackageArn, or omitted entirely for the base model.

class sagemaker.train.training_session.sampling_client.SamplingClient(training_session_name: str, runtime: Any, *, model_package_arn: str = '', base_model: str = '', tokenizer: Any | None = None, hub_client: Any | None = None, hf_model_id: str | None = None, service_client: Any | None = None, session_handle: _SessionHandle | None = None, submitter: BackgroundSubmitter | None = None)[source]#

Bases: object

Samples from a checkpoint; model_package_arn="" = base model.

Parameters:

training_session_name – Name of the backing TrainingSession resource.
runtime – Client for the training-session runtime.
model_package_arn – Model package ARN of the checkpoint to sample from; "" = base model.
base_model – Base model (HF model ID or SageMaker Hub ARN), used to resolve a tokenizer via get_tokenizer().
tokenizer – Optional pre-loaded tokenizer for string prompts.
hub_client – Optional SageMaker Hub client, used to look up the HuggingFace model id behind a Hub ARN when a tokenizer is needed. Only consulted when neither tokenizer nor hf_model_id is set, and borrowed from service_client at that point if omitted here.
hf_model_id – Optional HuggingFace model id to load the tokenizer from. Set it when base_model is an ARN whose id cannot be looked up, or to pin the tokenizer to a specific repo.
service_client – The ServiceClient that created this client, used to stop the session when this client owns it.
session_handle – Existing _SessionHandle to share, passed by TrainingClient.create_sampling_client() so this client and its TrainingClient refcount one session. When omitted a fresh handle is created for training_session_name.
submitter – Existing BackgroundSubmitter to share, passed by TrainingClient.create_sampling_client() so refreshing a sampler does not add an event loop and thread each time. When omitted this client builds and owns its own.

Whether this client owns its session is read off session_handle rather than passed separately: sharing a handle means a TrainingClient is driving the session, so this client must not stop it; a fresh handle means the session was created implicitly for this client alone and has no other owner. Ownership governs stop() and __exit__ only; the automatic stop is the handle’s job either way. Submitter ownership is read the same way but tracked separately: a client can share a handle and still own its submitter.

compute_logprobs(prompt: Any) → APIFuture[List[float | None]][source]#

Submit a compute-logprobs operation; returns an APIFuture.

Unlike sample(), this generates no new tokens; it scores the given prompt and returns one log probability per prompt token. Useful for evaluation, reranking, or importance weighting. Non-blocking, like sample().

Parameters:: prompt – A list of token IDs, a string (requires a tokenizer), or a pre-built ModelInput dict.
Returns:: An APIFuture resolving to a list with one logprob per prompt token (the first entry is None, since it has no preceding context).
Raises:: ResultMismatchError – On resolution, when the result’s length does not match the sequence sent.

async compute_logprobs_async(prompt: Any) → List[float | None][source]#: Coroutine twin of compute_logprobs(); awaits the logprobs directly.

get_tokenizer() → Any[source]#

Return a HuggingFace tokenizer for the session’s base model.

The runtime accepts only token IDs, so callers with text must tokenize client-side before sample(). The tokenizer is loaded once and cached (and reused for string prompts passed to sample()).

Returns:

A transformers.PreTrainedTokenizer.

Raises:

ValueError – If the base model has no tokenizer available (e.g. a model whose tokenizer is not published to HuggingFace); pass pre-tokenized token IDs in that case.
ImportError – If transformers is not installed (install the optional extra: pip install sagemaker-train[tokenizer]).

sample(prompt: Any, num_samples: int = 1, sampling_params: SamplingParams | Dict[str, Any] | Any | None = None, *, include_prompt_logprobs: bool = False, topk_prompt_logprobs: int | None = None) → APIFuture[SampleResponse][source]#

Submit a sampling operation; returns an APIFuture.

Parameters:

prompt – A list of token IDs, a string (requires a tokenizer), or a pre-built ModelInput dict.
num_samples – Number of sequences to generate.
sampling_params – Generation parameters. Uses defaults if omitted. Accepts a SamplingParams, a dict of the same fields, or any object carrying at least one of them, so an existing config object can be passed as-is.
include_prompt_logprobs – Request per-token logprobs over the prompt.
topk_prompt_logprobs – If set, request the top-k logprobs per prompt position.

Returns:

An APIFuture resolving to a SampleResponse. Call .result() on it (or await it) to get the response. The submit itself runs on a background thread: this returns immediately, many requests can be queued concurrently, and a submit error surfaces when the future is resolved, not here.

async sample_async(prompt: Any, num_samples: int = 1, sampling_params: SamplingParams | Dict[str, Any] | Any | None = None, *, include_prompt_logprobs: bool = False, topk_prompt_logprobs: int | None = None) → SampleResponse[source]#

Coroutine twin of sample(); awaits the response directly.

resp = await sc.sample_async(...) resolves to the SampleResponse. Sample ops are order-independent, so unlike the training *_async methods there is no sequence position to hold and therefore no intermediate future to return.

stop() → None[source]#

Stop this client’s training session.

Only meaningful for a client created by ServiceClient.create_sampling_client(), which creates a session implicitly and therefore owns it. For a client created by TrainingClient.create_sampling_client() the session belongs to the TrainingClient, so this retires only this client; stop that session through its TrainingClient instead.

Drains the background submitter first (best effort) when this client owns it, so every operation already queued reaches the service before the stop request. The client accepts no new operations afterwards.

Configuration and result types for the training_session client interface.

class sagemaker.train.training_session.types.AdamParams(*, learning_rate: float = 0.0001, beta1: float = 0.9, beta2: float = 0.999, eps: float = 1e-08, weight_decay: float = 0.01, grad_clip_norm: float | None = None, lora_plus_lr_ratio: float | None = None)[source]#

Bases: _ForbidExtraBase

Optimizer parameters for Adam.

Parameters:

learning_rate – Step size for parameter updates.
beta1 – Exponential decay rate for the first moment estimates.
beta2 – Exponential decay rate for the second moment estimates.
eps – Small constant for numerical stability.
weight_decay – L2 regularization coefficient.
grad_clip_norm – Gradient clipping norm.
lora_plus_lr_ratio – Learning-rate multiplier for LoRA+ B matrices relative to A matrices. Typical values are 8-32. When set, the service applies learning_rate * lora_plus_lr_ratio to B matrices while A matrices use learning_rate.

beta1: float#

beta2: float#

eps: float#

classmethod from_any(value: Any) → AdamParams[source]#

Coerce a params-like value to AdamParams.

Accepts an AdamParams (passthrough), a dict of fields, or any object sharing at least one field name. Raises TypeError for an object that shares none, instead of silently returning all defaults.

grad_clip_norm: float | None#

learning_rate: float#

lora_plus_lr_ratio: float | None#

model_config: ClassVar[ConfigDict] = {'extra': 'forbid', 'validate_assignment': True}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

weight_decay: float#

class sagemaker.train.training_session.types.Datum(*, model_input: ModelInput | None = None, prompt: Any | None = None, loss_fn_inputs: Dict[str, Any] | None = None, **extra_data: Any)[source]#

Bases: _AllowExtraBase

A single training datum (model input + optional loss function inputs).

Use as the batch element for TrainingClient.forward_backward() and TrainingClient.forward().

Can be constructed from: - A ModelInput with optional loss_fn_inputs - Simple token IDs (via prompt) for convenience

Parameters:

model_input – Structured input (chunks-based). Mutually exclusive with prompt.
prompt – Convenience shorthand for token IDs or text. Converted to a ModelInput internally.
loss_fn_inputs – Per-datum loss function inputs, keyed by name: target_tokens, weights, advantages, logprobs, ref_log_prob, temperature. Which are required depends on the loss function. Values can be TensorData objects, tensors, or raw lists. Per-token entries must be as long as the token sequence.

classmethod from_any(value: Any) → Datum[source]#

Coerce a datum-like value to a Datum.

Accepts a Datum (passthrough), a ModelInput or a bare chunk, a token ID list or text (the prompt shorthand), a wire dict (modelInput/lossFnInputs or snake_case keys), or any object with model_input / prompt / loss_fn_inputs attributes.

Raises TypeError for a value carrying none of those. Every field on this type is optional, so an unrecognised value would otherwise become an empty Datum serializing to {}, dropping a training example from the batch with nothing raised. A mixed batch reaches this path: the element kind is decided by the first entry, and later entries still come through here one at a time.

loss_fn_inputs: Dict[str, Any] | None#

model_config: ClassVar[ConfigDict] = {'extra': 'allow', 'protected_namespaces': ()}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

model_input: ModelInput | None#

prompt: Any | None#

class sagemaker.train.training_session.types.EncodedTextChunk(*, tokens: List[int])[source]#

Bases: _ForbidExtraBase

A chunk of token IDs representing encoded text.

This is the fundamental building block for model inputs: a sequence of token IDs. Multiple chunks can be combined into a ModelInputChunk list to represent structured prompts.

Parameters:: tokens – List of token IDs.

classmethod from_any(value: Any) → EncodedTextChunk[source]#

Coerce a chunk-like value to an EncodedTextChunk.

Accepts an EncodedTextChunk (passthrough), a ModelInputChunk wrapper (or any object with the same attributes), a wire dict, a plain token ID list, or any object with tokens.

Raises:: TypeError – If the value carries no token IDs. Rejected instead of coerced to an empty chunk, which would send a prompt the model never sees and report success.

property length: int#: Number of tokens in this chunk.

model_config: ClassVar[ConfigDict] = {'extra': 'forbid'}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

tokens: List[int]#

class sagemaker.train.training_session.types.ForwardBackwardOutput[source]#

Bases: _OperationResult

Result of a completed forward_backward / forward operation.

property loss_fn_outputs: List[Dict[str, Any]] | None#

Return loss-function tensor outputs from the pass, one entry per datum.

Entries carry the loss outputs under actor/* keys (e.g. actor/pg_loss, actor/ppo_kl, actor/pg_clipfrac, actor/grad_norm).

property metrics: Dict[str, float]#

Return scalar training metrics for the pass.

Contains the throughput metrics (promised_flops, global_token_num, mfu, estimated_flops) plus scalar mirrors of the loss diagnostics (actor/pg_loss, actor/grad_norm, …); the tensor-form diagnostics are in loss_fn_outputs.

model_config: ClassVar[ConfigDict] = {'extra': 'forbid'}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

model_post_init(context: Any, /) → None#

This function is meant to behave like a BaseModel method to initialize private attributes.

It takes context as an argument since that’s what pydantic-core passes when calling it.

Parameters:

self – The BaseModel instance.
context – The context.

class sagemaker.train.training_session.types.LoadStateResponse[source]#

Bases: _OperationResult

Result of a completed load_state / load_state_with_optimizer operation.

Successful resolution of the future is the signal that the state was loaded; the service returns no payload fields for load operations.

model_config: ClassVar[ConfigDict] = {'extra': 'forbid'}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

model_post_init(context: Any, /) → None#

This function is meant to behave like a BaseModel method to initialize private attributes.

It takes context as an argument since that’s what pydantic-core passes when calling it.

Parameters:

self – The BaseModel instance.
context – The context.

class sagemaker.train.training_session.types.LossFunction(value)[source]#

Bases: str, Enum

Built-in loss functions for training operations.

Supported values: cross_entropy, importance_sampling, ppo, cispo.

CISPO = 'cispo'#

CROSS_ENTROPY = 'cross_entropy'#

IMPORTANCE_SAMPLING = 'importance_sampling'#

PPO = 'ppo'#

class sagemaker.train.training_session.types.ModelInput(*, chunks: List[Any] | None = None, prompt: Any | None = None, **extra_data: Any)[source]#

Bases: _AllowExtraBase

Structured model input composed of one or more chunks.

Use as the prompt argument to SamplingClient.sample() or to build typed training datums. Accepts either pre-built chunks or a simple prompt (token IDs or text string) via the convenience constructors from_ints() and from_text().

Parameters:: chunks – List of input chunks. Each chunk is either a ModelInputChunk, an EncodedTextChunk, or a raw dict matching the wire format.

append(chunk: Any) → ModelInput[source]#

Return a new ModelInput with chunk appended (does not mutate).

An input built from a prompt of token IDs carries those tokens into the chunks first, so appending extends it rather than replacing it.

Raises:: ValueError – This input still holds untokenized text. Appending built the result from chunks alone, so the text was dropped and the appended chunk became the whole input.

append_int(token: int) → ModelInput[source]#: Return a new ModelInput with a single-token chunk appended.

chunks: List[Any] | None#

classmethod empty() → ModelInput[source]#: Build a ModelInput with no chunks.

classmethod from_any(value: Any) → ModelInput[source]#

Coerce a prompt-like value to a ModelInput.

Accepts a ModelInput (passthrough), an EncodedTextChunk, a wire dict (chunks, prompt, or a single chunk’s tokens), a token ID list, or any object with chunks.

Raises:: ValueError – A dict naming none of those keys. Such a dict once became an input with no chunks, which the service accepts, so a training example was dropped without anything to see.

classmethod from_ints(tokens: List[int]) → ModelInput[source]#: Build a ModelInput from a flat list of token IDs.

classmethod from_text(text: str) → ModelInput[source]#

Build a ModelInput from text.

The text is tokenized when the input is submitted, using the session’s tokenizer, the same handling a bare string gets. Passing one to a client that cannot resolve a tokenizer raises there.

property length: int#

Total number of tokens across all chunks.

Raises:: ValueError – This input has no token IDs to count, for the reasons to_ints() gives.

model_config: ClassVar[ConfigDict] = {'extra': 'allow'}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

prompt: Any | None#

to_ints() → List[int][source]#

Flatten all chunks to a single token ID list.

Raises:: ValueError – This input carries a prompt that is not token IDs, so it has none to return yet. Returning an empty list instead would read a full prompt as empty wherever a token budget is counted.

class sagemaker.train.training_session.types.ModelInputChunk(*, chunk: EncodedTextChunk, **extra_data: Any)[source]#

Bases: _AllowExtraBase

A single chunk within a structured model input.

Wraps an EncodedTextChunk (the only chunk type today). This layer exists so the type system can accommodate future chunk types (e.g. image embeddings) without breaking existing code.

Parameters:: chunk – The encoded text chunk.

chunk: EncodedTextChunk#

model_config: ClassVar[ConfigDict] = {'extra': 'allow'}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class sagemaker.train.training_session.types.ModelPackageDeletion(*, model_package_arn: str, model_package_deleted: bool, s3_uri: str | None = None, objects_deleted: int = 0, artifacts_retained: bool = True, retained_reason: str | None = None)[source]#

Bases: _ForbidExtraBase

Outcome of ServiceClient.delete_model_package().

Parameters:

model_package_arn – ARN of the model package that was deleted.
model_package_deleted – True when this call deleted the package. False means it was already gone, so there was nothing left to delete.
s3_uri – The S3 location the model package pointed at, or None when the package exposes none. When the artifacts were kept, this is the only remaining record of where they are.
objects_deleted – Number of S3 objects deleted under s3_uri. Zero when the artifacts were kept, or when the location was already empty. On a versioned bucket a delete adds a marker instead of reclaiming storage, so this counts objects, not bytes freed.
artifacts_retained – True when this call did not attempt to delete the artifacts, either because delete_artifacts was False or because they are not in a location you manage. False with objects_deleted at zero means the location was already empty.
retained_reason – Human-readable reason the artifacts were kept, or None when they were deleted. Diagnostic text for a log; read artifacts_retained to branch on.

artifacts_retained: bool#

model_config: ClassVar[ConfigDict] = {'extra': 'forbid'}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

model_package_arn: str#

model_package_deleted: bool#

objects_deleted: int#

retained_reason: str | None#

s3_uri: str | None#

class sagemaker.train.training_session.types.ModelPackageInfo(*, model_package_arn: str, version: int | None = None, creation_time: datetime | None = None, description: str | None = None, kind: ModelPackageKind | None = None)[source]#

Bases: _ForbidExtraBase

Summary of one model package version in the session’s group.

A lightweight view (ARN, version, creation time, kind), distinct from the full sagemaker.core.resources.ModelPackage resource. Returned by ServiceClient.list_model_packages() and ServiceClient.get_latest_model_package(). Pass model_package_arn straight to TrainingClient.load_state() / TrainingClient.load_state_with_optimizer() (resume training), or to ServiceClient.create_sampling_client() (inference).

Parameters:

model_package_arn – ARN of the model package version.
version – Model package version number (monotonic within the group).
creation_time – When the model package was registered.
description – The model package description set when it was saved (a customer-provided string, or a platform default).
kind – Whether this is a CHECKPOINT (save_state) or MODEL (save_weights_for_sampler) package. None when not resolved: the field is only available from DescribeModelPackage, so ServiceClient.list_model_packages() leaves it unset (it lists without describing); the get_latest_* helpers always populate it.

creation_time: datetime | None#

description: str | None#

kind: ModelPackageKind | None#

model_config: ClassVar[ConfigDict] = {'extra': 'forbid'}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

model_package_arn: str#

version: int | None#

class sagemaker.train.training_session.types.ModelPackageKind(value)[source]#

Bases: str, Enum

Kind of model package a save operation registers.

A session registers one of two kinds of model package:

CHECKPOINT: a save_state checkpoint: sharded LoRA weights plus optimizer/FSDP state. Use it to resume an interrupted training run via TrainingClient.load_state_with_optimizer() (weights + optimizer) or TrainingClient.load_state() (weights only, optimizer resets).
MODEL: a save_weights_for_sampler model: the full (unsharded) LoRA weights. This is the deployable artifact (no optimizer state), so use it for downstream tasks like inference/sampling or to start a new training run.

CHECKPOINT = 'checkpoint'#

MODEL = 'model'#

class sagemaker.train.training_session.types.OptimStepResponse[source]#

Bases: _OperationResult

Result of a completed optim_step operation.

property metrics: Dict[str, float]#: Return the optimizer-step metrics, such as the learning rate.

model_config: ClassVar[ConfigDict] = {'extra': 'forbid'}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

model_post_init(context: Any, /) → None#

This function is meant to behave like a BaseModel method to initialize private attributes.

It takes context as an argument since that’s what pydantic-core passes when calling it.

Parameters:

self – The BaseModel instance.
context – The context.

class sagemaker.train.training_session.types.SampleResponse(*, sequences: List[SampledSequence], prompt_logprobs: List[float] | None = None, topk_prompt_logprobs: List[List[TokenLogprob]] | None = None)[source]#

Bases: _ForbidExtraBase

Result of a completed sampling operation.

Parameters:

sequences – Generated sequences.
prompt_logprobs – Log probabilities for prompt tokens (if requested).
topk_prompt_logprobs – Top-k log probabilities per prompt position (if requested).

model_config: ClassVar[ConfigDict] = {'extra': 'forbid'}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

prompt_logprobs: List[float] | None#

sequences: List[SampledSequence]#

topk_prompt_logprobs: List[List[TokenLogprob]] | None#

class sagemaker.train.training_session.types.SampledSequence(*, stop_reason: StopReason, tokens: List[int], logprobs: List[float] | None = None)[source]#

Bases: _ForbidExtraBase

A single generated sequence from a sampling operation.

Parameters:

stop_reason – Why generation stopped.
tokens – Generated token IDs.
logprobs – Per-token log probabilities (if requested).

logprobs: List[float] | None#

model_config: ClassVar[ConfigDict] = {'extra': 'forbid'}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

stop_reason: StopReason#

tokens: List[int]#

class sagemaker.train.training_session.types.SamplingParams(*, max_tokens: int = 512, temperature: Annotated[float | None, Ge(ge=0)] = None, top_p: Annotated[float | None, Ge(ge=0), Le(le=1)] = None, top_k: int | None = None, seed: int | None = None, stop: List[str] | None = None)[source]#

Bases: _ForbidExtraBase

Parameters for sampling / inference operations.

Parameters:

max_tokens – Maximum number of tokens to generate. The service requires a minimum of 512.
temperature – Sampling temperature. Must not be negative; 0 is greedy decoding. The service applies its own upper limit.
top_p – Nucleus sampling probability threshold, between 0 and 1. It is a cumulative probability, so a percentage such as 95 is rejected here rather than at the service.
top_k – Top-k sampling parameter. -1 disables it.
seed – Random seed for reproducibility.
stop – Stop sequences.

classmethod from_any(value: Any) → SamplingParams[source]#

Coerce a params-like value to SamplingParams.

Accepts a SamplingParams (passthrough), a dict of fields, or any object sharing at least one field name. Raises TypeError for an object that shares none, instead of silently returning all defaults.

max_tokens: int#

model_config: ClassVar[ConfigDict] = {'extra': 'forbid', 'validate_assignment': True}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

seed: int | None#

stop: List[str] | None#

temperature: float | None#

top_k: int | None#

top_p: float | None#

class sagemaker.train.training_session.types.SaveStateResponse[source]#

Bases: _OperationResult

Result of a completed save_state operation.

model_config: ClassVar[ConfigDict] = {'extra': 'forbid'}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

property model_package_arn: str | None#: Return the ARN of the model package the checkpoint was registered as.

model_post_init(context: Any, /) → None#

This function is meant to behave like a BaseModel method to initialize private attributes.

It takes context as an argument since that’s what pydantic-core passes when calling it.

Parameters:

self – The BaseModel instance.
context – The context.

class sagemaker.train.training_session.types.SaveWeightsResponse[source]#

Bases: _OperationResult

Result of a completed save_weights_for_sampler operation.

model_config: ClassVar[ConfigDict] = {'extra': 'forbid'}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

property model_package_arn: str | None#: Return the ARN of the model package the weights were registered as.

model_post_init(context: Any, /) → None#

This function is meant to behave like a BaseModel method to initialize private attributes.

It takes context as an argument since that’s what pydantic-core passes when calling it.

Parameters:

self – The BaseModel instance.
context – The context.

class sagemaker.train.training_session.types.SessionStatus(value)[source]#

Bases: str, Enum

Status of a training session.

Use with ServiceClient.list_sessions() to filter by status without memorising the service’s raw PascalCase strings.

BLOCKED = 'Blocked'#

CREATED = 'Created'#

DELETING = 'Deleting'#

DELETION_FAILED = 'DeletionFailed'#

EXPIRED = 'Expired'#

FAILED = 'Failed'#

RUNNING = 'Running'#

SCHEDULED = 'Scheduled'#

STOPPED = 'Stopped'#

STOPPING = 'Stopping'#

class sagemaker.train.training_session.types.StopReason(value)[source]#

Bases: str, Enum

Reason why text generation stopped.

LENGTH = 'length'#

STOP = 'stop'#

class sagemaker.train.training_session.types.TensorData(data: Any = None, *, dtype: TensorDtype = TensorDtype.FLOAT32, shape: List[int] | None = None)[source]#

Bases: _ForbidExtraBase

A named tensor with dtype and shape metadata.

Use in loss_fn_inputs to pass typed tensor data (advantages, weights, logprobs, etc.) alongside training datums.

Tensor-like values (torch tensors, numpy arrays, anything with tolist()) are held as-is and only flattened to a plain list when the request is serialized or data is first read, so building datums from large tensors does not copy them up front.

Parameters:

data – Flat list of numeric values, or a tensor-like object converted lazily.
dtype – Tensor element type.
shape – Tensor dimensions (e.g. [batch_size, seq_len]).

property data: List[Any]#: Tensor values as a flat list (materialized from the source on first access).

dtype: TensorDtype#

classmethod from_any(value: Any) → TensorData[source]#

Coerce a tensor-like value to a TensorData.

Accepts a TensorData (passthrough), a wire-format dict, a torch tensor / numpy array (anything with tolist()), a plain list, or any object with data/dtype attributes.

model_config: ClassVar[ConfigDict] = {'extra': 'forbid'}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

model_post_init(context: Any, /) → None#

This function is meant to behave like a BaseModel method to initialize private attributes.

It takes context as an argument since that’s what pydantic-core passes when calling it.

Parameters:

self – The BaseModel instance.
context – The context.

shape: List[int] | None#

to_numpy() → Any[source]#: Convert to a numpy.ndarray (requires numpy).

to_torch() → Any[source]#: Convert to a torch.Tensor (requires torch).

tolist() → List[Any][source]#: Return the tensor values as a plain list.

class sagemaker.train.training_session.types.TensorDtype(value)[source]#

Bases: str, Enum

Data type for tensor values in loss function inputs.

Maps to the dtypes the runtime accepts in TensorData.dtype.

BFLOAT16 = 'bfloat16'#

FLOAT16 = 'float16'#

FLOAT32 = 'float32'#

INT32 = 'int32'#

INT64 = 'int64'#

class sagemaker.train.training_session.types.TokenLogprob(*, token_id: int, logprob: float)[source]#

Bases: _ForbidExtraBase

A token ID and its log probability.

logprob: float#

model_config: ClassVar[ConfigDict] = {'extra': 'forbid'}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

token_id: int#

class sagemaker.train.training_session.types.TrainingSessionSummary(*, training_session_arn: str, status: SessionStatus | str | None = None, creation_time: datetime | None = None, last_modified_time: datetime | None = None)[source]#

Bases: _ForbidExtraBase

Summary of one training session, as listed.

Returned by ServiceClient.list_sessions(). Carries the fields the list reports; anything else about a session comes from describing it.

Parameters:

training_session_arn – ARN of the session.
status – Current session status.
creation_time – When the session was created.
last_modified_time – When the session last changed state.

creation_time: datetime | None#

last_modified_time: datetime | None#

model_config: ClassVar[ConfigDict] = {'extra': 'forbid'}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

status: SessionStatus | str | None#: str is part of the type so a status this version does not know about still reaches the caller rather than failing the whole list call.

training_session_arn: str#

property training_session_name: str#: The session name, taken from the tail of training_session_arn.

class sagemaker.train.training_session.types.VpcConfig(*, subnets: List[str], security_group_ids: List[str])[source]#

Bases: _ForbidExtraBase

VPC configuration for a training session.

When set, the session’s compute runs inside your VPC: traffic to your Amazon S3 output location and to other resources flows through the subnets and security groups you specify, so make sure they allow the required access (for example an S3 gateway endpoint).

Parameters:

subnets – IDs of the subnets to run the session in.
security_group_ids – IDs of the security groups to apply.

classmethod from_any(value: Any) → VpcConfig[source]#

Coerce a VPC-config-like value to VpcConfig.

Accepts a VpcConfig (passthrough), a dict of fields, or any object with subnets / security_group_ids attributes.

model_config: ClassVar[ConfigDict] = {'extra': 'forbid'}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

security_group_ids: List[str]#

subnets: List[str]#

APIFuture and request-status model for async training/sample operations.

A single APIFuture serves training, sampling and checkpointing operations. It is generic in the type it resolves to: each method returns an APIFuture[T] whose result() is a T, the operation’s own response type. APIFuture is the one future type callers see; the two subclasses below are internal, so a method that may return either is still annotated APIFuture[T].

timeout is the only knob callers pass, and it bounds the whole wait on every future in this module. A future that resolves in stages (PendingAPIFuture waits for its submit and then polls; CombinedFuture waits for several sub-futures) turns timeout into a single deadline at the top of result() and hands each stage only the time still left, so result(timeout=60) returns or raises TimeoutError within 60 seconds regardless of how many stages ran. The default is timeout=None waits indefinitely; pass an explicit timeout to bound a specific wait.

Every future can be waited on either way. result() blocks the calling thread. await (or result_async()) suspends a coroutine instead: no thread is held while the operation runs, only for the individual status requests. So asyncio.gather over hundreds of futures keeps them all in flight.

The polling cadence is the SDK’s to pick, so it is set at construction rather than per wait: poll_interval, time_fn, sleep_fn, async_sleep_fn and rand_fn are constructor seams the SDK and its tests use (tests inject a fake clock and a fixed jitter draw, and run without real waits). Callers receive futures, never build them.

class sagemaker.train.training_session.futures.APIFuture(operation_id: str, poll_fn: ~typing.Callable[[], ~typing.Dict[str, ~typing.Any]], parse_fn: ~typing.Callable[[~typing.Dict[str, ~typing.Any]], ~sagemaker.train.training_session.futures._T], *, status_key: str = 'status', operation_type: str = 'operation', poll_interval: float = 2, time_fn: ~typing.Callable[[], float] = <built-in function monotonic>, sleep_fn: ~typing.Callable[[float], None] = <built-in function sleep>, async_sleep_fn: ~typing.Callable[[float], ~typing.Awaitable[None]] = <function sleep>, max_transient_poll_errors: int = 150, rand_fn: ~typing.Callable[[float, float], float] = <bound method Random.uniform of <random.Random object>>)[source]#

Bases: Generic[_T]

Future for an async runtime operation (training or sample).

Call result() to block until the operation reaches a terminal state. The same future type serves both training and sample operations; the caller supplies a poll_fn that fetches the latest response and a parse_fn that turns a successful response into the public result object.

Generic in the result type: parse_fn fixes it, so result() on an APIFuture[T] returns a T.

Parameters:

operation_id – Server-assigned id of the operation (for messages).
poll_fn – Zero-arg callable returning the latest operation response dict. Each call should hit Get*OperationResult for this operation.
parse_fn – Callable mapping a completed response dict to the result object returned by result().
status_key – Key under which the response carries its status string.
operation_type – Label used in error messages (e.g. "training").
poll_interval – Seconds between poll attempts.
time_fn – Monotonic clock, injectable for tests. Defaults to time.monotonic().
sleep_fn – Sleep function used by result(), injectable for tests. Defaults to time.sleep().
async_sleep_fn – Awaitable sleep used by result_async(), injectable for tests. Defaults to asyncio.sleep().
max_transient_poll_errors – Consecutive transient poll errors (throttling / 5xx on Get*OperationResult) tolerated before giving up. A successful poll resets the counter.
rand_fn – Jitter draw for the backoff between failed polls, injectable for tests. Defaults to random.uniform().

cancel() → bool[source]#

Always False: a submitted operation cannot be called back.

Present because the rest of this surface follows the standard future protocol, and a caller looping over futures should not have to special case this one. The service has the operation by the time a caller holds the future, so there is nothing to take back; stop the session to end the work it is doing.

cancelled() → bool[source]#: Always False: see cancel().

done() → bool[source]#

Whether this future has resolved, without waiting for it.

True once the operation’s outcome is known, whether that is a result or a failure, so a loop that drains futures by checking this terminates on a failed operation as well as a successful one. A future whose operation has finished on the service but has not been waited on yet still reports False: nothing has polled for it.

A failed request is not an outcome and is not recorded here. A call that could not be made or answered, credentials that expired mid-run being the clearest case, says nothing about the operation, which may still be running and may still succeed, so a later wait asks again rather than replaying an answer that was never about the operation.

exception(timeout: float | None = None) → BaseException | None[source]#

Wait for the operation and return the error it failed with.

Returns None if the operation succeeded, so a caller can inspect several futures without wrapping each in try.

A failed operation hands back the same error object on every wait, the way any settled future does, so identity comparisons hold. Re-raising one object repeatedly does lengthen its traceback each time, which is worth knowing if a retry loop waits on the same future many times over; read __cause__ and the message rather than the traceback depth.

Parameters:: timeout – Max seconds to wait before raising TimeoutError. None waits indefinitely.
Raises:: TimeoutError – The operation did not complete within timeout.

property operation_id: str#: Return the server-assigned operation id.

result(timeout: float | None = None) → _T[source]#

Block until the operation completes and return its parsed result.

Parameters:

timeout – Max seconds to wait before raising TimeoutError. None waits indefinitely.

Returns:

The object produced by parse_fn from the completed response.

Raises:

OperationError – The operation reached a terminal failure state.
TimeoutError – The operation did not complete within timeout.
Exception – The error botocore raised for a failed poll, either a transient one that recurred past max_transient_poll_errors in a row, or one that was never worth retrying.

async result_async(timeout: float | None = None) → _T[source]#

Await the operation result without occupying a thread while it runs.

Equivalent to await future, and the way to await with a timeout of your own: await future cannot take arguments, so it uses the default.

Parameters:

timeout – Max seconds to wait before raising TimeoutError. None waits indefinitely.

Raises:

OperationError – The operation reached a terminal failure state.
TimeoutError – The operation did not complete within timeout.
Exception – The error botocore raised for a failed poll, either a transient one that recurred past max_transient_poll_errors in a row, or one that was never worth retrying.

class sagemaker.train.training_session.futures.RequestStatus(value)[source]#

Bases: str, Enum

Lifecycle states of an asynchronous training or sampling operation.

An operation moves from Pending to InProgress to a terminal state: Completed (success), Failed (the operation errored), or Skipped (the session failed and abandoned its remaining queued operations).

COMPLETED = 'Completed'#

FAILED = 'Failed'#

IN_PROGRESS = 'InProgress'#

PENDING = 'Pending'#

SKIPPED = 'Skipped'#

Exceptions for the training_session client interface.

A call that fails on the wire raises the error botocore models for it: ValidationException, ConflictException, ResourceNotFoundException, ThrottlingException, InternalServerException, reachable both as attributes of the client (client.exceptions.ValidationException) and as botocore.exceptions.ClientError, whose response carries the service’s own error code and message. The SDK passes those through untouched, so nothing here has to be caught to read what the service said.

The exceptions defined here cover what a wire error cannot express: an asynchronous operation that was accepted and later failed (OperationError), and two conditions the SDK detects on the client side (ResultMismatchError, ArtifactDeletionError). A condition with nothing specific to report, such as an operation asked of a client that is already closed, raises the base TrainingSessionError itself, so catching that one catches every error defined here.

Running out of time waiting is the exception to that: a wait that outlives its timeout raises the built-in TimeoutError, matching what concurrent.futures and asyncio raise for the same thing, so an existing except TimeoutError keeps working. Catch it alongside TrainingSessionError to cover both. Use is_retryable() to decide whether a failed call is worth sending again.

exception sagemaker.train.training_session.exceptions.ArtifactDeletionError(model_package_arn: str, s3_uri: str, failed_keys: List[str])[source]#

Bases: TrainingSessionError

Raised when deleting a model package’s S3 artifacts did not fully succeed.

The model package is left in place, so it still points at the objects listed on failed_keys. Re-run ServiceClient.delete_model_package() once the cause is addressed (most often a missing S3 permission). The message truncates failed_keys; the attribute is complete.

exception sagemaker.train.training_session.exceptions.OperationError(operation_type: str, message: str, *, operation_id: str | None = None)[source]#

Bases: TrainingSessionError

Raised when an asynchronous operation completed in a failed state.

The operation was accepted, so nothing was raised when it was submitted; it then failed while running, and that surfaces here when the future for it is resolved. A request the service rejects outright raises ClientError instead, at submit time.

operation_id identifies the failed operation server-side and is the one detail to quote when asking for help with it.

exception sagemaker.train.training_session.exceptions.ResultMismatchError(operation_type: str, expected: int, received: int)[source]#

Bases: TrainingSessionError

Raised when a completed operation’s result does not match its request.

A per-token result must be as long as the sequence sent. A different length means the result belongs to some other request, so it is raised instead of returned: the values would otherwise be read as this request’s own.

exception sagemaker.train.training_session.exceptions.TrainingSessionError[source]#

Bases: Exception

Base exception for training session errors.

Every error class in this module derives from this one, and it is also raised directly where the message says all there is to say: asking a closed client for another operation, for instance. A wait that runs out of time raises the built-in TimeoutError instead, so bound waits are caught the same way as with any other future.

sagemaker.train.training_session.exceptions.is_retryable(error: BaseException) → bool[source]#

Return whether sending the failed call again may succeed.

True for throttling, an HTTP 5xx, and connection-level failures such as a read or connect timeout: the request did not arrive, or arrived and could not be served for a reason unrelated to what it asked for. False for everything else, including a rejected request, since resending that one gets the same answer.

The SDK already retries its own internal calls on this basis. Use it to decide the same thing for a call of your own. Resolve the future inside the try: submitting is non-blocking, so the operation’s own failure surfaces where the result is awaited, not where it was queued:

future = training_client.forward_backward(batch, loss_fn="cross_entropy")
try:
    result = future.result()
except Exception as error:
    if not is_retryable(error):
        raise
    ...

Parameters:: error – The exception raised by the failed call.
Returns:: True if the failure is transient, False otherwise.

Model Evaluation#

SageMaker Model Evaluation Module.

This module provides comprehensive evaluation capabilities for SageMaker models:

Classes:

BaseEvaluator: Abstract base class for all evaluators
BenchMarkEvaluator: Standard benchmark evaluations
CustomScorerEvaluator: Custom scorer and preset metrics evaluations
LLMAsJudgeEvaluator: LLM-as-judge evaluations
EvaluationPipelineExecution: Pipeline-based evaluation execution implementation
PipelineExecutionStatus: Combined status with step details and failure reason
StepDetail: Individual pipeline step information

Bases: BaseModel

Base class for SageMaker model evaluators.

Provides common functionality for all evaluators including model resolution, MLflow integration, and AWS resource configuration. Subclasses must implement the evaluate() method.

region#

AWS region for evaluation jobs. If not provided, will use SAGEMAKER_REGION env var or default region.

Type:: Optional[str]

role#

IAM execution role ARN for SageMaker pipeline and training jobs. If not provided, will be derived from the session’s caller identity. Use this when running outside SageMaker-managed environments (e.g., local notebooks, CI/CD) where the caller identity is not a SageMaker-assumable role.

Type:: Optional[str]

sagemaker_session#

SageMaker session object. If not provided, a default session will be created automatically.

Type:: Optional[Any]

model#

Model for evaluation. Can be: - JumpStart model ID (str): e.g., ‘llama3-2-1b-instruct’ - ModelPackage object: A fine-tuned model package - ModelPackage ARN (str): e.g., ‘arn:aws:sagemaker:region:account:model-package/name/version’ - S3 checkpoint path (str): e.g., ‘s3://bucket/path/to/checkpoint’ (for HyperPod outputs) - BaseTrainer object: A completed training job (i.e., it must have

_latest_training_job with output_model_package_arn populated)

Type:: Union[str, Any]

base_model_name#

Base model name for recipe lookup when using S3 checkpoint paths. Required when model is an S3 URI. E.g., ‘amazon.nova-lite-v2’ or ‘nova-textgeneration-lite-v2’.

Type:: Optional[str]

base_eval_name#

Optional base name for evaluation jobs. This name is used as the PipelineExecutionDisplayName when creating the SageMaker pipeline execution. The actual display name will be “{base_eval_name}-{timestamp}”. This parameter can be used to cross-reference the pipeline execution ARN with a human-readable display name in the SageMaker console. If not provided, a unique name will be generated automatically in the format “eval-{model_name}-{uuid}”.

Type:: Optional[str]

s3_output_path#

S3 location for evaluation outputs. Required.

Type:: str

mlflow_resource_arn#

MLflow resource ARN for experiment tracking. Optional. If not provided, the system will attempt to resolve it using the default MLflow app experience (checks domain match, account default, or creates a new app). Supported formats: - MLflow tracking server: arn:aws:sagemaker:region:account:mlflow-tracking-server/name - MLflow app: arn:aws:sagemaker:region:account:mlflow-app/app-id

Type:: Optional[str]

mlflow_experiment_name#

Optional MLflow experiment name for tracking evaluation runs.

Type:: Optional[str]

mlflow_run_name#

Optional MLflow run name for tracking individual evaluation executions.

Type:: Optional[str]

networking#

VPC configuration for evaluation jobs. Accepts a sagemaker_core.shapes.VpcConfig object with security_group_ids and subnets attributes. When provided, evaluation jobs will run within the specified VPC for enhanced security and access to private resources.

Type:: Optional[VpcConfig]

kms_key_id#

AWS KMS key ID for encrypting output data. When provided, evaluation job outputs will be encrypted using this KMS key for enhanced data security.

Type:: Optional[str]

model_package_group#

Model package group. Accepts: 1. ARN string (e.g., ‘arn:aws:sagemaker:region:account:model-package-group/name’) 2. ModelPackageGroup object (ARN will be extracted from model_package_group_arn attribute) 3. Model package group name string (will fetch the object and extract ARN) Required when model is a JumpStart model ID. Optional when model is a ModelPackage ARN/object (will be inferred automatically).

Type:: Optional[Union[str, ModelPackageGroup]]

tags#

Tags applied to the evaluation pipeline when it is created, which cascade to the pipeline’s step jobs.

Type:: Optional[List[TagsDict]]

class Config[source]#

Bases: object

Pydantic model configuration.

arbitrary_types_allowed = True#

base_eval_name: str | None#

base_model_name: str | None#

compute: Compute | HyperPodCompute | None#

evaluate(dry_run: bool = False) → Any[source]#

Create and start an evaluation execution.

This method must be implemented by subclasses to define the specific evaluation logic for different evaluation types (benchmark, custom scorer, LLM-as-judge, etc.).

Parameters:: dry_run (bool) – If True, runs all validation (IAM, model resolution, data paths) without submitting the evaluation. Returns None on success, raises on validation failure. Defaults to False.
Returns:: The created evaluation execution object, or None if dry_run=True.
Return type:: EvaluationPipelineExecution
Raises:: NotImplementedError – This is an abstract method that must be implemented by subclasses.

Example

>>> # In a subclass implementation
>>> class CustomEvaluator(BaseEvaluator):
...     def evaluate(self):
...         # Create pipeline definition
...         pipeline_definition = self._build_pipeline()
...         # Start execution
...         return EvaluationPipelineExecution.start(...)

get_resolved_recipe() → Dict[str, Any][source]#

Return the fully resolved evaluation recipe configuration.

Merges base defaults (from Hub hyperparameters spec) + user recipe YAML + overrides dict and returns the result. Callable before or after evaluate().

When neither recipe nor overrides were provided at construction time but hyperparameters have been set directly (e.g. evaluator.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[str, Any]
Raises:: ValueError – If no recipe, overrides, or direct hyperparameter assignments were provided.

kms_key_id: str | None#

mlflow_experiment_name: str | None#

mlflow_resource_arn: str | None#

mlflow_run_name: str | None#

model: str | BaseTrainer | AgentRFTJob | ModelPackage#

model_config: ClassVar[ConfigDict] = {'arbitrary_types_allowed': True}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

model_package_group: str | ModelPackageGroup | None#

model_post_init(context: Any, /) → None#

This function is meant to behave like a BaseModel method to initialize private attributes.

It takes context as an argument since that’s what pydantic-core passes when calling it.

Parameters:

self – The BaseModel instance.
context – The context.

networking: VpcConfig | None#

overrides: Dict[str, Any] | None#

recipe: str | None#

region: str | None#

role: str | None#

s3_output_path: str#

sagemaker_session: Any | None#

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

Stream CloudWatch logs for the latest evaluation execution.

Dispatches to the underlying execution object’s stream_logs() for pipeline-based evaluations, or streams directly from the HyperPod cluster log group for HyperPod evaluations.

Parameters:

poll – Seconds between CloudWatch polling cycles (1-300).
start_time – Stream from this timestamp. Accepts datetime or epoch milliseconds (int). If None, defaults depend on the backend.

Raises:

ValueError – If no evaluation has been executed yet or poll is out of range.

tags: List[Dict[str, str | PipelineVariable]] | None#

training_image: str | None#

Bases: BaseEvaluator

Benchmark evaluator for standard model evaluation tasks.

This evaluator accepts a benchmark enum and automatically deduces the appropriate metrics, strategy, and subtask availability based on the benchmark configuration. Supports various standard benchmarks like MMLU, BBH, MATH, MMMU, and others.

benchmark#

Benchmark type from the Benchmark enum obtained via get_benchmarks(). Required. Use get_benchmarks() to access available benchmark types.

Type:: _Benchmark

subtasks#

Benchmark subtask(s) to evaluate. Defaults to ‘ALL’ for benchmarks that support subtasks. Can be a single subtask string, a list of subtasks, or ‘ALL’ to run all subtasks. For benchmarks without subtask support, must be None.

Type:: Optional[Union[str, list[str]]]

mlflow_resource_arn#

ARN of the MLflow tracking server for experiment tracking. Optional. If not provided, the system will attempt to resolve it using the default MLflow app experience (checks domain match, account default, or creates a new app). Format: arn:aws:sagemaker:region:account:mlflow-tracking-server/name

Type:: Optional[str]

evaluate_base_model#

Whether to evaluate the base model in addition to the custom model. Set to False to skip base model evaluation and only evaluate the custom model. Defaults to True (evaluates both models).

Type:: bool

recipe#

Path to a user recipe YAML file (local or S3 URI) for evaluation configuration. Optional. When provided, values are merged with overrides.

Type:: Optional[str]

overrides#

Programmatic overrides dict (nested structure). These take highest precedence over recipe file and base defaults.

Type:: Optional[Dict[str, Any]]

region#

AWS region. Inherited from BaseEvaluator.

Type:: Optional[str]

sagemaker_session#

SageMaker session object. Inherited from BaseEvaluator.

Type:: Optional[Any]

model#

Model for evaluation. Inherited from BaseEvaluator.

Type:: Union[str, Any]

base_eval_name#

Base name for evaluation jobs. Inherited from BaseEvaluator.

Type:: Optional[str]

s3_output_path#

S3 location for evaluation outputs. Inherited from BaseEvaluator.

Type:: str

mlflow_experiment_name#

MLflow experiment name. Inherited from BaseEvaluator.

Type:: Optional[str]

mlflow_run_name#

MLflow run name. Inherited from BaseEvaluator.

Type:: Optional[str]

networking#

VPC configuration. Inherited from BaseEvaluator.

Type:: Optional[VpcConfig]

kms_key_id#

KMS key ID for encryption. Inherited from BaseEvaluator.

Type:: Optional[str]

model_package_group#

Model package group. Inherited from BaseEvaluator.

Type:: Optional[Union[str, ModelPackageGroup]]

Example

# Get available benchmarks
Benchmark = get_benchmarks()

# Create evaluator with benchmark and subtasks
evaluator = BenchMarkEvaluator(
    benchmark=Benchmark.MMLU,
    subtasks=["abstract_algebra", "anatomy", "astronomy"],
    model="llama3-2-1b-instruct",
    s3_output_path="s3://bucket/outputs/",
    mlflow_resource_arn="arn:aws:sagemaker:us-west-2:123456789012:mlflow-tracking-server/my-server"
)

# Run evaluation with configured subtasks
execution = evaluator.evaluate()
execution.wait()

# Or override subtasks at evaluation time
execution = evaluator.evaluate(subtask="abstract_algebra")

# With recipe file for evaluation config
evaluator = BenchMarkEvaluator(
    benchmark=Benchmark.MMLU,
    model="amazon-nova-pro-v2",
    s3_output_path="s3://bucket/outputs/",
    recipe="./eval-recipes/nova-pro-mmlu.yaml",
    overrides={"inference": {"max_new_tokens": 4096, "temperature": 0}},
)
resolved = evaluator.get_resolved_recipe()

base_eval_name: str | None#

base_model_name: str | None#

benchmark: _Benchmark#

compute: Compute | HyperPodCompute | None#

evaluate(subtask: str | List[str] | None = None, dry_run: bool = False) → EvaluationPipelineExecution[source]#

Create and start a benchmark evaluation job.

Supports multiple compute backends via the compute parameter set at construction time: - Serverless (default): Runs via SageMaker Pipelines. - SMTJ: Runs on user-managed instances via ModelTrainer. - HyperPod: Submits to a HyperPod cluster via the HyperPod CLI.

Parameters:

subtask (Optional[Union[str, list[str]]]) – Optional subtask(s) to evaluate. If not provided, uses the subtasks from constructor. Can be a single subtask string, a list of subtasks, or ‘ALL’ to run all subtasks.
dry_run (bool) – If True, runs all validation (IAM, model resolution, data paths) without submitting the evaluation. Returns None on success, raises on validation failure. Defaults to False.

Returns:

The created benchmark evaluation execution, or None if dry_run=True.

Return type:

EvaluationPipelineExecution

Example

Benchmark = get_benchmarks()
evaluator = BenchMarkEvaluator(
    benchmark=Benchmark.MMLU,
    subtasks="ALL",
    model="llama3-2-1b-instruct",
    s3_output_path="s3://bucket/outputs/"
)

# Evaluate single subtask
execution = evaluator.evaluate(subtask="abstract_algebra")

# Evaluate multiple subtasks
execution = evaluator.evaluate(subtask=["abstract_algebra", "anatomy"])

# Evaluate all subtasks (uses constructor default)
execution = evaluator.evaluate()

evaluate_base_model: bool#

classmethod get_all(session: Any | None = None, region: str | None = None) → Iterator[EvaluationPipelineExecution][source]#

Get all benchmark evaluation executions.

Uses EvaluationPipelineExecution.get_all() to retrieve all benchmark evaluation executions as an iterator.

Parameters:

session (Optional[Any]) – Optional boto3 session. If not provided, will be inferred.
region (Optional[str]) – Optional AWS region. If not provided, will be inferred.

Yields:

EvaluationPipelineExecution – Benchmark evaluation execution instances.

Example

# Get all benchmark evaluations as iterator
eval_iter = BenchMarkEvaluator.get_all()
all_executions = list(eval_iter)

# Or iterate directly
for execution in BenchMarkEvaluator.get_all():
    print(f"{execution.name}: {execution.status.overall_status}")

# With specific session/region
eval_iter = BenchMarkEvaluator.get_all(session=my_session, region='us-west-2')
all_executions = list(eval_iter)

property hyperparameters#

Get evaluation hyperparameters as a FineTuningOptions object.

This property provides access to evaluation hyperparameters with validation, type checking, and user-friendly information display. Hyperparameters are lazily loaded from the JumpStart Hub when first accessed.

Returns:: Dynamic object with evaluation hyperparameters
Return type:: FineTuningOptions
Raises:: ValueError – If base model name is not available or if hyperparameters cannot be loaded

Example

evaluator = BenchMarkEvaluator(...)

# Access current values
print(evaluator.hyperparameters.temperature)

# Modify values (with validation)
evaluator.hyperparameters.temperature = 0.5

# Get as dictionary
params = evaluator.hyperparameters.to_dict()

# Display parameter information
evaluator.hyperparameters.get_info()
evaluator.hyperparameters.get_info('temperature')

kms_key_id: str | None#

mlflow_experiment_name: str | None#

mlflow_resource_arn: str | None#

mlflow_run_name: str | None#

model: str | BaseTrainer | AgentRFTJob | ModelPackage#

model_config: ClassVar[ConfigDict] = {'arbitrary_types_allowed': True}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

model_package_group: str | ModelPackageGroup | None#

model_post_init(context: Any, /) → None#

This function is meant to behave like a BaseModel method to initialize private attributes.

It takes context as an argument since that’s what pydantic-core passes when calling it.

Parameters:

self – The BaseModel instance.
context – The context.

networking: VpcConfig | None#

overrides: Dict[str, Any] | None#

recipe: str | None#

region: str | None#

role: str | None#

s3_output_path: str#

sagemaker_session: Any | None#

subtasks: str | List[str] | None#

tags: List[TagsDict] | None#

training_image: str | None#

Bases: BaseEvaluator

Custom scorer evaluation job for preset or custom evaluator metrics.

This evaluator supports both preset metrics (via built-in metrics enum) and custom evaluator implementations for specialized evaluation needs.

evaluator#

Built-in metric enum value, Evaluator object, or Evaluator ARN string. Required. Use get_builtin_metrics() for available preset metrics.

Type:: Union[str, Any]

dataset#

Dataset for evaluation. Required. Accepts S3 URI, Dataset ARN, or DataSet object.

Type:: Any

mlflow_resource_arn#

ARN of the MLflow tracking server for experiment tracking. Optional. If not provided, the system will attempt to resolve it using the default MLflow app experience (checks domain match, account default, or creates a new app). Inherited from BaseEvaluator.

Type:: Optional[str]

evaluate_base_model#

Whether to evaluate the base model in addition to the custom model. Set to False to skip base model evaluation and only evaluate the custom model. Defaults to True (evaluates both models).

Type:: bool

region#

AWS region. Inherited from BaseEvaluator.

Type:: Optional[str]

sagemaker_session#

SageMaker session object. Inherited from BaseEvaluator.

Type:: Optional[Any]

model#

Model for evaluation. Inherited from BaseEvaluator.

Type:: Union[str, Any]

base_eval_name#

Base name for evaluation jobs. Inherited from BaseEvaluator.

Type:: Optional[str]

s3_output_path#

S3 location for evaluation outputs. Inherited from BaseEvaluator.

Type:: str

mlflow_experiment_name#

MLflow experiment name. Inherited from BaseEvaluator.

Type:: Optional[str]

mlflow_run_name#

MLflow run name. Inherited from BaseEvaluator.

Type:: Optional[str]

networking#

VPC configuration. Inherited from BaseEvaluator.

Type:: Optional[VpcConfig]

kms_key_id#

KMS key ID for encryption. Inherited from BaseEvaluator.

Type:: Optional[str]

model_package_group#

Model package group. Inherited from BaseEvaluator.

Type:: Optional[Union[str, ModelPackageGroup]]

Example

from sagemaker.train.evaluate.custom_scorer_evaluator import (
    CustomScorerEvaluator,
    get_builtin_metrics
)
from sagemaker.ai_registry.evaluator import Evaluator

# Using preset metric
BuiltInMetric = get_builtin_metrics()
evaluator = CustomScorerEvaluator(
    evaluator=BuiltInMetric.PRIME_MATH,
    dataset=my_dataset,
    base_model="my-model",
    s3_output_path="s3://bucket/output",
    mlflow_resource_arn="arn:aws:sagemaker:us-west-2:123456789012:mlflow-tracking-server/my-server"
)

# Using custom evaluator
my_evaluator = Evaluator.create(
    name="my-custom-evaluator",
    function_source="/path/to/evaluator.py",
    sub_type="AWS/Evaluator"
)
evaluator = CustomScorerEvaluator(
    evaluator=my_evaluator,
    dataset=my_dataset,
    base_model="my-model",
    s3_output_path="s3://bucket/output",
    mlflow_resource_arn="arn:aws:sagemaker:us-west-2:123456789012:mlflow-tracking-server/my-server"
)

# Using evaluator ARN string
evaluator = CustomScorerEvaluator(
    evaluator="arn:aws:sagemaker:us-west-2:123456789012:hub-content/AIRegistry/Evaluator/my-evaluator/1",
    dataset=my_dataset,
    base_model="my-model",
    s3_output_path="s3://bucket/output",
    mlflow_resource_arn="arn:aws:sagemaker:us-west-2:123456789012:mlflow-tracking-server/my-server"
)

job = evaluator.evaluate()

base_eval_name: str | None#

base_model_name: str | None#

compute: Compute | HyperPodCompute | None#

dataset: Any#

evaluate(dry_run: bool = False) → EvaluationPipelineExecution[source]#

Create and start a custom scorer evaluation job.

Supports multiple compute backends via the compute parameter set at construction time: - Serverless (default): Runs via SageMaker Pipelines. - SMTJ: Runs on user-managed instances via ModelTrainer. - HyperPod: Submits to a HyperPod cluster via the HyperPod CLI.

Parameters:: dry_run (bool) – If True, runs all validation (IAM, model resolution, data paths) without submitting the evaluation. Returns None on success, raises on validation failure. Defaults to False.
Returns:: The created custom scorer evaluation execution, or None if dry_run=True.
Return type:: EvaluationPipelineExecution

Example

evaluator = CustomScorerEvaluator(
    evaluator=BuiltInMetric.CODE_EXECUTIONS,
    dataset=my_dataset,
    base_model="my-model",
    s3_output_path="s3://bucket/output",
    mlflow_resource_arn="arn:..."
)
execution = evaluator.evaluate()
execution.wait()

evaluate_base_model: bool#

evaluator: str | Any#

classmethod get_all(session: Any | None = None, region: str | None = None)[source]#

Get all custom scorer evaluation executions.

Uses EvaluationPipelineExecution.get_all() to retrieve all custom scorer evaluation executions as an iterator.

Parameters:

session (Optional[Any]) – Optional boto3 session. If not provided, will be inferred.
region (Optional[str]) – Optional AWS region. If not provided, will be inferred.

Yields:

EvaluationPipelineExecution – Custom scorer evaluation execution instances

Example

# Get all custom scorer evaluations as iterator
evaluations = CustomScorerEvaluator.get_all()
all_executions = list(evaluations)

# Or iterate directly
for execution in CustomScorerEvaluator.get_all():
    print(f"{execution.name}: {execution.status.overall_status}")

# With specific session/region
evaluations = CustomScorerEvaluator.get_all(session=my_session, region='us-west-2')
all_executions = list(evaluations)

property hyperparameters#

Get evaluation hyperparameters as a FineTuningOptions object.

This property provides access to evaluation hyperparameters with validation, type checking, and user-friendly information display. Hyperparameters are lazily loaded from the JumpStart Hub when first accessed.

Returns:: Dynamic object with evaluation hyperparameters
Return type:: FineTuningOptions
Raises:: ValueError – If base model name is not available or if hyperparameters cannot be loaded

Example

evaluator = CustomScorerEvaluator(...)

# Access current values
print(evaluator.hyperparameters.temperature)

# Modify values (with validation)
evaluator.hyperparameters.temperature = 0.5

# Get as dictionary
params = evaluator.hyperparameters.to_dict()

# Display parameter information
evaluator.hyperparameters.get_info()
evaluator.hyperparameters.get_info('temperature')

kms_key_id: str | None#

mlflow_experiment_name: str | None#

mlflow_resource_arn: str | None#

mlflow_run_name: str | None#

model: str | BaseTrainer | AgentRFTJob | ModelPackage#

model_config: ClassVar[ConfigDict] = {'arbitrary_types_allowed': True}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

model_package_group: str | ModelPackageGroup | None#

model_post_init(context: Any, /) → None#

This function is meant to behave like a BaseModel method to initialize private attributes.

It takes context as an argument since that’s what pydantic-core passes when calling it.

Parameters:

self – The BaseModel instance.
context – The context.

networking: VpcConfig | None#

overrides: Dict[str, Any] | None#

recipe: str | None#

region: str | None#

role: str | None#

s3_output_path: str#

sagemaker_session: Any | None#

tags: List[TagsDict] | None#

training_image: str | None#

class sagemaker.train.evaluate.EvaluationPipelineExecution(*, arn: str | None = None, name: str, status: ~sagemaker.train.evaluate.execution.PipelineExecutionStatus = <factory>, last_modified_time: ~datetime.datetime | None = None, eval_type: ~sagemaker.train.evaluate.constants.EvalType | None = None, s3_output_path: str | None = None, steps: ~typing.List[~typing.Dict[str, ~typing.Any]] = <factory>)[source]#

Bases: BaseModel

Manages SageMaker pipeline-based evaluation execution lifecycle.

This class wraps SageMaker Pipeline execution to provide a simplified interface for running, monitoring, and managing evaluation jobs. Users typically don’t instantiate this class directly, but receive instances from evaluator classes.

Example

from sagemaker.train.evaluate import BenchmarkEvaluator
from sagemaker.train.evaluate.execution import EvaluationPipelineExecution

# Start evaluation through evaluator
evaluator = BenchmarkEvaluator(...)
execution = evaluator.evaluate()

# Monitor execution
print(f"Status: {execution.status.overall_status}")
print(f"Steps: {len(execution.status.step_details)}")

# Wait for completion
execution.wait()

# Display results
execution.show_results()

# Retrieve past executions
all_executions = list(EvaluationPipelineExecution.get_all())
specific_execution = EvaluationPipelineExecution.get(arn="arn:...")

Parameters:

arn (Optional[str]) – ARN of the pipeline execution.
name (str) – Name of the evaluation execution.
status (PipelineExecutionStatus) – Combined status with step details and failure reason.
last_modified_time (Optional[datetime]) – Last modification timestamp.
eval_type (Optional[EvalType]) – Type of evaluation (BENCHMARK, CUSTOM_SCORER, LLM_AS_JUDGE).
s3_output_path (Optional[str]) – S3 location where evaluation results are stored.
steps (List[Dict[str, Any]]) – Raw step information from SageMaker.

class Config[source]#

Bases: object

Pydantic model configuration.

arbitrary_types_allowed = True#

arn: str | None#

eval_type: EvalType | None#

classmethod get(arn: str, session: Session | None = None, region: str | None = None) → EvaluationPipelineExecution[source]#

Get a sagemaker pipeline execution instance by ARN.

Parameters:

arn (str) – ARN of the pipeline execution.
session (Optional[Session]) – Boto3 session. Will be inferred if not provided.
region (Optional[str]) – AWS region. Will be inferred if not provided.

Returns:

Retrieved pipeline execution instance.

Return type:

EvaluationPipelineExecution

Raises:

ClientError – If AWS service call fails.

Example

# Get execution by ARN
arn = "arn:aws:sagemaker:us-west-2:123456789012:pipeline/eval-pipeline/execution/abc123"
execution = EvaluationPipelineExecution.get(arn=arn)
print(execution.status.overall_status)

classmethod get_all(eval_type: EvalType | None = None, session: Session | None = None, region: str | None = None)[source]#

Get all pipeline executions, optionally filtered by evaluation type.

Searches for existing pipelines using prefix and tag validation, then retrieves executions from those pipelines.

Parameters:

eval_type (Optional[EvalType]) – Evaluation type to filter by (e.g., EvalType.BENCHMARK). If None, returns executions from all evaluation pipelines.
session (Optional[Session]) – Boto3 session. Will be inferred if not provided.
region (Optional[str]) – AWS region. Will be inferred if not provided.

Yields:

EvaluationPipelineExecution – Pipeline execution instances.

Example

# Get all evaluation executions as iterator
iter = EvaluationPipelineExecution.get_all()
all_executions = list(iter)

# Get only benchmark evaluations
iter = EvaluationPipelineExecution.get_all(eval_type=EvalType.BENCHMARK)
benchmark_executions = list(iter)

last_modified_time: datetime | None#

model_config: ClassVar[ConfigDict] = {'arbitrary_types_allowed': True}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

name: str#

refresh() → None[source]#: Describe a pipeline execution and update job status

s3_output_path: str | None#

classmethod start(eval_type: EvalType, name: str, pipeline_definition: str, role_arn: str, s3_output_path: str | None = None, session: Session | None = None, region: str | None = None, tags: List[Dict[str, str | PipelineVariable]] | None = []) → EvaluationPipelineExecution[source]#

Create sagemaker pipeline execution. Optionally creates pipeline.

Parameters:

eval_type (EvalType) – Type of evaluation (BENCHMARK, CUSTOM_SCORER, LLM_AS_JUDGE).
name (str) – Name for the evaluation execution.
pipeline_definition (str) – Complete rendered pipeline definition as JSON string.
role_arn (str) – IAM role ARN for pipeline execution.
s3_output_path (Optional[str]) – S3 location where evaluation results are stored.
session (Optional[Session]) – Boto3 session for API calls.
region (Optional[str]) – AWS region for the pipeline.
tags (Optional[List[TagsDict]]) – List of tags to include in pipeline

Returns:

Started pipeline execution instance.

Return type:

EvaluationPipelineExecution

Raises:

ValueError – If pipeline_definition is not valid JSON.
ClientError – If AWS service call fails.

status: PipelineExecutionStatus#

steps: List[Dict[str, Any]]#

stop() → None[source]#: Stop a pipeline execution

wait(target_status: Literal['Executing', 'Stopping', 'Stopped', 'Failed', 'Succeeded'] = 'Succeeded', poll: int = 5, timeout: int | None = None) → None[source]#

Wait for a pipeline execution to reach certain status.

This method provides a hybrid implementation that works in both Jupyter notebooks and terminal environments, with appropriate visual feedback for each.

Parameters:

target_status – The status to wait for
poll – The number of seconds to wait between each poll
timeout – The maximum number of seconds to wait before timing out

class sagemaker.train.evaluate.InspectAIEvaluator(*, region: str | None = None, role: str | None = None, sagemaker_session: Any | None = None, model: str | BaseTrainer | AgentRFTJob | ModelPackage, base_model_name: str | None = None, base_eval_name: str | None = None, s3_output_path: str, mlflow_resource_arn: str | None = None, mlflow_experiment_name: str | None = None, mlflow_run_name: str | None = None, networking: VpcConfig | None = None, kms_key_id: str | None = None, model_package_group: str | ModelPackageGroup | None = None, tags: List[Dict[str, str | PipelineVariable]] | None = None, compute: Compute | HyperPodCompute | None = None, training_image: str | None = None, recipe: str | None = None, overrides: Dict[str, Any] | None = None, benchmarks_path: str, tasks: List[Dict[str, Any]] | None = None, output_format: str | None = None, bedrock_model_id: str | None = None, endpoint_name: str | None = None, model_s3_uri: str | None = None, inference_image_uri: str | None = None, endpoint_instance_type: str | None = None, endpoint_instance_count: int = 1, endpoint_execution_role_arn: str | None = None, context_length: str | None = None, max_concurrency: str | None = None, cleanup_endpoint: bool = True, endpoint_prefix: str = 'inspectai', endpoint_environment: Dict[str, str] | None = None, extra_args: List[str] | None = None, environment: Dict[str, str] | None = None, image_uri: str | None = None, instance_type: str = 'ml.m5.large', max_runtime_seconds: int = 86400, max_connections: int = 16, max_retries: int = 100, timeout: int = 600, temperature: float = 0.0, top_p: float = 1.0, top_k: int = -1, max_tokens: int = 8192)[source]#

Bases: BaseEvaluator

InspectAI evaluation job.

Runs InspectAI tasks inside a SageMaker Training container, supporting three inference provider modes: Bedrock, existing SageMaker endpoint, or creating a new endpoint.

The evaluator serializes configuration to a YAML file (inspect_config.yaml), uploads it to S3, and launches a single-step SageMaker Pipeline that runs the InspectAI container with that config as input.

Supports resource chaining: a completed trainer (e.g., SFTTrainer, DPOTrainer, MultiTurnRLTrainer) can be passed directly as the model parameter. The evaluator will automatically resolve the trainer’s output model package artifacts and configure endpoint creation for evaluation.

benchmarks_path#

S3 URI pointing to benchmark .py files with @task decorators. Required.

Type:: str

tasks#

List of task configurations. Each dict must have a "name" key. Optional keys: "path" (must end with .py), "limit" (int >= 1), "epochs" (int >= 1), "task_args" (dict). If None or empty, all tasks at benchmarks_path are run.

Type:: Optional[List[Dict[str, Any]]]

output_format#

Output format for results. One of "eval", "csv", "jsonl", "json".

Type:: Optional[str]

bedrock_model_id#

Explicit Bedrock model ID for bedrock inference mode. Falls back to the model’s bedrock_model_id if not set.

Type:: Optional[str]

endpoint_name#

Existing SageMaker endpoint name. Mutually exclusive with model_s3_uri/inference_image_uri.

Type:: Optional[str]

model_s3_uri#

S3 URI of model artifacts for creating a new endpoint. Must be paired with inference_image_uri.

Type:: Optional[str]

inference_image_uri#

ECR image URI for creating a new endpoint. Must be paired with model_s3_uri.

Type:: Optional[str]

endpoint_instance_type#

Instance type for new endpoint (must start with ml.).

Type:: Optional[str]

endpoint_instance_count#

Instance count for new endpoint. Defaults to 1.

Type:: int

endpoint_execution_role_arn#

IAM role ARN for new endpoint.

Type:: Optional[str]

context_length#

Context length as string integer.

Type:: Optional[str]

max_concurrency#

Max concurrency as string integer.

Type:: Optional[str]

cleanup_endpoint#

Delete endpoint after evaluation. Defaults to True.

Type:: bool

endpoint_prefix#

Prefix for auto-created endpoint names.

Type:: str

endpoint_environment#

Env vars for the inference endpoint container.

Type:: Optional[Dict[str, str]]

extra_args#

Additional CLI args forwarded to inspect eval.

Type:: Optional[List[str]]

environment#

Env vars for the SageMaker Training Job container.

Type:: Optional[Dict[str, str]]

image_uri#

Override for the InspectAI container image URI.

Type:: Optional[str]

instance_type#

Instance type for the orchestrator Training Job (CPU-only). Defaults to "ml.m5.large".

Type:: str

max_runtime_seconds#

Max runtime for the Training Job in seconds. Defaults to 86400 (24 hours).

Type:: int

max_connections#

Max concurrent inference connections used by the InspectAI eval runner. Defaults to 16.

Type:: int

max_retries#

Max retries per inference request. Defaults to 100.

Type:: int

timeout#

Per-request timeout in seconds. Defaults to 600.

Type:: int

temperature#

Sampling temperature in [0.0, 2.0]. Defaults to 0.0.

Type:: float

top_p#

Nucleus sampling cutoff in [0.0, 1.0]. Defaults to 1.0.

Type:: float

top_k#

Top-k sampling cutoff. Use -1 to disable. Defaults to -1.

Type:: int

max_tokens#

Max tokens to generate per response. Defaults to 8192.

Type:: int

Example

from sagemaker.train.evaluate import InspectAIEvaluator

evaluator = InspectAIEvaluator(
    model="amazon-nova-lite-v1",
    benchmarks_path="s3://my-bucket/benchmarks/",
    tasks=[{"name": "boolq_pt", "limit": 10}],
    s3_output_path="s3://my-bucket/eval-output/",
)
execution = evaluator.evaluate()
execution.wait()
execution.show_results()

Resource chaining with a trainer:

from sagemaker.train import SFTTrainer
from sagemaker.train.evaluate import InspectAIEvaluator

# Train a model
trainer = SFTTrainer(model="llama3-2-1b-instruct", ...)
trainer.train(training_dataset="s3://bucket/data.jsonl")

# Evaluate the fine-tuned model directly
evaluator = InspectAIEvaluator(
    model=trainer,
    benchmarks_path="s3://my-bucket/benchmarks/",
    tasks=[{"name": "boolq_pt", "limit": 10}],
    s3_output_path="s3://my-bucket/eval-output/",
)
execution = evaluator.evaluate()

base_eval_name: str | None#

base_model_name: str | None#

bedrock_model_id: str | None#

benchmarks_path: str#

cleanup_endpoint: bool#

compute: Compute | HyperPodCompute | None#

context_length: str | None#

endpoint_environment: Dict[str, str] | None#

endpoint_execution_role_arn: str | None#

endpoint_instance_count: int#

endpoint_instance_type: str | None#

endpoint_name: str | None#

endpoint_prefix: str#

environment: Dict[str, str] | None#

evaluate(dry_run: bool = False) → EvaluationPipelineExecution | None[source]#

Create and start an InspectAI evaluation job.

Serializes the InspectAI configuration to YAML, uploads it to S3, and launches a single-step SageMaker Pipeline with the InspectAI container.

Parameters:

dry_run (bool) – If True, runs all validation (IAM, model resolution, config building, template rendering) without uploading config to S3 or submitting the evaluation. Returns None on success, raises on validation failure. Defaults to False.

Returns:

The started evaluation execution with: .wait(), .refresh(), and .show_results() methods, or None if dry_run=True.

Return type:

EvaluationPipelineExecution

Example

evaluator = InspectAIEvaluator(
    model="amazon-nova-lite-v1",
    benchmarks_path="s3://my-bucket/benchmarks/",
    tasks=[{"name": "boolq_pt", "limit": 10}],
    s3_output_path="s3://my-bucket/eval-output/",
)
execution = evaluator.evaluate()
execution.wait()
execution.show_results()

# Validate without submitting:
evaluator.evaluate(dry_run=True)

extra_args: List[str] | None#

classmethod get_all(session: Any | None = None, region: str | None = None) → Iterator[EvaluationPipelineExecution][source]#

Get all InspectAI evaluation executions.

Parameters:

session (Optional[Any]) – Optional boto3 session.
region (Optional[str]) – Optional AWS region.

Yields:

EvaluationPipelineExecution – InspectAI evaluation execution instances.

image_uri: str | None#

inference_image_uri: str | None#

instance_type: str#

kms_key_id: str | None#

max_concurrency: str | None#

max_connections: int#

max_retries: int#

max_runtime_seconds: int#

max_tokens: int#

mlflow_experiment_name: str | None#

mlflow_resource_arn: str | None#

mlflow_run_name: str | None#

model: str | BaseTrainer | AgentRFTJob | ModelPackage#

model_config: ClassVar[ConfigDict] = {'arbitrary_types_allowed': True}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

model_package_group: str | ModelPackageGroup | None#

model_post_init(context: Any, /) → None#

This function is meant to behave like a BaseModel method to initialize private attributes.

It takes context as an argument since that’s what pydantic-core passes when calling it.

Parameters:

self – The BaseModel instance.
context – The context.

model_s3_uri: str | None#

networking: VpcConfig | None#

output_format: str | None#

overrides: Dict[str, Any] | None#

recipe: str | None#

region: str | None#

role: str | None#

s3_output_path: str#

sagemaker_session: Any | None#

tags: List[TagsDict] | None#

tasks: List[Dict[str, Any]] | None#

temperature: float#

timeout: int#

top_k: int#

top_p: float#

training_image: str | None#

upload_benchmarks(local_path: str) → str[source]#

Upload local benchmark files to S3.

Uploads all files from a local directory to an S3 prefix under the configured output path. The uploaded path can be used as benchmarks_path.

Parameters:: local_path – Local directory path containing .py files with @task decorators.
Returns:: S3 URI prefix where benchmarks were uploaded.
Raises:: ValueError – If local_path does not exist or is not a directory.

Bases: BaseEvaluator

LLM-as-judge evaluation job.

This evaluator uses foundation models to evaluate LLM responses based on various quality and responsible AI metrics.

This feature is powered by Amazon Bedrock Evaluations. Your use of this feature is subject to pricing of Amazon Bedrock Evaluations, the Service Terms applicable to Amazon Bedrock, and the terms that apply to your usage of third-party models. Amazon Bedrock Evaluations may securely transmit data across AWS Regions within your geography for processing. For more information, access Amazon Bedrock Evaluations documentation.

Documentation: https://docs.aws.amazon.com/bedrock/latest/userguide/evaluation-judge.html

evaluator_model#

AWS Bedrock foundation model identifier to use as the judge. Required. For supported models, see: https://docs.aws.amazon.com/bedrock/latest/userguide/evaluation-judge.html#evaluation-judge-supported

Type:: str

dataset#

Evaluation dataset. Required. Accepts: - S3 URI (str): e.g., ‘s3://bucket/path/dataset.jsonl’ - Dataset ARN (str): e.g., ‘arn:aws:sagemaker:…:hub-content/AIRegistry/DataSet/…’ - DataSet object: sagemaker.ai_registry.dataset.DataSet instance (ARN inferred automatically)

Type:: Union[str, Any]

builtin_metrics#

List of built-in evaluation metric names to compute. The ‘Builtin.’ prefix from Bedrock documentation is optional and will be automatically removed if present. Examples: [‘Correctness’, ‘Faithfulness’] or [‘Builtin.Correctness’, ‘Builtin.Faithfulness’]. Optional.

Type:: Optional[List[str]]

custom_metrics#

JSON string containing array of custom metric definitions. Optional. For format details, see: https://docs.aws.amazon.com/bedrock/latest/userguide/model-evaluation-custom-metrics-prompt-formats.html

Type:: Optional[str]

mlflow_resource_arn#

ARN of the MLflow tracking server for experiment tracking. Optional. If not provided, the system will attempt to resolve it using the default MLflow app experience (checks domain match, account default, or creates a new app). Inherited from BaseEvaluator.

Type:: Optional[str]

evaluate_base_model#

Whether to evaluate the base model in addition to the custom model. Set to False to skip base model evaluation and only evaluate the custom model. Defaults to True (evaluates both models).

Type:: bool

region#

AWS region. Inherited from BaseEvaluator.

Type:: Optional[str]

sagemaker_session#

SageMaker session object. Inherited from BaseEvaluator.

Type:: Optional[Any]

model#

Model for evaluation. Inherited from BaseEvaluator.

Type:: Union[str, Any]

base_eval_name#

Base name for evaluation jobs. Inherited from BaseEvaluator.

Type:: Optional[str]

s3_output_path#

S3 location for evaluation outputs. Inherited from BaseEvaluator.

Type:: str

mlflow_experiment_name#

MLflow experiment name. Inherited from BaseEvaluator.

Type:: Optional[str]

mlflow_run_name#

MLflow run name. Inherited from BaseEvaluator.

Type:: Optional[str]

networking#

VPC configuration. Inherited from BaseEvaluator.

Type:: Optional[VpcConfig]

kms_key_id#

KMS key ID for encryption. Inherited from BaseEvaluator.

Type:: Optional[str]

model_package_group#

Model package group. Inherited from BaseEvaluator.

Type:: Optional[Union[str, ModelPackageGroup]]

Example

from sagemaker.train.evaluate import LLMAsJudgeEvaluator

# Example with built-in metrics (prefix optional)
# Both formats work - with or without 'Builtin.' prefix
evaluator = LLMAsJudgeEvaluator(
    base_model="llama-3-3-70b-instruct",
    evaluator_model="anthropic.claude-sonnet-4-5-20250929-v1:0",
    dataset="s3://my-bucket/my-dataset.jsonl",
    builtin_metrics=["Correctness", "Helpfulness"],  # Prefix optional
    mlflow_resource_arn="arn:aws:sagemaker:us-west-2:123456789012:mlflow-tracking-server/my-server",
    s3_output_path="s3://my-bucket/output"
)
execution = evaluator.evaluate()

# Example with custom metrics
custom_metrics = [
    {
        "customMetricDefinition": {
            "name": "PositiveSentiment",
            "instructions": "Assess if the response has positive sentiment. "
            "Prompt: {{prompt}}\nResponse: {{prediction}}",
            "ratingScale": [
                {"definition": "Good", "value": {"floatValue": 1.0}},
                {"definition": "Poor", "value": {"floatValue": 0.0}}
            ]
        }
    }
]

evaluator = LLMAsJudgeEvaluator(
    base_model="llama-3-3-70b-instruct",
    evaluator_model="anthropic.claude-3-haiku-20240307-v1:0",
    dataset="s3://my-bucket/dataset.jsonl",
    custom_metrics=custom_metrics,
    s3_output_path="s3://my-bucket/output"
)
execution = evaluator.evaluate()

# Example evaluating only custom model (skip base model)
evaluator = LLMAsJudgeEvaluator(
    base_model="llama-3-3-70b-instruct",
    evaluator_model="anthropic.claude-sonnet-4-5-20250929-v1:0",
    dataset="s3://my-bucket/my-dataset.jsonl",
    builtin_metrics=["Correctness"],  # Prefix optional
    evaluate_base_model=False,
    s3_output_path="s3://my-bucket/output"
)
execution = evaluator.evaluate()

base_eval_name: str | None#

base_model_name: str | None#

builtin_metrics: List[str] | None#

compute: Compute | HyperPodCompute | None#

custom_metrics: str | None#

dataset: str | Any#

evaluate(dry_run: bool = False)[source]#

Create and start an LLM-as-judge evaluation job.

This method initiates a 2-phase evaluation job:

Phase 1: Generate inference responses from base and custom models
Phase 2: Use judge model to evaluate responses with built-in and custom metrics

When the InspectAI path is active (custom model or Nova JumpStart model), Phase 1 runs inside an InspectAI container that generates inference responses and writes them to S3. Phase 2 remains unchanged — the LLMAJEvaluation judging step evaluates those responses with the judge model.

Parameters:: dry_run (bool) – If True, runs all validation (IAM, model resolution, data paths) without submitting the evaluation. Returns None on success, raises on validation failure. Defaults to False.
Returns:: The created LLM-as-judge evaluation execution, or None if dry_run=True.
Return type:: EvaluationPipelineExecution
Raises:: ValueError – If invalid model, dataset, or metric configurations are provided

Example

evaluator = LLMAsJudgeEvaluator(
    base_model="llama-3-3-70b-instruct",
    evaluator_model="anthropic.claude-sonnet-4-5-20250929-v1:0",
    dataset="s3://my-bucket/my-dataset.jsonl",
    builtin_metrics=["Correctness", "Helpfulness"],
    s3_output_path="s3://my-bucket/output"
)
execution = evaluator.evaluate()
execution.wait()

evaluate_base_model: bool#

evaluator_model: str#

classmethod get_all(session: Any | None = None, region: str | None = None)[source]#

Get all LLM-as-judge evaluation executions.

Uses EvaluationPipelineExecution.get_all() to retrieve all LLM-as-judge evaluation executions as an iterator.

Parameters:

session (Optional[Any]) – Optional boto3 session. If not provided, will be inferred.
region (Optional[str]) – Optional AWS region. If not provided, will be inferred.

Yields:

EvaluationPipelineExecution – LLM-as-judge evaluation execution instances

Example

# Get all LLM-as-judge evaluations as iterator
evaluations = LLMAsJudgeEvaluator.get_all()
all_executions = list(evaluations)

# Or iterate directly
for execution in LLMAsJudgeEvaluator.get_all():
    print(f"{execution.name}: {execution.status.overall_status}")

# With specific session/region
evaluations = LLMAsJudgeEvaluator.get_all(session=my_session, region='us-west-2')
all_executions = list(evaluations)

kms_key_id: str | None#

mlflow_experiment_name: str | None#

mlflow_resource_arn: str | None#

mlflow_run_name: str | None#

model: str | BaseTrainer | AgentRFTJob | ModelPackage#

model_config: ClassVar[ConfigDict] = {'arbitrary_types_allowed': True}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

model_package_group: str | ModelPackageGroup | None#

model_post_init(context: Any, /) → None#

This function is meant to behave like a BaseModel method to initialize private attributes.

It takes context as an argument since that’s what pydantic-core passes when calling it.

Parameters:

self – The BaseModel instance.
context – The context.

networking: VpcConfig | None#

overrides: Dict[str, Any] | None#

recipe: str | None#

region: str | None#

role: str | None#

s3_output_path: str#

sagemaker_session: Any | None#

tags: List[TagsDict] | None#

training_image: str | None#

Bases: BaseEvaluator

Evaluate a multi-turn RL agent model against a held-out prompt dataset.

The evaluator runs rollouts of the agent against an environment (Bedrock AgentCore runtime or a Lambda-wrapped agent) and computes aggregate metrics (pass@k, mean reward, etc.). Execution routes through SageMaker Pipelines using the new AgentRFT Job step type (JobCategory="AgentRFTEvaluation").

The evaluator supports three evaluation shapes, selected automatically based on the provided inputs:

Base-model only — pass a base model (JumpStart ID or ModelPackage) with an explicit agent_config.
Fine-tuned only — pass a MultiTurnRLTrainer or a fine-tuned ModelPackage; the evaluator extracts the source model package ARN and evaluates it only.
Base + fine-tuned comparison — pass evaluate_base_model=True along with a fine-tuned trainer / ModelPackage; both runs land in the same MLflow experiment for side-by-side comparison.

dataset#

Prompt dataset — S3 URI, hub-content DataSet ARN, or object exposing an .arn attribute. Required.

Type:: Union[str, Any]

agent_config#

Agent environment — Bedrock AgentCore ARN or Lambda ARN. Auto-resolved from a MultiTurnRLTrainer when provided as model.

Type:: Optional[Union[str, Any]]

agent_qualifier#

Bedrock AgentCore qualifier (e.g. "PROD"). Ignored when agent_config is a Lambda.

Type:: Optional[str]

accept_eula#

Forwarded to JobConfigDocument.EvaluationConfig.AcceptEula. Defaults to True (templates emit true unconditionally; flag kept for future backend schemas).

Type:: bool

evaluate_base_model#

When True and a fine-tuned model is present, render the comparison template (both base and fine-tuned are evaluated). Defaults to False — fine-tuned only.

Type:: bool

stopping_condition#

Maximum job duration in seconds. Default 86400 (24 hours); must be in (0, 259200].

Type:: int

tags#

Customer tags propagated to the pipeline + step Tags list.

Type:: Optional[List[Dict[str, str]]]

See :class:`BaseEvaluator` for inherited fields (``model``,

``s3_output_path``, ``mlflow_resource_arn``,

``mlflow_experiment_name``, ``networking``, ``kms_key_id``,

``model_package_group``, ``base_eval_name``, ``region``, ``role``,

``sagemaker_session``).

Example

from sagemaker.train.evaluate import MultiTurnRLEvaluator

# Evaluate a fine-tuned MTRL trainer output
evaluator = MultiTurnRLEvaluator(
    model=completed_mtrl_trainer,
    dataset='s3://my-bucket/eval-prompts.jsonl',
    s3_output_path='s3://my-bucket/mtrl-eval-output/',
)

execution = evaluator.evaluate()
execution.wait()
execution.show_results()

accept_eula: bool#

agent_config: Any | None#

agent_qualifier: str | None#

base_eval_name: str | None#

base_model_name: str | None#

compute: Compute | HyperPodCompute | None#

dataset: Any#

evaluate(dry_run: bool = False) → MTRLEvaluationExecution | None[source]#

Render the MTRL pipeline and start a non-blocking execution.

Parameters:: dry_run (bool) – If True, runs all validation (IAM, agent resolution, model resolution, template rendering) without submitting the evaluation. Returns None on success, raises on validation failure. Defaults to False.
Returns:: The started pipeline execution, or None if dry_run=True. Call .wait() to block until completion and .show_results() to render the aggregate report.
Return type:: MTRLEvaluationExecution

Example

execution = evaluator.evaluate()
execution.wait()
execution.show_results()

# Validate without submitting:
evaluator.evaluate(dry_run=True)

evaluate_base_model: bool#

classmethod get_all(session=None, region=None)[source]#

List all MTRL evaluation executions in the account / region.

Parameters:

session – Optional boto3 session.
region – Optional AWS region.

Yields:

EvaluationPipelineExecution – MTRL evaluation execution instances.

property hyperparameters#

Lazy-load evaluation hyperparameters from the JumpStart hub.

Returns a FineTuningOptions object exposing to_dict(), get_info(), and attribute-style read/write access with hub-sourced validation (type + range).

Supported parameters (sourced from the AgentRFT evaluation recipe): eval_group_size, sampling_temperature, top_p, max_tokens, pass_k_values, success_threshold.

Raises:: ValueError – If the base model name is not available or the hub does not expose an AgentRFTEvaluation override spec for the model.

kms_key_id: str | None#

static list_bedrock_agentcore_runtimes(session=None) → list[source]#

List Bedrock AgentCore runtimes.

Parameters:: session – Optional boto3 session.
Returns:: List of dicts, each with keys name, runtime_id, arn, and status.

static list_supported_models(session=None) → list[source]#

Return the list of models that support MTRL evaluation.

Queries SageMakerPublicHub to discover all models with MTRL recipes in their RecipeCollection.

Parameters:: session – Optional boto3 session.
Returns:: List of hub content model names supporting MTRL evaluation.

mlflow_experiment_name: str | None#

mlflow_resource_arn: str | None#

mlflow_run_name: str | None#

model: str | BaseTrainer | AgentRFTJob | ModelPackage#

model_config: ClassVar[ConfigDict] = {'arbitrary_types_allowed': True}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

model_package_group: str | ModelPackageGroup | None#

model_post_init(context: Any, /) → None#

This function is meant to behave like a BaseModel method to initialize private attributes.

It takes context as an argument since that’s what pydantic-core passes when calling it.

Parameters:

self – The BaseModel instance.
context – The context.

networking: VpcConfig | None#

overrides: Dict[str, Any] | None#

recipe: str | None#

region: str | None#

role: str | None#

s3_output_path: str#

sagemaker_session: Any | None#

stopping_condition: int#

tags: List[Dict[str, str]] | None#

training_image: str | None#

class sagemaker.train.evaluate.PipelineExecutionStatus(*, overall_status: str, step_details: ~typing.List[~sagemaker.train.evaluate.execution.StepDetail] = <factory>, failure_reason: str | None = None)[source]#

Bases: BaseModel

Combined pipeline execution status with step details and failure reason.

Aggregates the overall execution status along with detailed information about individual pipeline steps and any failure reasons.

Parameters:

overall_status (str) – Overall execution status (Starting, Executing, Completed, Failed, etc.).
step_details (List[StepDetail]) – List of individual pipeline step details.
failure_reason (Optional[str]) – Detailed reason if the execution failed.

failure_reason: str | None#

model_config: ClassVar[ConfigDict] = {}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

overall_status: str#

step_details: List[StepDetail]#

Bases: BaseModel

Pipeline step details for tracking execution progress.

Represents the status and timing information for a single step in a SageMaker pipeline execution.

Parameters:

name (str) – Name of the pipeline step.
status (str) – Status of the step (Completed, Executing, Waiting, Failed).
start_time (Optional[str]) – ISO format timestamp when step started.
end_time (Optional[str]) – ISO format timestamp when step ended.
display_name (Optional[str]) – Human-readable display name for the step.
failure_reason (Optional[str]) – Detailed reason if the step failed.

display_name: str | None#

end_time: str | None#

failure_reason: str | None#

job_arn: str | None#

model_config: ClassVar[ConfigDict] = {}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

name: str#

start_time: str | None#

status: str#

sagemaker.train.evaluate.get_benchmark_properties(benchmark: _Benchmark) → Dict[str, Any][source]#

Get properties for a specific benchmark.

This utility method returns the properties associated with a given benchmark as a dictionary, including information about modality, metrics, strategy, and available subtasks.

Parameters:

benchmark (_Benchmark) – The benchmark to get properties for (from get_benchmarks()).

Returns:

Dictionary containing benchmark properties with keys:

modality (str): The modality type (e.g., “Text”, “Multi-Modal”)

description (str): Description of the benchmark

metrics (list[str]): List of supported metrics

strategy (str): The evaluation strategy used

subtask_available (bool): Whether subtasks are supported

subtasks (Optional[list[str]]): List of available subtasks, if applicable

Return type:

Dict[str, Any]

Raises:

ValueError – If the provided benchmark is not found in the configuration.

Example

Benchmark = get_benchmarks()
props = get_benchmark_properties(Benchmark.MMLU)
print(props['description'])
# 'Multi-task Language Understanding – Tests knowledge across 57 subjects.'
print(props['subtasks'][:3])
# ['abstract_algebra', 'anatomy', 'astronomy']

Note

In the future, this will be extended to dynamically fetch benchmark properties from a backend API call instead of using the internal static configuration.

sagemaker.train.evaluate.get_benchmarks() → Type[_Benchmark][source]#

Get the Benchmark enum for selecting available benchmarks.

This utility method provides access to the internal Benchmark enum, allowing users to reference available benchmarks without directly accessing internal implementation details.

Returns:: The Benchmark enum class containing all available benchmarks.
Return type:: Type[_Benchmark]

Example

Benchmark = get_benchmarks()
evaluator = BenchMarkEvaluator(
    benchmark=Benchmark.MMLU,
    sagemaker_session=session,
    s3_output_path="s3://bucket/output"
)

Note

In the future, this will be extended to dynamically generate the enum from a backend API call to fetch the latest available benchmarks.

sagemaker.train.evaluate.get_builtin_metrics() → Type[_BuiltInMetric][source]#

Get the built-in metrics enum for custom scorer evaluation.

This utility function provides access to preset metrics for custom scorer evaluation.

Returns:: The built-in metric enum class
Return type:: Type[_BuiltInMetric]

Example

from sagemaker.train.evaluate import get_builtin_metrics

BuiltInMetric = get_builtin_metrics()
evaluator = CustomScorerEvaluator(
    evaluator=BuiltInMetric.PRIME_MATH,
    dataset=my_dataset,
    base_model="my-model",
    s3_output_path="s3://bucket/output",
    mlflow_resource_arn="arn:..."
)

Multi-Turn RL Evaluation#

MultiTurnRLEvaluator — evaluate MTRL agents on held-out prompts.

This module implements MultiTurnRLEvaluator, the SDK surface for evaluating Multi-Turn Reinforcement Learning (MTRL) agent models via the AgentRFT CreateJob pipeline step. Mirrors the architecture of sagemaker.train.evaluate.BenchMarkEvaluator, with MTRL-specific fields, validators, and the three-template rendering surface defined in sagemaker.train.evaluate.mtrl_pipeline_templates.

Bases: BaseEvaluator

Evaluate a multi-turn RL agent model against a held-out prompt dataset.

The evaluator runs rollouts of the agent against an environment (Bedrock AgentCore runtime or a Lambda-wrapped agent) and computes aggregate metrics (pass@k, mean reward, etc.). Execution routes through SageMaker Pipelines using the new AgentRFT Job step type (JobCategory="AgentRFTEvaluation").

The evaluator supports three evaluation shapes, selected automatically based on the provided inputs:

Base-model only — pass a base model (JumpStart ID or ModelPackage) with an explicit agent_config.
Fine-tuned only — pass a MultiTurnRLTrainer or a fine-tuned ModelPackage; the evaluator extracts the source model package ARN and evaluates it only.
Base + fine-tuned comparison — pass evaluate_base_model=True along with a fine-tuned trainer / ModelPackage; both runs land in the same MLflow experiment for side-by-side comparison.

dataset#

Prompt dataset — S3 URI, hub-content DataSet ARN, or object exposing an .arn attribute. Required.

Type:: Union[str, Any]

agent_config#

Agent environment — Bedrock AgentCore ARN or Lambda ARN. Auto-resolved from a MultiTurnRLTrainer when provided as model.

Type:: Optional[Union[str, Any]]

agent_qualifier#

Bedrock AgentCore qualifier (e.g. "PROD"). Ignored when agent_config is a Lambda.

Type:: Optional[str]

accept_eula#

Forwarded to JobConfigDocument.EvaluationConfig.AcceptEula. Defaults to True (templates emit true unconditionally; flag kept for future backend schemas).

Type:: bool

evaluate_base_model#

When True and a fine-tuned model is present, render the comparison template (both base and fine-tuned are evaluated). Defaults to False — fine-tuned only.

Type:: bool

stopping_condition#

Maximum job duration in seconds. Default 86400 (24 hours); must be in (0, 259200].

Type:: int

tags#

Customer tags propagated to the pipeline + step Tags list.

Type:: Optional[List[Dict[str, str]]]

See :class:`BaseEvaluator` for inherited fields (``model``,

``s3_output_path``, ``mlflow_resource_arn``,

``mlflow_experiment_name``, ``networking``, ``kms_key_id``,

``model_package_group``, ``base_eval_name``, ``region``, ``role``,

``sagemaker_session``).

Example

from sagemaker.train.evaluate import MultiTurnRLEvaluator

# Evaluate a fine-tuned MTRL trainer output
evaluator = MultiTurnRLEvaluator(
    model=completed_mtrl_trainer,
    dataset='s3://my-bucket/eval-prompts.jsonl',
    s3_output_path='s3://my-bucket/mtrl-eval-output/',
)

execution = evaluator.evaluate()
execution.wait()
execution.show_results()

accept_eula: bool#

agent_config: Any | None#

agent_qualifier: str | None#

base_eval_name: str | None#

base_model_name: str | None#

compute: Compute | HyperPodCompute | None#

dataset: Any#

evaluate(dry_run: bool = False) → MTRLEvaluationExecution | None[source]#

Render the MTRL pipeline and start a non-blocking execution.

Parameters:: dry_run (bool) – If True, runs all validation (IAM, agent resolution, model resolution, template rendering) without submitting the evaluation. Returns None on success, raises on validation failure. Defaults to False.
Returns:: The started pipeline execution, or None if dry_run=True. Call .wait() to block until completion and .show_results() to render the aggregate report.
Return type:: MTRLEvaluationExecution

Example

execution = evaluator.evaluate()
execution.wait()
execution.show_results()

# Validate without submitting:
evaluator.evaluate(dry_run=True)

evaluate_base_model: bool#

classmethod get_all(session=None, region=None)[source]#

List all MTRL evaluation executions in the account / region.

Parameters:

session – Optional boto3 session.
region – Optional AWS region.

Yields:

EvaluationPipelineExecution – MTRL evaluation execution instances.

property hyperparameters#

Lazy-load evaluation hyperparameters from the JumpStart hub.

Returns a FineTuningOptions object exposing to_dict(), get_info(), and attribute-style read/write access with hub-sourced validation (type + range).

Supported parameters (sourced from the AgentRFT evaluation recipe): eval_group_size, sampling_temperature, top_p, max_tokens, pass_k_values, success_threshold.

Raises:: ValueError – If the base model name is not available or the hub does not expose an AgentRFTEvaluation override spec for the model.

kms_key_id: str | None#

static list_bedrock_agentcore_runtimes(session=None) → list[source]#

List Bedrock AgentCore runtimes.

Parameters:: session – Optional boto3 session.
Returns:: List of dicts, each with keys name, runtime_id, arn, and status.

static list_supported_models(session=None) → list[source]#

Return the list of models that support MTRL evaluation.

Queries SageMakerPublicHub to discover all models with MTRL recipes in their RecipeCollection.

Parameters:: session – Optional boto3 session.
Returns:: List of hub content model names supporting MTRL evaluation.

mlflow_experiment_name: str | None#

mlflow_resource_arn: str | None#

mlflow_run_name: str | None#

model: str | BaseTrainer | AgentRFTJob | ModelPackage#

model_config: ClassVar[ConfigDict] = {'arbitrary_types_allowed': True}#: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

model_package_group: str | ModelPackageGroup | None#

model_post_init(context: Any, /) → None#

This function is meant to behave like a BaseModel method to initialize private attributes.

It takes context as an argument since that’s what pydantic-core passes when calling it.

Parameters:

self – The BaseModel instance.
context – The context.

networking: VpcConfig | None#

overrides: Dict[str, Any] | None#

recipe: str | None#

region: str | None#

role: str | None#

s3_output_path: str#

sagemaker_session: Any | None#

stopping_condition: int#

tags: List[Dict[str, str]] | None#

training_image: str | None#

SageMaker Train

Contents

SageMaker Train#

Model Training#

Distributed Training#

Multi-Turn RL Training#

Interactive Training Sessions#

Model Evaluation#

Multi-Turn RL Evaluation#