amazon-sagemaker-examples Documentation

Example ๐Ÿ““ Jupyter notebooks that demonstrate how to build, train, and deploy machine learning models using ๐Ÿง  Amazon SageMaker.

11,000 stars Jupyter Notebook 10 files ยท ~8,615 tokens #aws#data-science#deep-learning#examples#inference#jupyter-notebook#machine-learning#mlops
RAW Doc

Technical Documentation: aws/amazon-sagemaker-examples

โ„น๏ธ Provenance: Multi-Tier Ingestion ยท CodeWiki Reference

Overview (`README.md`)

:exclamation::fire: Announcing SageMaker-Core: A New Python SDK for Amazon SageMaker :fire::exclamation:

Introduction

Today, Amazon SageMaker is excited to announce the release of SageMaker-Core, a new Python SDK that provides an object-oriented interface for interacting with SageMaker resources such as TrainingJob, Model, and Endpoint. This SDK introduces the resource chaining feature, allowing developers to pass resource objects as parameters, eliminating manual parameter specification and simplifying code management. SageMaker-Core abstracts low-level details like resource state transitions and polling logic, achieving full parity with SageMaker APIs. It also includes usability improvements such as auto code completion, comprehensive documentation, and type hints, enhancing the overall developer experience.

Use Case

SageMaker-Core is ideal for ML practitioners who seek full customization of AWS primitives for their ML workloads. SageMaker-Core is an improvement over Boto3, providing a more intuitive and efficient way to manage SageMaker resources. By providing an intuitive object-oriented interface and resource chaining, the SDK allows for seamless integration and management of SageMaker resources. This flexibility, combined with intelligent defaults enables developers to tailor their ML workloads according to their needs. Comprehensive documentation, and type hints help developers write code faster and with fewer errors without navigating complex API documentation.

Call to Action

To learn more about SageMaker-Core, visit the documentation and example notebooks. Get started today by integrating SageMaker-Core into your machine learning workflows and experience the benefits of a streamlined and efficient development process.

Amazon SageMaker Examples

Example Jupyter notebooks that demonstrate how to build, train, and deploy machine learning models using Amazon SageMaker.

:books: Read this before you proceed further

Amazon SageMaker examples are divided in two repositories:

  • SageMaker example notebooks is the official repository, containing examples that demonstrate the usage of Amazon SageMaker. This repository is entirely focussed on covering the breadth of features provided by SageMaker, and is maintained directly by the Amazon SageMaker team.

  • Sagemaker Example Community repository is another SageMaker repository which contains additional examples and reference solutions, beyond the examples showcased in the official repository. This repository is maintained by community of engineers and solution architects at AWS.

Planning to submit a PR to this repository? Read this first:

  • This repository will only accept notebooks/examples which demonstrate a feature of SageMaker, not yet covered anywhere in this repository. PR submitters are requested to check this before submitting the PR to avoid getting it rejected.

  • If you still would like to contribute your example, please submit a PR to Sagemaker Example Community repository instead.

:hammer_and_wrench: Setup

The quickest setup to run example notebooks includes:

:computer: Usage

These example notebooks are automatically loaded into SageMaker Notebook Instances.
They can be accessed by clicking on the SageMaker Examples tab in Jupyter or the SageMaker logo in JupyterLab.

Although most examples utilize key Amazon SageMaker functionality like distributed, managed training or real-time hosted endpoints, these notebooks can be run outside of Amazon SageMaker Notebook Instances with minimal modification (updating IAM role definition and installing the necessary libraries).

:notebook: Example Notebook Categories

The notebooks are organized by ML capability, following the lifecycle order of a typical
project: train a model, customize it, evaluate it, deploy it, then operationalize it.
Every notebook is self-contained: it lives with the scripts, data and images it needs, and
all of its references are relative to its own folder.

[Training](https://github.com/aws/amazon-sagemaker-examples/tree/default/%20%20%20%20%20%20training)

Amazon SageMaker Training is a fully managed service that helps you train ML models at
scale. It containerizes your workload and manages the AWS compute for you, so you can focus
on the model rather than the infrastructure. These examples train with ModelTrainer from
SageMaker Python SDK v3 and cover script and framework training, distributed training,
managed spot training with checkpointing, heterogeneous clusters, bringing your own
container, submitting work through AWS Batch training queues, and running local code as a
training job with @remote.

[Model Customization](https://github.com/aws/amazon-sagemaker-examples/tree/default/%20%20%20%20%20model_customization)

Model customization adapts a pre-trained foundation model to your data and your task, which
is usually far cheaper and faster than training from scratch. These examples cover the
fine-tuning techniques SageMaker supports โ€” supervised fine-tuning (SFT), direct preference
optimization (DPO), reinforcement learning with verifiable rewards (RLVR), RL from AI
feedback (RLAIF), multi-turn RL (MTRL) and continued pre-training (CPT) โ€” along with recipe
overrides, data mixing, JumpStart fine-tuning with a private model hub, and distributed
fine-tuning across serverless, serverful and HyperPod compute.

[Evaluation](https://github.com/aws/amazon-sagemaker-examples/tree/default/%20%20%20%20evaluation)

Before you promote a customized model you need to know whether it actually got better.
These examples use the SageMaker evaluator surface to score models: standard benchmark
evaluation, your own scoring function, an LLM acting as a judge, and Inspect AI. They show
how to evaluate a base model and a fine-tuned model on the same footing so the comparison
is meaningful.

[Inference](https://github.com/aws/amazon-sagemaker-examples/tree/default/%20%20%20inference)

With Amazon SageMaker you can get predictions from your trained models through a broad set
of deployment options, and scale them without taking on the operational burden yourself.
These examples deploy with ModelBuilder and sagemaker-core, and cover real-time,
serverless and asynchronous endpoints; in-process and local-container modes for fast
iteration; inference pipelines; multi-model and multi-container endpoints; safe rollout
through A/B testing, shadow variants and deployment guardrails; autoscaling with inference
components; latency and cost optimization with optimize(); batch transform; deployment to
Amazon Bedrock; AWS Marketplace model packages; and human review with Augmented AI.

[MLOps](https://github.com/aws/amazon-sagemaker-examples/tree/default/%20%20mlops)

MLOps applies DevOps practices to machine learning: continuous integration and deployment,
reproducibility, lineage and governance. It accounts for the parts of an AI/ML project that
ordinary software delivery does not, which shortens delivery time, reduces defects and makes
data science more productive. These examples cover SageMaker Pipelines with the Model
Registry, lineage tracking, experiment tracking, Feature Store, processing jobs, EMR
Serverless steps, bias and explainability analysis with Clarify, and MLflow.

:balance_scale: License

This library is licensed under the Apache 2.0 License.
For more details, please take a look at the LICENSE file.

:handshake: Contributing

Although we're extremely excited to receive contributions from the community, we're still working on the best mechanism to take in examples from external sources. Please bear with us in the short-term if pull requests take longer than expected or are closed.
Please read our contributing guidelines
if you'd like to open an issue or submit a pull request.


Overview (` training/README.md`)

Amazon SageMaker Examples

Training

Train models on Amazon SageMaker with the v3 SDK: ModelTrainer for script and framework training, distributed training, managed spot with checkpointing, heterogeneous clusters, bring-your-own-container, AWS Batch training queues, and @remote local-code execution.


Overview (` model_customization/README.md`)

Amazon SageMaker Examples

Model Customization

Fine-tune and customize foundation models: SFT, DPO, RLVR, RLAIF, MTRL and CPT trainers, recipe overrides, data mixing, JumpStart fine-tuning with a private model hub, and distributed fine-tuning across serverless, serverful and HyperPod compute.


Overview (` evaluation/README.md`)

Amazon SageMaker Examples

Evaluation

Evaluate models with the v3 evaluator surface: benchmark evaluation, custom scorers, LLM-as-a-judge, and Inspect AI.


Overview (` inference/README.md`)

Amazon SageMaker Examples

Inference

Deploy and serve models with ModelBuilder and sagemaker-core: real-time, serverless and async endpoints, in-process and local-container modes, inference pipelines, multi-model and multi-container endpoints, A/B testing, shadow variants, deployment guardrails, autoscaling with inference components, model optimization, batch transform, Bedrock deployment, AWS Marketplace model packages and Augmented AI.


Overview (` mlops/README.md`)

Amazon SageMaker Examples

MLOps

Operationalize ML on SageMaker: Pipelines with Model Registry, lineage tracking, experiment tracking, Feature Store, processing jobs, EMR Serverless steps, Clarify, and MLflow.


Overview (` model_customization/sm-training_sessions_grpo_gsm8k/README.md`)

GRPO on GSM8K with Amazon SageMaker AI training session APIs

This example writes a GRPO reinforcement learning loop against SageMaker AI training session
APIs. It has two notebooks: the first trains a LoRA adapter for GPT-OSS 20B on GSM8K math
problems, and the second deploys the trained checkpoint to a SageMaker real-time endpoint.

About training session APIs

Fine-tuning foundation models usually forces a tradeoff between control and operational
simplicity. Managed fine-tuning services handle the infrastructure but limit you to a fixed set of
techniques. Working at a lower level gives you algorithmic freedom, but you also take on GPU
clusters, distributed training frameworks, and fault recovery.

A training session is a live, GPU-backed session that you drive from your own Python process.
Your loop decides each step, and SageMaker AI handles the GPUs, base model loading, LoRA adapter
registration, distributed execution, fault recovery, and checkpoint management. You can implement
any training algorithm: supervised fine-tuning, reinforcement learning with custom reward
functions, multi-stage pipelines, or new techniques. The session is stateful, so the model stays
in memory across steps, and RL workflows can sample from the current policy without a separate
inference endpoint.

How it works

You write the training logic in Python, in a notebook or on your local machine, using three
primitives:

  • forward_backward computes gradients over a batch
  • optim_step applies the optimizer update
  • sample generates model responses from the current policy

The loss function is an argument to each training call: cross_entropy for supervised
fine-tuning, and ppo, importance_sampling, or cispo for RL. The same session and API run
either kind of training.

This example uses all three for RL. In each step, the notebook samples several responses per
question with sample, scores them locally against the reference answer, converts the scores
into group-relative advantages, and runs forward_backward with the ppo loss followed by
optim_step. It then calls save_weights_and_get_sampling_client so the next step samples from
the updated policy.

Quickstart

  1. Install the dependencies. requirements.txt installs the SageMaker Python SDK v3.25.0 or
    later, which includes the training session APIs (sagemaker.train.training_session):

    bash
    pip install -r requirements.txt
  2. Create the IAM execution role:

    bash
    python infrastructure/setup_infra.py --region us-west-2

    setup_infra.py creates service-role/SageMakerTrainingSessionRole, trusted by the training
    sessions service and SageMaker AI Inference. It attaches the
    AmazonSageMakerTrainingSessionFullAccess managed policy and an inline policy for hosting the
    deployed model (JumpStart cache read, Amazon ECR image pull, CloudWatch logs). The role is
    created under /service-role/ because the AmazonSageMakerTrainingSessionAccess caller policy
    only allows passing roles under that path. If an earlier version of this example created the
    role at the root path, delete it with python infrastructure/cleanup_infra.py --delete-roles
    and run setup_infra.py again.

  3. Run sm-training_sessions_1_train_grpo_gsm8k.ipynb, then
    sm-training_sessions_2_deploy_grpo_checkpoint.ipynb.

The notebooks resolve the Region, account, execution role, and default S3 bucket at runtime, and
create the model package group on first run.

The core API

python
from sagemaker.train.training_session import ServiceClient, AdamParams, SamplingParams

svc = ServiceClient(region=..., role_arn=..., model_package_group_arn=..., s3_output_path=...)
tc = svc.create_lora_training_client("openai/gpt-oss-20b", rank=32, alpha=64)
sc = tc.create_sampling_client()

# Sample from the current policy
out = sc.sample(prompt_tokens, num_samples=8, sampling_params=SamplingParams(max_tokens=512)).result()

# Train
fb = tc.forward_backward(datums, loss_fn="ppo")         # returns a future
opt = tc.optim_step(AdamParams(learning_rate=1e-5))     # pipelined behind it
fb.result(), opt.result()

# Save current weights and sample from them
sc = tc.save_weights_and_get_sampling_client()
logprobs = sc.compute_logprobs(prompt_tokens).result()

# Or sample any saved checkpoint directly
sc = svc.create_sampling_client(model_package_arn=...)

# Checkpoints
tc.save_weights_for_sampler(merge_weights=True).result()   # merged, deployable model
print(tc.save_state().result().model_package_arn)          # resumable state

tc.stop()

Every GPU-bound call returns a future immediately, so a step can submit all of its sampling
requests before waiting on any of them. Each operation also has an *_async variant, such as
sample_async and forward_backward_async.

Contents

text
.
โ”œโ”€โ”€ infrastructure/
โ”‚   โ”œโ”€โ”€ setup_infra.py                                   # create the IAM execution role
โ”‚   โ””โ”€โ”€ cleanup_infra.py                                 # remove what the example created
โ”œโ”€โ”€ data/                                                # GSM8K splits, written by the training notebook
โ”œโ”€โ”€ sm-training_sessions_1_train_grpo_gsm8k.ipynb        # GRPO training on GSM8K
โ”œโ”€โ”€ sm-training_sessions_2_deploy_grpo_checkpoint.ipynb  # deploy the trained checkpoint
โ””โ”€โ”€ requirements.txt

sm-training_sessions_1_train_grpo_gsm8k.ipynb opens a session, measures baseline accuracy on
held-out problems, runs the RL loop, measures accuracy again, and saves two checkpoints: a merged
model for deployment and a resumable training state.

sm-training_sessions_2_deploy_grpo_checkpoint.ipynb finds the merged checkpoint, hosts it on a
real-time endpoint, invokes it, and deletes the endpoint.

Prerequisites

  • Linux or macOS, Python 3.10, 3.11, or 3.12
  • SageMaker Python SDK v3.25.0 or later
  • Quota for training session APIs in US East (N. Virginia), US West (Oregon), or Asia Pacific
    (Mumbai)
  • AWS credentials:
    • to run setup_infra.py: permissions to create an IAM role and attach policies to it
    • to run the training notebook: the AmazonSageMakerTrainingSessionAccess managed policy
    • to run the deploy notebook: SageMaker AI hosting permissions (create and invoke models,
      endpoint configurations, endpoints, and inference components) and iam:PassRole on the
      execution role for sagemaker.amazonaws.com

Clean up

If you run both notebooks to completion, the training notebook stops the session and the deploy
notebook deletes the endpoint and its resources. The S3 outputs, the registered checkpoints, and
the IAM role remain.

bash
python infrastructure/cleanup_infra.py                    # list only, deletes nothing
python infrastructure/cleanup_infra.py --stop-sessions
python infrastructure/cleanup_infra.py --delete-endpoints <endpoint-name>
python infrastructure/cleanup_infra.py --all              # sessions, S3, model package group, role

A training session holds GPU resources while it is running, so stop it when you are done.

Further reading

  • SageMaker training sessions cookbook:
    complete training loops built on training session APIs, including supervised fine-tuning on chat
    datasets, RL on GSM8K, RL on competitive programming graded by sandboxed code execution, and
    KTO preference alignment, with checkpoint resume, validation, and MLflow metrics.
  • SageMaker Python SDK: source for the
    sagemaker.train.training_session module.
  • SageMaker AI model customization:
    the managed serverless customization options.

Overview (` inference/sm-a_b_testing/README.md`)

A/B Testing with Amazon SageMaker

In production ML workflows, data scientists and data engineers frequently try to improve their models in various ways, such as by performing Perform Automatic Model Tuning, training on additional or more-recent data, and improving feature selection. Performing A/B testing between a new model and an old model with production traffic can be an effective final step in the validation process for a new model. In A/B testing, you test different variants of your models and compare how each variant performs relative to each other. You then choose the best-performing model to replace a previously-existing model new version delivers better performance than the previously-existing version.

Amazon SageMaker enables you to test multiple models or model versions behind the same endpoint using production variants. Each production variant identifies a machine learning (ML) model and the resources deployed for hosting the model. You can distribute endpoint invocation requests across multiple production variants by providing the traffic distribution for each variant, or you can invoke a specific variant directly for each request.

In this notebook we'll:

  • Evaluate models by invoking specific variants
  • Gradually release a new model by specifying traffic distribution

Overview (` training/sm-scikit_build_your_own_container/container/ReadMe.md`)

Bring-your-own Algorithm Sample

This example shows how to package an algorithm for use with SageMaker. We have chosen a simple scikit-learn implementation of decision trees to illustrate the procedure.

SageMaker supports two execution modes: training where the algorithm uses input data to train a new model and serving where the algorithm accepts HTTP requests and uses the previously trained model to do an inference (also called "scoring", "prediction", or "transformation").

The algorithm that we have built here supports both training and scoring in SageMaker with the same container image. It is perfectly reasonable to build an algorithm that supports only training or scoring as well as to build an algorithm that has separate container images for training and scoring.v

In order to build a production grade inference server into the container, we use the following stack to make the implementer's job simple:

  1. nginx is a light-weight layer that handles the incoming HTTP requests and manages the I/O in and out of the container efficiently.
  2. gunicorn is a WSGI pre-forking worker server that runs multiple copies of your application and load balances between them.
  3. flask is a simple web framework used in the inference app that you write. It lets you respond to call on the /ping and /invocations endpoints without having to write much code.

The Structure of the Sample Code

The components are as follows:

  • Dockerfile: The Dockerfile describes how the image is built and what it contains. It is a recipe for your container and gives you tremendous flexibility to construct almost any execution environment you can imagine. Here. we use the Dockerfile to describe a pretty standard python science stack and the simple scripts that we're going to add to it. See the Dockerfile reference for what's possible here.

  • build_and_push.sh: The script to build the Docker image (using the Dockerfile above) and push it to the Amazon EC2 Container Registry (ECR) so that it can be deployed to SageMaker. Specify the name of the image as the argument to this script. The script will generate a full name for the repository in your account and your configured AWS region. If this ECR repository doesn't exist, the script will create it.

  • decision-trees: The directory that contains the application to run in the container. See the next session for details about each of the files.

  • local-test: A directory containing scripts and a setup for running a simple training and inference jobs locally so that you can test that everything is set up correctly. See below for details.

The application run inside the container

When SageMaker starts a container, it will invoke the container with an argument of either train or serve. We have set this container up so that the argument in treated as the command that the container executes. When training, it will run the train program included and, when serving, it will run the serve program.

  • train: The main program for training the model. When you build your own algorithm, you'll edit this to include your training code.
  • serve: The wrapper that starts the inference server. In most cases, you can use this file as-is.
  • wsgi.py: The start up shell for the individual server workers. This only needs to be changed if you changed where predictor.py is located or is named.
  • predictor.py: The algorithm-specific inference server. This is the file that you modify with your own algorithm's code.
  • nginx.conf: The configuration for the nginx master server that manages the multiple workers.

Setup for local testing

The subdirectory local-test contains scripts and sample data for testing the built container image on the local machine. When building your own algorithm, you'll want to modify it appropriately.

  • train-local.sh: Instantiate the container configured for training.
  • serve-local.sh: Instantiate the container configured for serving.
  • predict.sh: Run predictions against a locally instantiated server.
  • test-dir: The directory that gets mounted into the container with test data mounted in all the places that match the container schema.
  • payload.csv: Sample data for used by predict.sh for testing the server.

The directory tree mounted into the container

The tree under test-dir is mounted into the container and mimics the directory structure that SageMaker would create for the running container during training or hosting.

  • input/config/hyperparameters.json: The hyperparameters for the training job.
  • input/data/training/leaf_train.csv: The training data.
  • model: The directory where the algorithm writes the model file.
  • output: The directory where the algorithm can write its success or failure file.

Environment variables

When you create an inference server, you can control some of Gunicorn's options via environment variables. These
can be supplied as part of the CreateModel API call.

text
Parameter                Environment Variable              Default Value
---------                --------------------              -------------
number of workers        MODEL_SERVER_WORKERS              the number of CPU cores
timeout                  MODEL_SERVER_TIMEOUT              60 seconds

Overview (` model_customization/estimator-pytorch-gpu/smtj_batch_utils/README.md`)

Create AWS Batch resources for SageMaker Training job

Define AWS Batch Queue for SageMaker Training Job

To create your batch queues, edit the file config.py. The configuration system uses a dynamic naming convention that automatically discovers and creates resources based on attribute names.

Configuration Naming Convention

The system automatically detects resources using these naming patterns:

Service Environments:

  • *_SE_NAME - The name of the service environment
  • *_SE_CREATE_REQUEST - The creation request dictionary

Scheduling Policies:

  • *_SP_NAME - The name of the scheduling policy
  • *_SP_CREATE_REQUEST - The creation request dictionary

Job Queues:

  • *_QUEUE_NAME - The name of the job queue
  • *_QUEUE_CREATE_REQUEST - The creation request dictionary

Example Configuration:

python
class ServiceEnvs:
    C5_SE_NAME = "ml-c5-xlarge-fifo"
    C5_SE_CREATE_REQUEST = {
        "serviceEnvironmentName": C5_SE_NAME,
        "serviceEnvironmentType": "SAGEMAKER_TRAINING",
        "state": "ENABLED",
        "capacityLimits": [{"maxCapacity": 3, "capacityUnit": "NUM_INSTANCES"}]
    }

class SchedulingPolicies:
    G6_SP_NAME = "ml-g6-scheduling-policy"
    G6_SP_CREATE_REQUEST = {
        "name": G6_SP_NAME,
        "fairsharePolicy": {
            "shareDecaySeconds": 3600,
            "shareDistribution": [
                {"shareIdentifier": "HIGHPRI", "weightFactor": 1}
            ]
        }
    }

class TrainingJobQueues:
    G6_TJ_QUEUE_NAME = "ml-g6-queue"
    G6_TJ_QUEUE_CREATE_REQUEST = {
        "jobQueueName": G6_TJ_QUEUE_NAME,
        "jobQueueType": "SAGEMAKER_TRAINING",
        "state": "ENABLED",
        "priority": 1,
        "serviceEnvironmentOrder": [{"order": 1, "serviceEnvironment": "ml-g6-service-env"}],
        "schedulingPolicyArn": "ml-g6-scheduling-policy"  # References SP_NAME
    }

Queue Types for SageMaker Training

  1. FIFO Queue - No scheduling policy required
  2. Fair-Share Queue - Requires scheduling policy ARN (reference the SP_NAME)

The system automatically links scheduling policies to queues by matching the policy name.

Create AWS Batch resources

To create your Batch job queues and service environments using boto3, run the create_resources.py script from a terminal as shown below.

text
python3 create_resources.py

You can view your new resources from the CLI with the following commands.


--- METRICS ---