Bring your own model with Amazon SageMaker AI: Script mode in SDK v3 Amazon Web Services (AWS) announced that the Amazon SageMaker AI Python SDK v3 introduces a unified ModelTrainer class for training and ModelBuilder for deployment, replacing framework-specific estimators like SKLearn, PyTorch, and XGBoost. The new SDK supports script mode with a SourceCode object that syncs local code into any container image from Amazon ECR, AWS Deep Learning Containers, or third-party sources, eliminating the need to rebuild Docker images for each iteration. The post demonstrates two examples: training a scikit-learn Random Forest on the diabetes dataset and fine-tuning Stable Diffusion 3.5 with LoRA using Hugging Face Accelerate. Artificial Intelligence /blogs/machine-learning/ 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 /blogs/machine-learning/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 /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 https://github.com/aws/deep-learning-containers , 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: 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 https://docs.djl.ai/master/index.html DJL Serving, a high-performance model server. 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 https://github.com/aws-samples/sample-sagemaker-pysdkv3-script-mode . What changed from SDK v2 to v3? The following table summarizes the architectural shift: SDK v2 Estimator pattern | SDK v3 ModelTrainer pattern | | Training class | SKLearn , PyTorch , XGBoost , … | ModelTrainer one single class | Deployment class | Model + Predictor | ModelBuilder to deploy the endpoint, prediction handled as part of invoke | Container | AWS managed framework image | Any image: yours, AWS DLC, or third-party | Code injection | entry point + source dir , framework-specific | SourceCode object with source dir + command / entry script | Dependencies | requirements.txt in source dir | requirements.txt in source dir | Prerequisites To follow along, you need: - An AWS account with Amazon SageMaker AI /sagemaker/ access. - An AWS Identity and Access Management /iam/ IAM execution role with Amazon SageMaker AI and Amazon Simple Storage Service Amazon S3 permissions https://docs.aws.amazon.com/sagemaker/latest/dg/security-iam.html . - The SageMaker Python SDK v3 https://sagemaker.readthedocs.io/en/stable/ installed pip install sagemaker =3.0 . - A training container image pushed to Amazon ECR /ecr/ this post shows an example of building and pushing a container to Amazon ECR . - An Amazon S3 /s3/ bucket for training data and model artifacts. - Optional An MLflow app or tracking server on Amazon SageMaker AI https://docs.aws.amazon.com/sagemaker/latest/dg/mlflow.html 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 https://docs.aws.amazon.com/sagemaker/latest/dg/studio-updated-local-get-started.html . 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 https://www.openml.org/d/37 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: 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: 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 https://docs.aws.amazon.com/sagemaker/latest/APIReference/API DockerSettings.html sagemaker-Type-DockerSettings-EnableDockerAccess 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. 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 https://aws.github.io/deep-learning-containers/reference/available images/ . 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 https://docs.aws.amazon.com/sagemaker/latest/dg/mlflow.html 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. 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. 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 https://docs.aws.amazon.com/sagemaker/latest/dg/amazon-sagemaker-toolkits.html for the folder structure and the environment variables that SageMaker sets https://docs.aws.amazon.com/sagemaker/latest/dg/model-train-storage-env-var-summary.html . Step 3: Deploy to a real-time endpoint with ModelBuilder After training completes, we deploy the model artifact to a SageMaker real-time endpoint https://docs.aws.amazon.com/sagemaker/latest/dg/realtime-endpoints.html . 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 https://aws.github.io/deep-learning-containers/reference/available images/ . 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 https://docs.aws.amazon.com/sagemaker/latest/dg/how-it-works-modelbuilder-creation.html . With the model built, we call deploy to stand up the real-time endpoint, which returns an Endpoint interface: 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 https://docs.djl.ai/master/docs/serving/serving/docs/modes.html python-mode , 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 https://sagemaker.readthedocs.io/en/stable/api/sagemaker serve.html sagemaker.serve.ModelServer . - 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 https://docs.aws.amazon.com/sagemaker/latest/dg/realtime-endpoints.html and Inference components https://docs.aws.amazon.com/sagemaker/latest/dg/realtime-endpoints-deploy-models.html deployed-shared-utilization . Step 4: Test the endpoint Send a sample CSV request to confirm the endpoint is healthy: 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: As with the preceding container, the example notebook includes Docker build and push commands by using two shell scripts: 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: 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 https://docs.aws.amazon.com/sagemaker/latest/dg/how-it-works-training.html . 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 https://huggingface.co/datasets/google/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: 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/