{"slug": "bring-your-own-model-with-amazon-sagemaker-ai-script-mode-in-sdk-v3", "title": "Bring your own model with Amazon SageMaker AI: Script mode in SDK v3", "summary": "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.", "body_md": "[Artificial Intelligence](/blogs/machine-learning/)\n\n# Bring your own model with Amazon SageMaker AI: Script mode in SDK v3\n\nIn 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/).\n\nThe 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`\n\n, `PyTorch`\n\n, `XGBoost`\n\n) with a single, unified `ModelTrainer`\n\nfor training and `ModelBuilder`\n\nfor deployment.\n\nIn v3, the SDK syncs a local source code directory into the training job at runtime using the new `SourceCode`\n\nconfiguration 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.\n\nThis means:\n\n**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.\n\n## Solution overview\n\nIn this post, we walk through two end-to-end examples that demonstrate how script mode works in the SageMaker Python SDK v3:\n\n**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.\n\nBoth examples use the same two core classes:\n\n`ModelTrainer`\n\nreplaces the v2 Estimator family. Configures and launches a SageMaker training job.`ModelBuilder`\n\nreplaces the v2 Model/Predictor pattern. Packages your inference handler and deploys to an endpoint.\n\nA key concept is the `SourceCode`\n\nobject. It accepts a `source_dir`\n\n(a path to your local code directory) and either a `command`\n\nstring (for training) or an `entry_script`\n\n(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.\n\nYou can find the example code for this blog post in the [GitHub repository](https://github.com/aws-samples/sample-sagemaker-pysdkv3-script-mode).\n\n### What changed from SDK v2 to v3?\n\nThe following table summarizes the architectural shift:\n\nSDK v2 (Estimator pattern) |\nSDK v3 (ModelTrainer pattern) |\n|\nTraining class |\n`SKLearn` , `PyTorch` , `XGBoost` , … |\n`ModelTrainer` (one single class) |\nDeployment class |\nModel + Predictor | `ModelBuilder` to deploy the endpoint, prediction handled as part of `invoke()` |\nContainer |\nAWS managed framework image | Any image: yours, AWS DLC, or third-party |\nCode injection |\n`entry_point` + `source_dir` , framework-specific |\n`SourceCode` object with `source_dir` + `command` /`entry_script` |\nDependencies |\n`requirements.txt` in `source_dir` |\n`requirements.txt` in `source_dir` |\n\n## Prerequisites\n\nTo follow along, you need:\n\n- An AWS account with\n[Amazon SageMaker AI](/sagemaker/)access. - An\n[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\n[SageMaker Python SDK v3](https://sagemaker.readthedocs.io/en/stable/)installed (`pip install sagemaker>=3.0`\n\n). - A training container image pushed to\n[Amazon ECR](/ecr/)(this post shows an example of building and pushing a container to Amazon ECR). - An\n[Amazon S3](/s3/)bucket for training data and model artifacts. - (Optional) An\n[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\n[Local mode support in Amazon SageMaker Studio](https://docs.aws.amazon.com/sagemaker/latest/dg/studio-updated-local-get-started.html).\n\n## Example 1: Train and deploy a scikit-learn model\n\nLet’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.\n\n### Step 1: Building the Docker container\n\nThe 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.\n\nThe complete Dockerfile for our scikit-learn container is:\n\nThe container is a stable, version-controlled runtime environment. The algorithm-specific code lives in your `source_dir`\n\n, and the SDK injects it at runtime.\n\nBuild this container once, push it to Amazon ECR, and iterate on your training code as many times as you want without touching Docker again.\n\nThe example notebook includes Docker build and push commands by using two shell scripts:\n\nNote 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.\n\n### Step 1a: Configuration\n\nFirst, 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.\n\nIn the following snippet, we point `TRAINING_IMAGE_URI`\n\nat 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/).\n\nOptionally, 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`\n\nand `MLFLOW_EXPERIMENT_NAME`\n\nvariables in the following snippet to automatically enable logging. This is an optional step, and you can set the values to `None`\n\nto skip experiment tracking instead.\n\n### Step 2: Launch a training job with ModelTrainer\n\nThe `SourceCode`\n\nobject takes your local `source_dir`\n\nand a `command`\n\nstring. At job launch, SageMaker syncs the entire `source_dir`\n\ninto the container and runs your `command`\n\n. This decouples your code from your container image. Change the script, re-launch with no container rebuild needed.\n\nA few things to note:\n\n`source_dir`\n\ncan contain your files such as utility modules, config files, and shell scripts, which are synced into the container.`command`\n\nis 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`\n\nturns on SageMaker warm pools. The instance stays warm for 1 hour, so iterative re-runs launch in seconds rather than minutes.`OutputDataConfig`\n\nsets the S3 destination where SageMaker uploads your training results when the job finishes. Anything your script saves to`/opt/ml/model`\n\n(the`SM_MODEL_DIR`\n\nenvironment variable) is packaged as`model.tar.gz`\n\nunder 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).\n\n### Step 3: Deploy to a real-time endpoint with ModelBuilder\n\nAfter 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`\n\npackages 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`\n\nobject. 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/).\n\nThe `build()`\n\nstep assembles a deployable model without launching any infrastructure. `ModelBuilder`\n\ntakes 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`\n\ncan 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).\n\nWith the model built, we call `deploy()`\n\nto stand up the real-time endpoint, which returns an Endpoint interface:\n\nNotice the same `SourceCode`\n\npattern for inference: point at a local directory containing your handler and specify the `entry_script`\n\n. The SDK repacks the handler into the model archive so DJL Serving can find it at runtime.\n\nA few notes on the preceding code snippets:\n\n- The\n`inference.py`\n\nscript implements a single`handle(inputs)`\n\nfunction 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`\n\n, which corresponds to the`SM_MODEL_DIR`\n\nenvironment 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.\n- The\n`model_server`\n\nargument tells`ModelBuilder`\n\nwhich 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`\n\nis written. Here, we choose`ModelServer.DJL_SERVING`\n\n, 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)`\n\ncontract described earlier. For other model serving choices exposed by`ModelServer`\n\n, see[the ModelServer API reference](https://sagemaker.readthedocs.io/en/stable/api/sagemaker_serve.html#sagemaker.serve.ModelServer). - The\n`mode`\n\nparameter controls*where*your model runs. Here we use`Mode.SAGEMAKER_ENDPOINT`\n\n, which deploys to a fully managed real-time endpoint.`ModelBuilder`\n\nalso supports`Mode.LOCAL_CONTAINER`\n\n(run in a Docker container on your machine) and`Mode.IN_PROCESS`\n\n(run directly in your current Python process) for testing and iterating on your handler locally.\n\nIn 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).\n\n### Step 4: Test the endpoint\n\nSend a sample CSV request to confirm the endpoint is healthy:\n\nExample 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.\n\n## Example 2: Fine-tune Stable Diffusion 3.5 with LoRA\n\nThe 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.\n\n### Step 1: Building the Docker container\n\nAs in the scikit-learn example, the container is purely a runtime environment. The complete Dockerfile for our Stable Diffusion container is:\n\nAs with the preceding container, the example notebook includes Docker build and push commands by using two shell scripts:\n\nThe 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`\n\nand are synced at runtime:\n\nYou 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.\n\n### Step 1a: Prepare the training data\n\nIn 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`\n\nand populate it into our working bucket, which we then pass into the training job as `InputData`\n\n, as shown in the following code:\n\nThe `channel_name`\n\nyou assign in `InputData`\n\ncontrols 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>`\n\nand exposes the path through a matching `SM_CHANNEL_<CHANNEL_NAME>`\n\n[environment variable](https://docs.aws.amazon.com/sagemaker/latest/dg/model-train-storage-env-var-summary.html).\n\nHere we define a single `train`\n\nchannel, so the dataset lands at `/opt/ml/input/data/train`\n\n. 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`\n\n, `validation`\n\n, and `test`\n\nchannels, 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](https://docs.aws.amazon.com/sagemaker/latest/dg/your-algorithms-training-algo-running-container.html).\n\n### Step 2: Launch a training job with ModelTrainer\n\nThe `command`\n\nlaunches a bash script (`base.sh`\n\n) that detects the GPU count, runs `accelerate launch`\n\n, 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:\n\nAt this point, we’re ready to create the `ModelTrainer`\n\nclass following a similar paradigm as the previous example. However, we now include a `SECRETS_ARN`\n\ncorresponding to an entry in [AWS Secrets Manager](/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.\n\nWhen the container starts up, the `base.sh`\n\nscript 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.\n\n### Step 3: Deploy to a real-time endpoint with ModelBuilder\n\nAfter the training step completes, we locate the trained LoRA weights, create a `ModelBuilder`\n\nobject, and deploy our fine-tuned model to a real-time endpoint, following a similar pattern as the previous example:\n\n### Step 4: Test the endpoint\n\nLastly, we test by sending a sample request to confirm the endpoint is healthy:\n\nKey takeaways from this example:\n\n**Bash launchers work well**– your command can be a shell command, not only`python script.py`\n\n. 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`\n\n. 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](/secrets-manager/)and pass the corresponding ARN through the`environment`\n\nparameter. 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](https://docs.aws.amazon.com/secretsmanager/latest/userguide/auth-and-access.html).**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.\n\n## Clean up\n\nTo avoid incurring future charges, delete the resources you created:\n\n- Delete the SageMaker endpoint:\n- Delete the S3 training data and model artifacts if they are no longer needed.\n- Delete the ECR container images if you no longer need them.\n- (Optional) Delete the MLflow tracking server if you created one for this walkthrough.\n\n## Conclusion\n\nThe 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:\n\n**One API for multiple workloads**–`ModelTrainer`\n\nand`ModelBuilder`\n\nreplace a dozen framework-specific classes. Less to learn, less to maintain.**Code-container decoupling**–`SourceCode`\n\nsyncs your local code directory into the container at runtime. Change your algorithm without rebuilding your image.**Structured configuration**–`Compute`\n\n,`InputData`\n\n,`OutputDataConfig`\n\n, and`StoppingCondition`\n\nobjects 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.\n\nFor 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.\n\nTo learn more about building and deploying your own models using the new SageMaker Python SDK v3, refer to the [SageMaker Python SDK v3 documentation](https://sagemaker.readthedocs.io/en/stable/) and supporting [GitHub repository](https://github.com/aws/sagemaker-python-sdk).\n\n## Related resources\n\n[SageMaker Python SDK v3 documentation](https://sagemaker.readthedocs.io/en/stable/).[SageMaker Python SDK GitHub repository](https://github.com/aws/sagemaker-python-sdk).[Original script mode blog post (2021)](/blogs/machine-learning/bring-your-own-model-with-amazon-sagemaker-script-mode/).[Amazon SageMaker AI](/sagemaker/).", "url": "https://wpnews.pro/news/bring-your-own-model-with-amazon-sagemaker-ai-script-mode-in-sdk-v3", "canonical_source": "https://aws.amazon.com/blogs/machine-learning/bring-your-own-model-with-amazon-sagemaker-ai-script-mode-in-sdk-v3/", "published_at": "2026-08-26 16:31:32+00:00", "updated_at": "2026-08-26 16:44:12.051984+00:00", "lang": "en", "topics": ["machine-learning", "artificial-intelligence", "generative-ai", "developer-tools"], "entities": ["Amazon Web Services", "Amazon SageMaker AI", "Amazon Elastic Container Registry", "AWS Deep Learning Container", "Deep Java Library", "Hugging Face", "Stable Diffusion 3.5", "scikit-learn"], "alternates": {"html": "https://wpnews.pro/news/bring-your-own-model-with-amazon-sagemaker-ai-script-mode-in-sdk-v3", "markdown": "https://wpnews.pro/news/bring-your-own-model-with-amazon-sagemaker-ai-script-mode-in-sdk-v3.md", "text": "https://wpnews.pro/news/bring-your-own-model-with-amazon-sagemaker-ai-script-mode-in-sdk-v3.txt", "jsonld": "https://wpnews.pro/news/bring-your-own-model-with-amazon-sagemaker-ai-script-mode-in-sdk-v3.jsonld"}}