2026/10/12 - SageMaker Training Session Runtime - 6 new api methods
Changes Amazon SageMaker Training Session Runtime provides APIs for executing operations within active training sessions, including asynchronous training operations, sampling and inference requests, and streaming real-time model responses, enabling programmatic control over the training loop.
Retrieves the current status and result of a previously submitted training operation. Poll this operation until the status reaches a terminal state (Completed, Failed, or Skipped).
See also: AWS API Documentation
Request Syntax
client.get_training_operation_result(
trainingSessionName='string',
operationId='string'
)
string
[REQUIRED]
The name of the training session that owns the operation.
string
[REQUIRED]
The unique identifier of the operation to retrieve, returned by ExecuteTrainingOperation.
dict
Response Syntax
{
'status': 'Pending'|'InProgress'|'Completed'|'Failed'|'Skipped',
'operationResult': {
'forward': {
'lossFnOutputType': 'string',
'lossFnOutputs': [
{
'string': {
'data': {...}|[...]|123|123.4|'string'|True|None,
'dtype': 'string',
'shape': [
123,
]
}
},
],
'metrics': {
'string': 123.0
},
'tokenUsage': {
'trainingTokens': 123
}
},
'forwardBackward': {
'lossFnOutputType': 'string',
'lossFnOutputs': [
{
'string': {
'data': {...}|[...]|123|123.4|'string'|True|None,
'dtype': 'string',
'shape': [
123,
]
}
},
],
'metrics': {
'string': 123.0
},
'tokenUsage': {
'trainingTokens': 123
}
},
'optimizerStep': {
'metrics': {
'string': 123.0
}
},
'saveWeights': {
'modelPackageArn': 'string'
},
'saveState': {
'modelPackageArn': 'string'
},
'loadState': {}
},
'errorDetails': 'string',
'tokenUsage': {
'trainingTokens': 123
}
}
Response Structure
(dict) --
status (string) --
The current status of the operation. Possible values are Pending, InProgress, Completed, Failed, or Skipped.
operationResult (dict) --
The result payload of the operation. The service populates this field only when the status is Completed. The populated member of the union corresponds to the operation type that was submitted.
forward (dict) --
The result from a forward-only pass.
lossFnOutputType (string) --
The data type of the loss function output tensors.
lossFnOutputs (list) --
The loss function output tensors, keyed by name.
(dict) --
(string) --
(dict) --
A tensor with data, data type, and shape metadata.
data (:ref:`document<document>`) --
The raw tensor data as JSON. The format is determined by the dtype field: integer arrays for int64, floating-point arrays for float32.
dtype (string) --
The data type of the tensor elements, such as int64 or float32.
shape (list) --
The dimensions of the tensor, such as [batch_size, seq_length].
(integer) --
metrics (dict) --
The training metrics emitted by the engine, such as loss values.
(string) --
(float) --
tokenUsage (dict) --
The token usage statistics for metering.
trainingTokens (integer) --
The number of training tokens consumed by the operation.
forwardBackward (dict) --
The result from a forward and backward pass.
lossFnOutputType (string) --
The data type of the loss function output tensors.
lossFnOutputs (list) --
The loss function output tensors, keyed by name.
(dict) --
(string) --
(dict) --
A tensor with data, data type, and shape metadata.
data (:ref:`document<document>`) --
The raw tensor data as JSON. The format is determined by the dtype field: integer arrays for int64, floating-point arrays for float32.
dtype (string) --
The data type of the tensor elements, such as int64 or float32.
shape (list) --
The dimensions of the tensor, such as [batch_size, seq_length].
(integer) --
metrics (dict) --
The training metrics emitted by the engine, such as loss values.
(string) --
(float) --
tokenUsage (dict) --
The token usage statistics for metering.
trainingTokens (integer) --
The number of training tokens consumed by the operation.
optimizerStep (dict) --
The result from an optimizer step.
metrics (dict) --
The training metrics emitted by the engine after the optimizer step.
(string) --
(float) --
saveWeights (dict) --
The result from saving model weights.
modelPackageArn (string) --
The ARN of the model package created for the saved weights.
saveState (dict) --
The result from saving full training state.
modelPackageArn (string) --
The ARN of the model package created for the saved state checkpoint.
loadState (dict) --
The result from loading training state.
errorDetails (string) --
A message describing why the operation failed. The service populates this field only when the status is Failed.
tokenUsage (dict) --
The token usage statistics for metering purposes. The service populates this field for forward and forwardBackward operations upon completion.
trainingTokens (integer) --
The number of training tokens consumed by the operation.
Sends an inference request to the model associated with the specified training session and returns the response synchronously. The service passes the request and response bodies to and from the model without modification.
See also: AWS API Documentation
Request Syntax
client.sample_training_session(
trainingSessionName='string',
contentType='string',
accept='string',
modelPackageArn='string',
body=b'bytes'|file
)
string
[REQUIRED]
The name of the training session whose model to invoke.
string
The MIME type of the input data in the request body.
string
The desired MIME type of the inference body in the response.
string
The ARN of the model package whose checkpoint to invoke.
bytes or seekable file-like object
The raw inference request body.
dict
Response Syntax
{
'contentType': 'string',
'body': b'bytes'|file
}
Response Structure
(dict) --
contentType (string) --
The MIME type of the inference result in the response body.
body (bytes or seekable file-like object) --
The raw inference response body from the model.
Retrieves the current status and result of a previously submitted sampling operation. Poll this operation until the status reaches a terminal state (Completed or Failed).
See also: AWS API Documentation
Request Syntax
client.get_sampling_operation_result(
trainingSessionName='string',
operationId='string'
)
string
[REQUIRED]
The name of the training session that owns the sample operation.
string
[REQUIRED]
The unique identifier of the sample operation to retrieve, returned by ExecuteSamplingOperation.
dict
Response Syntax
{
'status': 'Pending'|'InProgress'|'Completed'|'Failed'|'Skipped',
'operationResult': {
'sample': {
'sequences': [
{
'stopReason': 'stop'|'length',
'tokens': [
123,
],
'logprobs': [
123.0,
],
'tokenUsage': {
'prefill': 123,
'decode': 123
}
},
],
'promptLogprobs': [
123.0,
],
'topkPromptLogprobs': [
[
{
'tokenId': 123,
'logprob': 123.0
},
],
]
},
'computeLogprobs': {
'logprobs': [
123.0,
]
},
'loadWeightsForSampling': {}
},
'errorDetails': 'string'
}
Response Structure
(dict) --
status (string) --
The current status of the sample operation. Possible values are Pending, InProgress, Completed, or Failed.
operationResult (dict) --
The result payload of the sample operation. The service populates this field only when the status is Completed.
sample (dict) --
The generated text samples from the model.
sequences (list) --
The list of generated sequences.
(dict) --
A single generated sequence from the model.
stopReason (string) --
The reason why generation stopped. A value of stop indicates a stop sequence was encountered. A value of length indicates the maximum token limit was reached.
tokens (list) --
The token IDs of the generated sequence.
(integer) --
logprobs (list) --
The per-token log-probabilities for the generated tokens.
(float) --
tokenUsage (dict) --
The token usage statistics for this individual sample.
prefill (integer) --
The number of tokens in the input prompt (prefill phase).
decode (integer) --
The number of tokens generated (decode phase).
promptLogprobs (list) --
The per-token log-probabilities for the input prompt. Populated only when includePromptLogprobs is true.
(float) --
topkPromptLogprobs (list) --
The top-k log-probabilities per prompt token, as a list of lists. Populated only when topkPromptLogprobs is greater than 0.
(list) --
(dict) --
The log-probability for a single token.
tokenId (integer) --
The vocabulary token ID.
logprob (float) --
The log-probability of the token.
computeLogprobs (dict) --
The computed log-probabilities for the prompt tokens.
logprobs (list) --
The per-token log-probabilities for the input prompt.
(float) --
loadWeightsForSampling (dict) --
The result from loading weights for sampling.
errorDetails (string) --
A message describing why the sample operation failed. The service populates this field only when the status is Failed.
Streams inference responses from the model associated with the specified training session. Use this operation for real-time token-by-token generation. The service passes the request body to the model without modification.
See also: AWS API Documentation
Request Syntax
client.sample_training_session_with_response_stream(
trainingSessionName='string',
contentType='string',
accept='string',
modelPackageArn='string',
body=b'bytes'|file
)
string
[REQUIRED]
The name of the training session whose model to invoke.
string
The MIME type of the input data in the request body.
string
The desired MIME type of the inference body in the response.
string
The ARN of the model package whose checkpoint to invoke.
bytes or seekable file-like object
The raw inference request body.
dict
Response Syntax
{
'contentType': 'string',
'body': StreamingBody()
}
Response Structure
(dict) --
contentType (string) --
The MIME type of the streaming inference result.
body (:class:`.StreamingBody`) --
The streaming inference response from the model, delivered as a series of chunks.
Submits an asynchronous sampling operation to generate text or compute log-probabilities from the model associated with the specified training session. The session must be in Running or Blocked state. Use GetSamplingOperationResult to poll for the result.
See also: AWS API Documentation
Request Syntax
client.execute_sampling_operation(
trainingSessionName='string',
operationConfig={
'sample': {
'prompt': {
'chunks': [
{
'tokens': [
123,
],
'type': 'encoded_text'
},
]
},
'numSamples': 123,
'samplingParams': {
'maxTokens': 123,
'temperature': 123.0,
'topK': 123,
'topP': 123.0,
'seed': 123,
'stop': [
'string',
]
},
'includePromptLogprobs': True|False,
'topkPromptLogprobs': 123,
'modelPackageArn': 'string'
},
'computeLogprobs': {
'prompt': {
'chunks': [
{
'tokens': [
123,
],
'type': 'encoded_text'
},
]
},
'modelPackageArn': 'string'
},
'loadWeightsForSampling': {
'modelPackageArn': 'string'
}
}
)
string
[REQUIRED]
The name of the training session to sample from.
dict
[REQUIRED]
The configuration for the sample operation. Specify either a sample configuration to generate text or a computeLogprobs configuration to compute token log-probabilities.
sample (dict) --
The configuration for generating text samples from the model.
prompt (dict) -- [REQUIRED]
The input prompt to send to the model for text generation.
chunks (list) -- [REQUIRED]
The ordered list of encoded text chunks forming the input sequence.
(dict) --
A chunk of pre-tokenized text represented as a list of token IDs.
tokens (list) -- [REQUIRED]
The token IDs representing the encoded text.
(integer) --
type (string) -- [REQUIRED]
The type of the chunk encoding.
numSamples (integer) -- [REQUIRED]
The number of independent samples to generate for the prompt.
samplingParams (dict) -- [REQUIRED]
The parameters controlling the sampling strategy.
maxTokens (integer) -- [REQUIRED]
The maximum number of tokens to generate per sample.
temperature (float) --
The sampling temperature controlling randomness. Higher values produce more diverse outputs.
topK (integer) --
The number of highest-probability vocabulary tokens to keep for top-k filtering. Set to -1 to disable.
topP (float) --
The cumulative probability threshold for nucleus (top-p) sampling. Only tokens with cumulative probability below this threshold are considered.
seed (integer) --
The random seed for reproducible sampling. When set, the same seed with the same input produces deterministic outputs.
stop (list) --
The list of stop sequences that terminate generation when encountered in the output.
(string) --
includePromptLogprobs (boolean) --
Specifies whether to include per-token log-probabilities for the input prompt in the response.
topkPromptLogprobs (integer) --
The number of top-k log-probabilities to return per prompt token.
modelPackageArn (string) --
The ARN of the model package whose checkpoint this sample should run against. Must be a valid SageMaker model package ARN in the format arn:aws:sagemaker:REGION:ACCOUNT:model-package/PACKAGE_NAME/VERSION. When absent, the sample runs against the session's latest checkpoint.
computeLogprobs (dict) --
The configuration for computing log-probabilities of prompt tokens without generating new text.
prompt (dict) -- [REQUIRED]
The input prompt whose token log-probabilities will be computed.
chunks (list) -- [REQUIRED]
The ordered list of encoded text chunks forming the input sequence.
(dict) --
A chunk of pre-tokenized text represented as a list of token IDs.
tokens (list) -- [REQUIRED]
The token IDs representing the encoded text.
(integer) --
type (string) -- [REQUIRED]
The type of the chunk encoding.
modelPackageArn (string) --
The ARN of the model package whose checkpoint this operation should run against. Must be a valid SageMaker model package ARN in the format arn:aws:sagemaker:REGION:ACCOUNT:model-package/PACKAGE_NAME/VERSION. When absent, the operation runs against the session's latest checkpoint.
loadWeightsForSampling (dict) --
The configuration for loading model weights for inference and sampling without restoring optimizer state.
modelPackageArn (string) -- [REQUIRED]
The ARN of the model package containing the weights to load for sampling. Must be a valid SageMaker model package ARN in the format arn:aws:sagemaker:REGION:ACCOUNT:model-package/PACKAGE_NAME/VERSION. The model package must belong to the same AWS account as the training session.
dict
Response Syntax
{
'operationId': 'string'
}
Response Structure
(dict) --
operationId (string) --
The unique identifier for the submitted sample operation. Use this identifier with GetSamplingOperationResult to retrieve the operation status and result.
Submits an asynchronous training operation to the specified training session. The service executes operations in sequence number order. The session must be in Running or Blocked state. Use GetTrainingOperationResult to poll for the result. Supported operation types: forward, forwardBackward, optimizerStep, saveWeights, saveState, and loadState.
See also: AWS API Documentation
Request Syntax
client.execute_training_operation(
trainingSessionName='string',
sequenceNumber=123,
clientToken='string',
operationConfig={
'forward': {
'lossFnName': 'cross_entropy'|'importance_sampling'|'ppo'|'cispo',
'datumList': [
{
'modelInput': {
'chunks': [
{
'tokens': [
123,
],
'type': 'encoded_text'
},
]
},
'lossFnInputs': {
'string': {
'data': {...}|[...]|123|123.4|'string'|True|None,
'dtype': 'string',
'shape': [
123,
]
}
}
},
],
'lossFnConfig': {...}|[...]|123|123.4|'string'|True|None
},
'forwardBackward': {
'lossFnName': 'cross_entropy'|'importance_sampling'|'ppo'|'cispo',
'datumList': [
{
'modelInput': {
'chunks': [
{
'tokens': [
123,
],
'type': 'encoded_text'
},
]
},
'lossFnInputs': {
'string': {
'data': {...}|[...]|123|123.4|'string'|True|None,
'dtype': 'string',
'shape': [
123,
]
}
}
},
],
'lossFnConfig': {...}|[...]|123|123.4|'string'|True|None
},
'optimizerStep': {
'adamParamsConfig': {
'learningRate': 123.0,
'beta1': 123.0,
'beta2': 123.0,
'eps': 123.0,
'weightDecay': 123.0,
'gradClipNorm': 123.0,
'loraPlusLrRatio': 123.0
}
},
'saveWeights': {
'mergeWeights': True|False,
'description': 'string'
},
'saveState': {
'description': 'string'
},
'loadState': {
'modelPackageArn': 'string',
'loadOptimizer': True|False
}
}
)
string
[REQUIRED]
The name of the training session to submit the operation to.
integer
[REQUIRED]
The monotonically increasing sequence number that determines execution order. The training session processes operations strictly in sequence number order. If a gap exists, the session blocks until the missing operation arrives or times out. Must be at least 1.
string
A unique, case-sensitive identifier that you provide to ensure that the operation completes no more than one time. If you retry a request with the same client token and the same parameters, the service ignores the request but does not return an error.
dict
[REQUIRED]
The configuration that specifies the type and parameters of the training operation. Exactly one member of the union must be set.
forward (dict) --
The configuration for a forward-only pass through the model without gradient computation.
lossFnName (string) -- [REQUIRED]
The name of the loss function to evaluate during the forward pass.
datumList (list) -- [REQUIRED]
The list of input data to feed through the model. Must contain at least one datum, and each datum must have a non-empty modelInput with at least one chunk.
(dict) --
A single training datum containing model input and optional loss function inputs.
modelInput (dict) -- [REQUIRED]
The tokenized input to feed to the model.
chunks (list) -- [REQUIRED]
The ordered list of encoded text chunks forming the input sequence.
(dict) --
A chunk of pre-tokenized text represented as a list of token IDs.
tokens (list) -- [REQUIRED]
The token IDs representing the encoded text.
(integer) --
type (string) -- [REQUIRED]
The type of the chunk encoding.
lossFnInputs (dict) --
The tensor inputs for the loss function, keyed by name. Unsupported keys are rejected. Supported keys: target_tokens, weights, advantages, logprobs, ref_log_prob, temperature. Required keys depend on the loss function type: SFT ( cross_entropy) requires target_tokens and weights; RL ( ppo, cispo, importance_sampling) requires target_tokens, weights, advantages, and logprobs.
(string) --
(dict) --
A tensor with data, data type, and shape metadata.
data (:ref:`document<document>`) -- [REQUIRED]
The raw tensor data as JSON. The format is determined by the dtype field: integer arrays for int64, floating-point arrays for float32.
dtype (string) -- [REQUIRED]
The data type of the tensor elements, such as int64 or float32.
shape (list) --
The dimensions of the tensor, such as [batch_size, seq_length].
(integer) --
lossFnConfig (:ref:`document<document>`) --
Additional loss function configuration passed to the engine. Supported keys are clip_low_threshold (valid range: 0 to 1) and clip_high_threshold (valid range: 1 to 10). Unsupported keys are rejected.
forwardBackward (dict) --
The configuration for a forward and backward pass that computes gradients.
lossFnName (string) -- [REQUIRED]
The name of the loss function to evaluate during the forward-backward pass.
datumList (list) -- [REQUIRED]
The list of input data to feed through the model. Must contain at least one datum, and each datum must have a non-empty modelInput with at least one chunk.
(dict) --
A single training datum containing model input and optional loss function inputs.
modelInput (dict) -- [REQUIRED]
The tokenized input to feed to the model.
chunks (list) -- [REQUIRED]
The ordered list of encoded text chunks forming the input sequence.
(dict) --
A chunk of pre-tokenized text represented as a list of token IDs.
tokens (list) -- [REQUIRED]
The token IDs representing the encoded text.
(integer) --
type (string) -- [REQUIRED]
The type of the chunk encoding.
lossFnInputs (dict) --
The tensor inputs for the loss function, keyed by name. Unsupported keys are rejected. Supported keys: target_tokens, weights, advantages, logprobs, ref_log_prob, temperature. Required keys depend on the loss function type: SFT ( cross_entropy) requires target_tokens and weights; RL ( ppo, cispo, importance_sampling) requires target_tokens, weights, advantages, and logprobs.
(string) --
(dict) --
A tensor with data, data type, and shape metadata.
data (:ref:`document<document>`) -- [REQUIRED]
The raw tensor data as JSON. The format is determined by the dtype field: integer arrays for int64, floating-point arrays for float32.
dtype (string) -- [REQUIRED]
The data type of the tensor elements, such as int64 or float32.
shape (list) --
The dimensions of the tensor, such as [batch_size, seq_length].
(integer) --
lossFnConfig (:ref:`document<document>`) --
Additional loss function configuration passed to the engine. Supported keys are clip_low_threshold (valid range: 0 to 1) and clip_high_threshold (valid range: 1 to 10). Unsupported keys are rejected.
optimizerStep (dict) --
The configuration for an optimizer step that updates model weights using accumulated gradients.
adamParamsConfig (dict) -- [REQUIRED]
The Adam optimizer hyperparameters for this step.
learningRate (float) -- [REQUIRED]
The learning rate for the optimizer. Must be greater than 0.
beta1 (float) --
The exponential decay rate for the first moment estimates.
beta2 (float) --
The exponential decay rate for the second moment estimates.
eps (float) --
A small constant added to the denominator for numerical stability.
weightDecay (float) --
The weight decay (L2 regularization) coefficient.
gradClipNorm (float) --
The maximum gradient norm for gradient clipping. Gradients are scaled down if their norm exceeds this value.
loraPlusLrRatio (float) --
The learning rate multiplier for LoRA+ B matrices relative to A matrices. When set, B matrices use learningRate * loraPlusLrRatio while A matrices use the base learning rate.
saveWeights (dict) --
The configuration for saving model weights as a checkpoint.
mergeWeights (boolean) --
Specifies whether to merge LoRA adapter weights into the base model before saving. When true, the saved checkpoint contains the full merged model weights.
description (string) --
A human-readable description for the saved checkpoint.
saveState (dict) --
The configuration for saving full training state, including model weights and optimizer state.
description (string) --
A human-readable description for the saved state checkpoint.
loadState (dict) --
The configuration for loading training state from a previously saved checkpoint.
modelPackageArn (string) -- [REQUIRED]
The ARN of the model package containing the checkpoint to load. Must be a valid SageMaker model package ARN in the format arn:aws:sagemaker:REGION:ACCOUNT:model-package/PACKAGE_NAME/VERSION. Model package group ARNs are not accepted.
loadOptimizer (boolean) -- [REQUIRED]
Specifies whether to also restore the optimizer state alongside model weights. When true, optimizer state is restored to resume training from where it left off.
dict
Response Syntax
{
'operationId': 'string'
}
Response Structure
(dict) --
operationId (string) --
The unique identifier for the submitted operation. Use this identifier with GetTrainingOperationResult to retrieve the operation status and result.