SageMaker Python SDK v3: Use SourceCode, ModelTrainer, and ModelBuilder to speed ML development

What changed and why it matters

SDK v3 is a ground-up redesign that consolidates framework-specific classes into clearer primitives. The important pieces are:

• ModelTrainer, one class to configure and launch training jobs (replaces SKLearn/PyTorch/XGBoost Estimators).
• ModelBuilder, packages inference handlers and model artifacts, then creates and deploys endpoints (replaces Model + Predictor).
• SourceCode, the script-mode primitive: point it at a local source_dir plus a shell command for training or an entry_script for inference; the SDK makes that code available inside the container at job start so you can iterate without rebuilding the entire image for most script changes.
• Typed config objects (Compute, InputData, OutputDataConfig, StoppingCondition) replace ad‑hoc dictionaries and make specs easier to validate and reuse.

These changes deliver faster iterations, full container control, and “one API for multiple frameworks, ” as described by the AWS authors Bobby Lindsey and Hazim Qudah.

Two sample workflows (what you can run immediately)

Run the examples in the AWS sample repo to follow along: github.com/aws-samples/sample-sagemaker-pysdkv3-script-mode.

1) scikit‑learn Random Forest (CPU)

Example base image Dockerfile excerpt used in the sample:

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

Example SourceCode command used to start training:

“python random_forest.py –n_jobs 4 –max_depth 10 –n_estimators 120 –mlflow_arn {MLFLOW_ARN} –mlflow_experiment_name {MLFLOW_EXPERIMENT_NAME}”

Example compute settings in the sample:

• instance_type: ml.m5.2xlarge
• instance_count: 1
• volume_size_in_gb: 30
• keep_alive_period_in_seconds: 3600 (warm pool)
• StoppingCondition: max_runtime_in_seconds = 3600

The inference example in the repo uses DJL Serving and sends a CSV payload such as:

“6, 148, 72, 35, 0, 33.6, 0.627, 50”

Code comments in the sample indicate the model returns labels where 1 = tested_positive and 0 = tested_negative.

2) Stable Diffusion 3.5 fine‑tune with LoRA (multi‑GPU)

GPU base image Dockerfile excerpt used in the sample:

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

The example orchestrates distributed training with a bash launcher (base.sh) and Hugging Face Accelerate. The sample uses an ml.g5.12xlarge instance (the G5 family provides A10G GPUs; the g5.12xlarge configuration in the example provides multiple A10G GPUs) and runs across the GPUs on that instance.

Example SourceCode command for the SD job:

“/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 –mlflow-arn {SD_MLFLOW_ARN} –mlflow-experiment-name {SD_MLFLOW_EXPERIMENT_NAME}”

Data is staged into SageMaker training channels with objects like InputData(channel_name=”train”, data_source=SD_TRAINING_DATA_S3_PATH). SageMaker exposes channels at /opt/ml/input/data/<channel_name> and via environment variables SM_CHANNEL_. The post notes you can use up to 20 channels per training job.

Inference example deploys to ml.g5.4xlarge in the sample. Predictor JSON payload example (exact from the repo):

{
“prompt”:”a boy Malcom and his dog Ben”,
“num_inference_steps”:30,
“guidance_scale”:7.5,
“seed”:42
}

The result contains a base64‑encoded generated_image.

Real benefits, and realistic limits

These primitives let teams iterate faster on code changes that are limited to scripts and Python modules. But there are important nuances:

• When SourceCode helps: change a Python training loop, tweak a loss function, or modify hyperparameter parsing and you can usually re-run without rebuilding the base image.
• When you still must rebuild: any change to system packages, OS libraries, CUDA/toolkit versions, compiled native extensions, or driver stacks requires rebuilding the image. SourceCode injects your source files at runtime; it does not modify the installed system/runtime image.
• Startup overhead: syncing large source trees at job start or frequent transfers can add latency that reduces warm‑pool benefits. Consider keeping source_dir minimal for fast restarts.
• Warm pools are billable and bounded: the sample uses keep_alive_period_in_seconds=3600 (60 minutes). Warm pools can improve iteration latency but incur charges while instances are kept alive; warm pools can persist across jobs when matching criteria are met and are governed by SageMaker warm‑pool rules and limits.

Reproducibility, versioning and observability, concrete guidance

SourceCode makes development faster, but you must discipline your workflow to keep runs reproducible and debuggable. Practical steps:

• Pin exact pip package versions in requirements.txt and generate a lock file (pip‑freeze, poetry.lock, or equivalent). Treat the lock file as part of your job inputs.
• Record the git commit SHA used for each run and attach it to job metadata or as an MLflow tag (the examples show optional MLflow integration via an MLFLOW_ARN). That creates a direct mapping from a training job to the exact code snapshot.
• Store the exact source snapshot (tarball) used for each job in S3 as an archival artifact and reference its S3 URI in job metadata so you can rehydrate the same source later.
• Log to CloudWatch and capture training metrics and stdout/stderr; use MLflow (as shown in the examples) or an equivalent tracking backend to centralize experiment metadata and artifacts.
• Be mindful of warm‑pool persistent caches (for example, pip caches in /opt/ml/sagemaker/warmpoolcache). Cache reuse speeds restarts but can mask dependency drift, clear caches as part of reproducibility testing or pin dependency hashes.

Security: practical controls

Examples use Secrets Manager (SECRETS_ARN) to inject Hugging Face tokens. Harden your workflow:

• Give training jobs the least‑privilege IAM role they need; avoid broad S3 or Secrets Manager permissions.
• Restrict the Secrets Manager resource policy so only authorized roles or principals can read a given secret.
• Enable CloudTrail auditing for Secrets Manager and ECR actions to retain an audit trail of who accessed tokens and images.
• Consider VPC endpoints for S3 and ECR access when processing private data to avoid egress over the public internet.
• Rotate tokens regularly and treat long‑lived credentials as high risk; prefer short‑lived session credentials where possible.

When to use SourceCode vs. baking code into an image, a simple heuristic

• Use SourceCode + warm pools for rapid experimentation, iterative model development, and early‑stage fine‑tuning where you change scripts frequently but not the runtime stack.
• Bake code into a tagged image for production training or low‑latency serving where bit‑for‑bit reproducibility, audited images, or custom OS/native library changes matter.
• Hybrid approach: use SourceCode in development to shorten iteration loops; promote the tested code snapshot and produce a minimal, versioned image for production runs.

Migration checklist (v2 → v3)

• Inventory existing v2 uses of Estimator.fit, Model, and Predictor and map them to ModelTrainer, ModelBuilder, and SourceCode equivalents.
• Update CI/CD to pin sagemaker package versions and to capture git SHAs and lockfiles as training inputs.
• Run integration tests in a staging account that exercises warm pools, secrets, and ECR image permissions.
• Validate inference handlers under ModelBuilder (Mode.SAGEMAKER_ENDPOINT, Mode.LOCAL_CONTAINER, Mode.IN_PROCESS) and confirm serving runtimes (DJL, TorchServe, custom) behave as expected.
• Plan a phased rollout: migrate low‑risk workloads first, then promote successful patterns to production.

Operational tips and patterns

• Keep source_dir small: move large assets (datasets, large model checkpoints) to S3 and reference them as InputData channels.
• Tag every training job with pipeline/run identifiers and git SHAs so experiments are traceable.
• Use secrets via SECRETS_ARN and least‑privilege roles; avoid baking tokens into images or source files.
• Test multi‑GPU configs locally with Mode.LOCAL_CONTAINER or small instances before scaling to larger GPUs; the SD example uses Hugging Face Accelerate with an ml.g5.12xlarge instance containing multiple A10G GPUs.
• Integrate cleanup steps in CI/CD pipelines: delete endpoints, endpoint configs, models, S3 artifacts, and ECR images when no longer needed to avoid surprise charges. The sample shows calls like Endpoint.get(endpoint_name=ENDPOINT_NAME).delete(), EndpointConfig.get(endpoint_config_name=ENDPOINT_NAME).delete(), and Model.get(model_name=ENDPOINT_NAME).delete().

“Bring your own model with Amazon SageMaker script mode”

, Bobby Lindsey and Hazim Qudah, AWS

Resources

• Sample repo (Random Forest and Stable Diffusion examples)
• SageMaker Python SDK (GitHub) and the SDK v3 documentation (see release notes and API reference)
• SageMaker warm pools and persistent cache documentation (see AWS SageMaker docs for KeepAlivePeriodInSeconds and warm pool behavior)

Key questions (and honest answers)

• How do I stop rebuilding images for every code change?

  Use the SourceCode primitive in SDK v3 to point at your local source_dir and supply a training command or inference entry_script. The SDK arranges for that code to be available inside the runtime container at job start so many script changes don’t require rebuilding the base image. Rebuilds are still required for changes to system packages or native libraries.

• Can I still use my own image or do I have to use AWS images?

  Yes, ModelTrainer/ModelBuilder support any image: your custom image, AWS Deep Learning Containers, or third‑party images in ECR. You keep full control of the runtime image while using SourceCode to inject scripts.

• Will this pattern work for both tabular models and large generative models?

  Yes. The sample repo demonstrates a scikit‑learn Random Forest running on ml.m5.2xlarge and a LoRA fine‑tune of Stable Diffusion on ml.g5.12xlarge using Hugging Face Accelerate; the same ModelTrainer/ModelBuilder/SourceCode primitives apply across those workloads.

• Are warm pools free, and how long can I keep instances warm?

  Warm pools are billable. The sample uses keep_alive_period_in_seconds = 3600 (60 minutes). Warm pools can persist across jobs when matching criteria are met and are subject to SageMaker warm‑pool rules, evaluate the cost benefit for your iteration cadence.

• How do I keep runs reproducible if code is injected at runtime?

  Pin dependencies, commit and record a git SHA for each run, store a copy of the exact source snapshot used (for example, a tarball in S3), and capture lockfiles. For production runs that require strict immutability, bake the tested code into a versioned image.

Install the SDK (pip install sagemaker>=3.0) and run the samples in the repo to see the pattern in action. Use SourceCode for fast iteration; bake images for hardened production runs, and make reproducibility, security, and cost part of the trade-off conversation when you choose which path to follow.