Prerequisites - Amazon SageMaker AI
Managed policiesPass the execution roleExecution role trust policyCreate a model package groupOther setup
Services or capabilities described in AWS documentation might vary by Region. To see the differences applicable to the AWS European Sovereign Cloud Region, see the AWS European Sovereign Cloud User Guide.

Prerequisites

The training session APIs extend serverless model customization with the SageMaker AI CreateTrainingSession API and a set of runtime operations that run training and sampling against an active session. Before you create your first session, configure three things: the IAM identity that calls the API (the caller identity), the IAM execution role that the session assumes (the RoleArn you pass to CreateTrainingSession), and a model package group for your checkpoints.

Managed policies

The training session APIs provide two AWS managed policies. Attach each to the appropriate identity before you create a session. For the full permission lists, see AWS managed policies for SageMaker AI training session APIs.

  • AmazonSageMakerTrainingSessionAccess – attach to the IAM users and roles that create and interact with training sessions. It grants control-plane actions (create, describe, stop, delete, list), the runtime operations that run training and sampling, model package management for checkpoint publication, iam:PassRole to hand the execution role to the service, and the AWS KMS permissions needed to use a customer managed key through SageMaker AI.

  • AmazonSageMakerTrainingSessionFullAccess – attach to the IAM role you pass as the session execution role. It grants the session the permissions it needs while running: read and write access to your training data and checkpoints in Amazon S3, AWS KMS permissions for Amazon S3-side encryption, access to SageMaker AI hub content and model packages, and Amazon EC2 network interface management for sessions that run in Amazon VPC mode.

Pass the execution role

Allow the caller identity to pass the execution role to the Training Sessions service principal. The AmazonSageMakerTrainingSessionAccess managed policy already includes this permission, scoped to roles under the service-role path whose names contain SageMaker. If you use a custom caller policy, include the following statement:

{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Sid": "PassExecutionRolePermissions",
      "Effect": "Allow",
      "Action": "iam:PassRole",
      "Resource": "arn:aws:iam::*:role/service-role/*SageMaker*",
      "Condition": {
        "StringEquals": {
          "iam:PassedToService": "training-session.sagemaker.amazonaws.com"
        }
      }
    }
  ]
}

Execution role trust policy

Add the Training Sessions service principal to the trust policy of your execution role so that the service can assume it:

{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Effect": "Allow",
      "Principal": {
        "Service": "training-session.sagemaker.amazonaws.com"
      },
      "Action": ["sts:AssumeRole", "sts:TagSession"]
    }
  ]
}

Create a model package group

The service registers each checkpoint you save as a model package in a SageMaker AI model package group. Create a group for your checkpoints and pass its ARN as model_package_group_arn when you create the ServiceClient:

aws sagemaker create-model-package-group \
    --model-package-group-name my-training-session-checkpoints \
    --region us-west-2

For more information, see Create a Model Group.

Other setup