Bring your own model with Amazon SageMaker AI: Script mode in SDK v3

In 2021, we published Bring your own model with Amazon SageMaker script mode. That post showed how to use script mode on managed framework containers from AWS to write custom training and inference code. Script mode was a leap forward: you didn’t need to build or maintain Docker images to run your own algorithm on Amazon SageMaker AI.

The v3 SDK delivers a redesign from scratch that makes many workflows like the bring-your-own-model workflow even more streamlined. The new SDK replaces framework-specific estimator classes (SKLearn, PyTorch, XGBoost) with a single, unified ModelTrainer for training and ModelBuilder for deployment.

In v3, the SDK syncs a local source code directory into the training job at runtime using the new SourceCode configuration object. You bring a container image from Amazon Elastic Container Registry (Amazon ECR): one you build, an AWS Deep Learning Container, or a third-party image. The SDK handles injecting your code at runtime.

This means:

• Faster iterations: Change your training script, rerun. No container rebuild necessary.
• Full container control: Install system packages or CUDA libraries in your image. The SDK doesn’t assume what’s inside.
• One API for multiple frameworks: Whether you’re training with frameworks like scikit-learn, PyTorch, Stable Diffusion, or a custom C++ inference binary, the interface is identical.

Solution overview

In this post, we walk through two end-to-end examples that demonstrate how script mode works in the SageMaker Python SDK v3:

1. Train and deploy a scikit-learn Random Forest – a classic tabular machine learning (ML) workflow that trains on the diabetes dataset and deploys to a real-time endpoint using Deep Java Library (DJL) Serving, a high-performance model server.
2. Fine-tune Stable Diffusion 3.5 with LoRA – a generative AI workflow that uses Hugging Face Accelerate for multi-GPU distributed training.

Both examples use the same two core classes:

• ModelTrainer replaces the v2 Estimator family. Configures and launches a SageMaker training job.
• ModelBuilder replaces the v2 Model/Predictor pattern. Packages your inference handler and deploys to an endpoint.

A key concept is the SourceCode object. It accepts a source_dir (a path to your local code directory) and either a command string (for training) or an entry_script (for inference). At job launch, SageMaker syncs this directory into the container, and your code runs inside the container without being baked into the image.

You can find the example code for this blog post in the GitHub repository.

What changed from SDK v2 to v3?

The following table summarizes the architectural shift:

Prerequisites

To follow along, you need:

• An AWS account with Amazon SageMaker AI access.
• An AWS Identity and Access Management (IAM) execution role with Amazon SageMaker AI and Amazon Simple Storage Service (Amazon S3) permissions.
• The SageMaker Python SDK v3 installed (pip install sagemaker>=3.0).
• A training container image pushed to Amazon ECR (this post shows an example of building and pushing a container to Amazon ECR).
• An Amazon S3 bucket for training data and model artifacts.
• (Optional) An MLflow app or tracking server on Amazon SageMaker AI for experiment tracking.
• (Optional) If you plan to build and run the example containers from a JupyterLab space in Amazon SageMaker Studio rather than a local machine, Docker access must be enabled at the domain level. For details, see Local mode support in Amazon SageMaker Studio.

Example 1: Train and deploy a scikit-learn model

Let’s start with a classic ML workflow. We train a Random Forest classifier on the diabetes dataset and deploy it to a real-time SageMaker endpoint.

Step 1: Building the Docker container

The training container is intentionally minimal. It contains only the runtime and framework libraries and no training code, so that we can reuse it for other scikit-learn models we might want to build.

The complete Dockerfile for our scikit-learn container is:

FROM python:3.13-slim

RUN apt-get update && apt-get install -y 
    build-essential jq git 
    && rm -rf /var/lib/apt/lists/*

COPY requirements.txt .

RUN pip install -r requirements.txt --no-cache-dir

The container is a stable, version-controlled runtime environment. The algorithm-specific code lives in your source_dir, and the SDK injects it at runtime.

Build this container once, push it to Amazon ECR, and iterate on your training code as many times as you want without touching Docker again.

The example notebook includes Docker build and push commands by using two shell scripts:

./build.sh --env .env.docker.sklearn
./push.sh --env .env.docker.sklearn

Note that you need Docker installed on the environment you’re using to run the code samples. If you’re running this on a JupyterLab space within Amazon SageMaker AI, you need to enable Docker on the domain-level settings.

Step 1a: Configuration

First, we auto-detect the account-level configuration. Note that we omit the import statements required in the following code snippet for brevity, but the full code is available in the GitHub repository.

...
boto_session = boto3.Session()
sm_session = Session(boto_session=boto_session)
AWS_REGION = boto_session.region_name
AWS_ACCOUNT_ID = boto3.client("sts").get_caller_identity()["Account"]
SAGEMAKER_EXECUTION_ROLE = sm_session.get_caller_identity_arn() # works for role or user
S3_BUCKET = sm_session.default_bucket()

In the following snippet, we point TRAINING_IMAGE_URI at the container we built ourselves in the previous step. This gives you full control over installed packages and runtime versions. For deployment, we show that you can also use a pre-existing managed DJL framework container if you don’t want to build your own. For more information about pre-built containers, see available Deep Learning Containers images.

# Your custom training image pushed to ECR
TRAINING_IMAGE_URI = (
    f"{AWS_ACCOUNT_ID}.dkr.ecr.{AWS_REGION}.amazonaws.com/sklearn:latest"
)

MODEL_OUTPUT_S3_PATH = f"s3://{S3_BUCKET}/random-forest/model-output"

Optionally, if you’d like to track hyperparameters, metrics, and model artifacts across training runs, the example training script is already instrumented for fully managed MLflow on Amazon SageMaker AI. Set the MLFLOW_ARN and MLFLOW_EXPERIMENT_NAME variables in the following snippet to automatically enable logging. This is an optional step, and you can set the values to None to skip experiment tracking instead.

MLFLOW_ARN = "arn:aws:sagemaker:{AWS_REGION}:{AWS_ACCOUNT_ID}:mlflow-app/{XYZ}"
MLFLOW_EXPERIMENT_NAME = "random-forest-experiment"

Step 2: Launch a training job with ModelTrainer

The SourceCode object takes your local source_dir and a command string. At job launch, SageMaker syncs the entire source_dir into the container and runs your command. This decouples your code from your container image. Change the script, re-launch with no container rebuild needed.

source_code = SourceCode(
    source_dir="./train/random_forest",
    command=(
        "python random_forest.py"
        " --n_jobs 4 --max_depth 10 --n_estimators 120"
        f" --mlflow_arn {MLFLOW_ARN}"
        f" --mlflow_experiment_name {MLFLOW_EXPERIMENT_NAME}"
    ),
)

compute = Compute(
    instance_type="ml.m5.2xlarge",
    instance_count=1,
    volume_size_in_gb=30,
    keep_alive_period_in_seconds=3600, # Warm pool for faster re-runs
)

stopping_condition = StoppingCondition(max_runtime_in_seconds=3600)
output_config = OutputDataConfig(s3_output_path=MODEL_OUTPUT_S3_PATH)

model_trainer = ModelTrainer(
    training_image=TRAINING_IMAGE_URI,
    source_code=source_code,
    compute=compute,
    output_data_config=output_config,
    stopping_condition=stopping_condition,
    role=SAGEMAKER_EXECUTION_ROLE,
    base_job_name="random-forest-training",
    sagemaker_session=Session(),
)

model_trainer.train(wait=True)

A few things to note:

• source_dir can contain your files such as utility modules, config files, and shell scripts, which are synced into the container.
• command is a shell command that runs inside the container. You can call a Python script, a bash script, or anything else your container supports.
• keep_alive_period_in_seconds turns on SageMaker warm pools. The instance stays warm for 1 hour, so iterative re-runs launch in seconds rather than minutes.
• OutputDataConfig sets the S3 destination where SageMaker uploads your training results when the job finishes. Anything your script saves to /opt/ml/model (the SM_MODEL_DIR environment variable) is packaged as model.tar.gz under this path, and that’s the model artifact we deploy in Step 3. For more information, see Using the SageMaker training and inference toolkits for the folder structure and the environment variables that SageMaker sets.

Step 3: Deploy to a real-time endpoint with ModelBuilder

After training completes, we deploy the model artifact to a SageMaker real-time endpoint. ModelBuilder packages your inference handler, repacks it with the model artifact, and creates the endpoint in a few lines. You can use the metadata from the training job to find the S3 path of the final model artifact, then supply that to the ModelBuilder object. For serving, we use a pre-built AWS Deep Learning Container rather than building a custom one, though you can bring your own if needed. For other pre-built containers, see available Deep Learning Containers images.

# Locate the trained model artifact
sm_session = Session()
sm_client = sm_session.boto_session.client("sagemaker")
training_job_desc = sm_client.describe_training_job(
    TrainingJobName=training_job_name
)
model_artifact_s3_uri = training_job_desc["ModelArtifacts"]["S3ModelArtifacts"]

# Define inference source code
inference_source_code = SourceCode(
    source_dir="./deploy/random_forest",
    entry_script="inference.py",
)

# Build and deploy
model_builder = ModelBuilder(
    image_uri=INFERENCE_IMAGE_URI,
    model_server=ModelServer.DJL_SERVING,
    source_code=inference_source_code,
    s3_model_data_url=model_artifact_s3_uri,
    env_vars={"OPTION_ENTRYPOINT": "code/inference.py"},
    sagemaker_session=Session(),
    role_arn=SAGEMAKER_EXECUTION_ROLE,
)

model_builder.build(model_name="random-forest-endpoint", mode=Mode.SAGEMAKER_ENDPOINT)

The build() step assembles a deployable model without launching any infrastructure. ModelBuilder takes your inference handler and model artifact and packages them together according to the conventions of your chosen model server (here, DJL Serving). It then registers a SageMaker model that points at your inference image and repacked artifact in Amazon S3. ModelBuilder can also do more than we illustrate here, such as auto-selecting a container, auto-capturing dependencies, and generating serialization code from a raw framework model. For more information, see Create a model in Amazon SageMaker AI with ModelBuilder.

With the model built, we call deploy() to stand up the real-time endpoint, which returns an Endpoint interface:

predictor = model_builder.deploy(
    endpoint_name="random-forest-endpoint",
    initial_instance_count=1,
)

Notice the same SourceCode pattern for inference: point at a local directory containing your handler and specify the entry_script. The SDK repacks the handler into the model archive so DJL Serving can find it at runtime.

A few notes on the preceding code snippets:

• The inference.py script implements a single handle(inputs) function per the DJL Python mode documentation, which SageMaker calls for every request. When the inference worker first starts up, an empty request is sent to the handler to complete a one-time model loading process. The concept is to load once and map future requests to a prediction. By default, the model is located at /opt/ml/model, which corresponds to the SM_MODEL_DIR environment variable. After the initial model loading, for each incoming request the inference script determines the Content-Type, deserializes the payload, and returns the prediction result as a JSON object.
• Load the model once during cold start and reuse it across requests, because reloading per request adds latency to every call. Also validate the request’s Content-Type so the endpoint rejects unexpected input with a clear, immediate error.
• The model_server argument tells ModelBuilder which serving runtime to package your model for and run inside the endpoint. The model server is the process that loads your model, exposes the endpoints SageMaker expects, and dispatches each request to your handler. This is why it corresponds to how our inference.py is written. Here, we choose ModelServer.DJL_SERVING, a flexible, high-performance server well-suited to general Python inference and large-model serving. This is also why our handler follows the DJL handle(inputs) contract described earlier. For other model serving choices exposed by ModelServer, see the ModelServer API reference.
• The mode parameter controls where your model runs. Here we use Mode.SAGEMAKER_ENDPOINT, which deploys to a fully managed real-time endpoint. ModelBuilder also supports Mode.LOCAL_CONTAINER (run in a Docker container on your machine) and Mode.IN_PROCESS (run directly in your current Python process) for testing and iterating on your handler locally.

In this case, we deploy to a real-time endpoint. Depending on your workload, you can host a single model on its own endpoint or pack multiple models behind one endpoint using inference components, so you can allocate resources and scale each model independently. For more information, see Real-time inference and Inference components.

Step 4: Test the endpoint

Send a sample CSV request to confirm the endpoint is healthy:

sample_csv = "6,148,72,35,0,33.6,0.627,50"
response = predictor.invoke(
    body=sample_csv.encode("utf-8"),
    content_type="text/csv",
)
print("Prediction:", response.body.read().decode("utf-8"))
# The model returns a class label: 1 = tested_positive (diabetes), 0 = tested_negative.

Example 1 covers a traditional ML use case, but this same pattern also works for generative AI use cases and for distributed training if needed, as we explore in the following example.

Example 2: Fine-tune Stable Diffusion 3.5 with LoRA

The same primitives used in the previous example can be extended for more complex training scenarios, including multi-GPU or multi-node generative AI jobs. In this example, we fine-tune Stable Diffusion 3.5 Medium using LoRA (Low-Rank Adaptation) on a custom image/caption dataset. The training job uses Hugging Face Accelerate for multi-GPU distributed training across 4 A10G GPUs on an ml.g5.12xlarge instance.

Step 1: Building the Docker container

As in the scikit-learn example, the container is purely a runtime environment. The complete Dockerfile for our Stable Diffusion container is:

FROM pytorch/pytorch:2.7.1-cuda12.8-cudnn9-devel

RUN apt-get update && apt-get install -y 
    build-essential jq git 
    && rm -rf /var/lib/apt/lists/*

COPY requirements.txt .

RUN pip install -r requirements.txt --no-cache-dir

As with the preceding container, the example notebook includes Docker build and push commands by using two shell scripts:

./build.sh --env .env.docker.stablediffusion
./push.sh --env .env.docker.stablediffusion

The requirements include the deep learning stack (PyTorch, diffusers, transformers, accelerate, PEFT, DeepSpeed) but again, no training scripts. The LoRA fine-tuning logic, Accelerate launcher script, recipe configs, and orchestration code live in source_dir and are synced at runtime:

train/stable_diffusion/
├── base.sh                        # Launcher: pip installs, GPU detection, accelerate launch
├── train_text_to_image_lora.py    # Main training script
├── requirements.txt               # Runtime deps installed by base.sh
├── recipes/
│   └── default-medium-g5_12x.yaml # Hyperparameter recipe
└── accelerate_configs/
    └── ddp.yaml                   # Distributed training config

You can swap recipes, adjust the LoRA rank, change the base model, or modify the training loop code by editing your local files, without rebuilding the Docker container.

Step 1a: Prepare the training data

In this example, we follow a slightly different paradigm for training data preparation to demonstrate the flexibility of SageMaker Training Jobs. In the previous example, the training script fetches the training data at runtime without staging it in Amazon S3. However, in this example, we retrieve the dreambooth dataset from Hugging Face using load_dataset and populate it into our working bucket, which we then pass into the training job as InputData, as shown in the following code:

from datasets import DatasetDict, load_dataset

LOCAL_DATA_DIR = "/tmp/sd-training-data"
dataset = load_dataset("google/dreambooth", "dog", split="train")
...
_bucket, _prefix = SD_TRAINING_DATA_S3_PATH.replace("s3://", "").split("/", 1)
s3 = boto_session.client("s3")
for root, _, files in os.walk(LOCAL_DATA_DIR):
    for fname in files:
        local_path = os.path.join(root, fname)
        s3_key = f"{_prefix}/{os.path.relpath(local_path, LOCAL_DATA_DIR)}"
        s3.upload_file(local_path, _bucket, s3_key)
...
sd_input_config = [
    InputData(channel_name="train", data_source=SD_TRAINING_DATA_S3_PATH),
]

The channel_name you assign in InputData controls where SageMaker stages that data inside the training container. At job startup, SageMaker automatically downloads the contents of each channel to /opt/ml/input/data/<channel_name> and exposes the path through a matching SM_CHANNEL_<CHANNEL_NAME> environment variable.

Here we define a single train channel, so the dataset lands at /opt/ml/input/data/train. However, channels are fully customizable. You can define multiple channels (up to 20 per training job) and name them whatever fits your workflow. For example, you can create separate train, validation, and test channels, where each is staged into its own directory automatically. Your training script can then reference data by a stable local path without hardcoding any S3 locations. For more information about defining and accessing input data channels, see the SageMaker input data documentation.

Step 2: Launch a training job with ModelTrainer

The command launches a bash script (base.sh) that detects the GPU count, runs accelerate launch, fetches secret values, and kicks off the training script. The hyperparameters live in a YAML recipe file. With this approach, you can tune the model training parameters by editing the recipe, not the container:

sd_source_code = SourceCode(
    source_dir="./train/stable_diffusion",
    command=(
        "/bin/bash base.sh"
        " --config recipes/default-medium-g5_12x.yaml"
        " --training-script train_text_to_image_lora.py"
        " --accelerate-config accelerate_configs/ddp.yaml"
        f" --mlflow-arn {SD_MLFLOW_ARN}"
        f" --mlflow-experiment-name {SD_MLFLOW_EXPERIMENT_NAME}"
    ),
)

sd_compute = Compute(
    instance_type="ml.g5.12xlarge", # 4x A10G GPUs
    instance_count=1,
    volume_size_in_gb=30,
    keep_alive_period_in_seconds=3600,
)

At this point, we’re ready to create the ModelTrainer class following a similar paradigm as the previous example. However, we now include a SECRETS_ARN corresponding to an entry in AWS Secrets Manager, which contains our Hugging Face token, required to download the gated Stable Diffusion model. The GitHub repository contains a sample AWS CloudFormation template you can use to deploy your own secret, fetch the Amazon Resource Name (ARN), and populate it into the following snippet. This is a more secure approach than including those sensitive inputs (such as Hugging Face tokens) in plain text or in environment variables.

sd_model_trainer = ModelTrainer(
    training_image=SD_TRAINING_IMAGE_URI,
    source_code=sd_source_code,
    compute=sd_compute,
    output_data_config=sd_output_config,
    stopping_condition=sd_stopping_condition,
    role=SAGEMAKER_EXECUTION_ROLE,
    environment={"SECRETS_ARN": SECRETS_ARN},
    base_job_name=SD_JOB_NAME,
    sagemaker_session=Session(),
)

sd_model_trainer.train(input_data_config=input_config, wait=True)

When the container starts up, the base.sh script contains logic to retrieve the secret by its ARN, parse the values, and set them as environment variables for future use. In this approach, we do not expose the sensitive values in our notebook or logs.

Step 3: Deploy to a real-time endpoint with ModelBuilder

After the training step completes, we locate the trained LoRA weights, create a ModelBuilder object, and deploy our fine-tuned model to a real-time endpoint, following a similar pattern as the previous example:

sd_inference_source_code = SourceCode(
    source_dir="./deploy/stable_diffusion",
    entry_script="inference.py",
)

sd_model_builder = ModelBuilder(
    image_uri=SD_INFERENCE_IMAGE_URI,
    model_server=ModelServer.DJL_SERVING,
    source_code=sd_inference_source_code,
    s3_model_data_url=sd_model_artifact_s3_uri,
    instance_type="ml.g5.4xlarge",
    env_vars={
        "OPTION_ENGINE": "Python",
        "OPTION_ENTRYPOINT": "code/inference.py",
        "SECRETS_ARN": SECRETS_ARN,
    },
    sagemaker_session=Session(),
    role_arn=SAGEMAKER_EXECUTION_ROLE,
)

Step 4: Test the endpoint

Lastly, we test by sending a sample request to confirm the endpoint is healthy:

payload = {
    "prompt": "a boy Malcom and his dog Ben",
    "num_inference_steps": 30,
    "guidance_scale": 7.5,
    "seed": 42,
}
response = sd_predictor.invoke(
    body=json.dumps(payload).encode("utf-8"),
    content_type="application/json",
)
result = json.loads(response.body.read().decode("utf-8"))
image = Image.open(io.BytesIO(base64.b64decode(result["generated_image"])))
image

Key takeaways from this example:

• Bash launchers work well – your command can be a shell command, not only python script.py. This is recommended for multi-step launchers that set up Accelerate, install runtime deps, or orchestrate distributed training.
• Recipe-driven training – hyperparameters, model IDs, and LoRA settings live in YAML recipe files inside source_dir. Change hyperparameters without touching the container or the training script.
• Secrets with AWS Secrets Manager – store Hugging Face tokens, API keys, or other secrets in AWS Secrets Manager and pass the corresponding ARN through the environment parameter. Then, handle parsing and environment configuration from within your training or deployment pipeline. Note that the continuous integration and continuous delivery (CI/CD) or Training Job principal need to have permissions to read from AWS Secrets Manager.
• Same API, different scale – the interface is identical whether you’re training a random forest on a CPU instance or fine-tuning a diffusion model on multi-GPU.

Clean up

To avoid incurring future charges, delete the resources you created:

1. Delete the SageMaker endpoint:

   from sagemaker.core.resources import Endpoint, EndpointConfig, Model
   
   Endpoint.get(endpoint_name=ENDPOINT_NAME).delete()
   EndpointConfig.get(endpoint_config_name=ENDPOINT_NAME).delete()
   Model.get(model_name=ENDPOINT_NAME).delete()

2. Delete the S3 training data and model artifacts if they are no longer needed.
3. Delete the ECR container images if you no longer need them.
4. (Optional) Delete the MLflow tracking server if you created one for this walkthrough.

Conclusion

The SageMaker Python SDK v3 re-imagines script mode for the modern ML practitioner. The core principles remain the same. Bring your own training and inference code, run it on managed infrastructure, and let SageMaker handle the undifferentiated heavy lifting. What’s new in v3 is how quickly you can go from code to a running training job and inference endpoint:

• One API for multiple workloads – ModelTrainer and ModelBuilder replace a dozen framework-specific classes. Less to learn, less to maintain.
• Code-container decoupling – SourceCode syncs your local code directory into the container at runtime. Change your algorithm without rebuilding your image.
• Structured configuration – Compute, InputData, OutputDataConfig, and StoppingCondition objects replace ad-hoc parameter dictionaries, with IDE auto-complete and type safety.
• Scales from tabular to generative AI – The same pattern trains a scikit-learn classifier on a single CPU and fine-tunes Stable Diffusion 3.5 across multiple GPUs.

For those familiar with script mode or new to SageMaker AI model training, the v3 SDK offers a simplified approach to training and deployment. Its clearly defined, consistent set of primitives speeds up development, no matter the model type.

To learn more about building and deploying your own models using the new SageMaker Python SDK v3, refer to the SageMaker Python SDK v3 documentation and supporting GitHub repository.

Related resources

• SageMaker Python SDK v3 documentation.
• SageMaker Python SDK GitHub repository.
• Original script mode blog post (2021).
• Amazon SageMaker AI.

About the authors

​