Train and Deploy ML Models with Amazon SageMaker AI and AWS Batch

This guide shows how to train a machine learning model on Amazon SageMaker AI, how AWS Batch fits in to queue and prioritise training jobs, and how to deploy the trained model as a prediction endpoint. It uses a small scikit-learn model so you can see the whole workflow, and it points out the cost traps that catch people out.

Updated October 2026: I rewrote this March 2024 guide because much of it no longer matches AWS. Amazon SageMaker was renamed Amazon SageMaker AI on 3 December 2024, and the newer Amazon SageMaker is now a wider unified data and AI platform. The SageMaker Python SDK is now on version 3, with a different API from the Estimator classes the old post used. And AWS Batch gained a direct integration with SageMaker: you can queue SageMaker Training jobs through Batch “service jobs”. I checked commands and request shapes against the AWS documentation and the current SDK and boto3 service definitions, and tested the training and inference scripts locally. I did not run anything against a real AWS account, so training and endpoints cost real money when you do: read the cost section first.

What Is What Now

• Amazon SageMaker AI (the service formerly just called SageMaker) is the managed service for building, training and deploying ML and foundation models, per AWS’s documentation. The API namespaces, CLI commands and IAM policy names that start with sagemaker were deliberately left unchanged.
• Amazon SageMaker (the next generation, from December 2024) is a broader platform that bundles SageMaker AI with a lakehouse, data governance, SQL analytics, data processing, SageMaker Unified Studio and Amazon Bedrock. You don’t need it for the workflow here.
• SageMaker Studio is the web IDE (JupyterLab, Code Editor and so on). AWS recommends the current Studio experience over Studio Classic, which is based on the older JupyterLab 3 and stopped receiving critical fixes for its JupyterLab 3 notebooks after 31 December 2024 (see AWS’s migration guide). Everything below also works from your own laptop or CI with the SDK and CLI, with no Studio domain at all.
• AWS Batch is a general batch-job scheduler that runs containers on EC2, Spot or Fargate. It can also now queue and prioritise SageMaker Training jobs through service jobs (next sections).

Do You Need AWS Batch With SageMaker?

Often not. SageMaker Training already provisions the instances, runs your container and shuts everything down. Batch becomes useful when many training jobs compete for limited, expensive capacity and you want queueing, priorities and a cap on how many instances run at once. Per AWS’s Batch on SageMaker AI guide, you create a service environment that sets the maximum number of training instances, attach a job queue to it, and submit training jobs as service jobs: Batch handles scheduling and priority, while the training itself runs on SageMaker infrastructure. AWS describes the benefits as removing manual coordination, avoiding accidental overspending and improving resource use across teams.

The old post instead had you containerise a training application and run it on a regular Batch compute environment. That still works as a separate approach (Batch runs any container, outside SageMaker), but it isn’t “Batch with SageMaker”, and you give up SageMaker’s managed training features.

Cost Warning, Before Anything Else

• Training instances are billed for the instance type you pick while the job runs. A GPU instance can cost many times more than a small CPU one, so start with a small CPU instance and a short maximum runtime.
• Real-time endpoints bill for as long as they exist, whether or not they receive traffic, according to the instance type. The most common surprise bill is an endpoint left running after a tutorial. Delete endpoints when you finish, and set a billing alarm.
• Serverless inference is billed differently: per AWS’s pricing page, you pay for the compute used to process requests (billed by the millisecond) and the data processed, which suits spiky traffic that can tolerate cold starts.
• Free tier: that page describes a limited free-tier allowance for new SageMaker AI users for two months (hours of notebook, training and inference instances). Check the current terms and your account’s eligibility; they are small.
• S3 storage, data transfer, CloudWatch logs and ECR images add small ongoing amounts.

Step 1: Prepare the Account

1. Use an administrative identity (not the root user) with MFA, set a budget and billing alerts, and choose one Region.
2. Create an S3 bucket for training data and model output (private, with public access blocked). SageMaker needs the bucket in the same Region.
3. Create a SageMaker execution role, the role SageMaker assumes to read your data, write the model and pull the container. AWS’s Batch tutorial creates one with this trust policy and managed policies (a quick way to start; see the note below):

# sagemaker-trust-policy.json
{
  "Version": "2012-10-17",
  "Statement": [{
    "Effect": "Allow",
    "Principal": { "Service": "sagemaker.amazonaws.com" },
    "Action": "sts:AssumeRole"
  }]
}

aws iam create-role --role-name SageMakerExecutionRole \
  --assume-role-policy-document file://sagemaker-trust-policy.json

aws iam attach-role-policy --role-name SageMakerExecutionRole \
  --policy-arn arn:aws:iam::aws:policy/AmazonSageMakerFullAccess

Least privilege: AWS’s tutorial also attaches AmazonS3FullAccess for convenience, which lets the role read and write every bucket in the account. For anything real, replace it with an inline policy that allows only your bucket (see AWS’s guidance on SageMaker execution roles), and review what AmazonSageMakerFullAccess grants before using it beyond a sandbox.

Step 2: The Training Script

SageMaker’s framework containers run a script you provide. This one trains a random-forest classifier on synthetic data, prints a validation accuracy (which SageMaker can capture as a metric), saves the model to the directory SageMaker uploads, and also copies the inference handlers (next section) into the model artifact so the same artifact can be deployed as is. SageMaker mounts input channels under /opt/ml/input/data/<channel> and uploads whatever you write to /opt/ml/model (the SM_MODEL_DIR environment variable) as model.tar.gz. Put this in src/train.py:

"""Minimal training script run inside the SageMaker training container."""
import argparse, json, os, shutil
import joblib
from sklearn.datasets import make_classification
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--n-estimators", type=int, default=100)
    args = parser.parse_args()

    # SageMaker mounts channels under /opt/ml/input/data/<channel>; models go to /opt/ml/model
    model_dir = os.environ.get("SM_MODEL_DIR", "/opt/ml/model")
    X, y = make_classification(n_samples=2000, n_features=20, random_state=0)
    split = 1500
    clf = RandomForestClassifier(n_estimators=args.n_estimators, random_state=0).fit(X[:split], y[:split])
    acc = accuracy_score(y[split:], clf.predict(X[split:]))
    print(f"validation_accuracy={acc:.4f}")           # SageMaker can scrape this with a metric regex

    os.makedirs(model_dir, exist_ok=True)
    joblib.dump(clf, os.path.join(model_dir, "model.joblib"))
    # Ship the inference handlers inside the model artifact so the same file can be deployed as-is
    code_dir = os.path.join(model_dir, "code")
    os.makedirs(code_dir, exist_ok=True)
    shutil.copy(os.path.join(os.path.dirname(os.path.abspath(__file__)), "inference.py"), code_dir)

    with open(os.path.join(model_dir, "metrics.json"), "w") as f:
        json.dump({"validation_accuracy": acc}, f)

if __name__ == "__main__":
    main()

And the inference handlers, which the scikit-learn serving container calls, in src/inference.py:

"""Inference handlers for the SageMaker scikit-learn serving container."""
import json
import os

import joblib
import numpy as np


def model_fn(model_dir):
    """Load the model once when the container starts."""
    return joblib.load(os.path.join(model_dir, "model.joblib"))


def input_fn(request_body, content_type):
    if content_type != "application/json":
        raise ValueError(f"Unsupported content type: {content_type}")
    payload = json.loads(request_body)
    return np.array(payload["instances"], dtype=float)


def predict_fn(data, model):
    return model.predict(data)


def output_fn(prediction, accept):
    return json.dumps({"predictions": prediction.tolist()}), "application/json"

I ran both locally: training reported a validation accuracy of 0.99 on the synthetic data, and the handlers loaded the model and returned predictions for JSON input and rejected unsupported content types. Replace the synthetic data with your own: read a CSV from the training channel directory and keep the same structure.

Step 3: Train With the SageMaker Python SDK 3

Install the SDK (version 3.x requires Python 3.10 or newer; see its PyPI page): pip install sagemaker. SDK 3 replaces the framework-specific Estimator classes with a ModelTrainer built from small configuration objects. This is the same job expressed with it (the parameter names come from the installed SDK, and the training image URI comes from the SDK’s own image lookup for scikit-learn 1.2-1 in us-east-1):

from sagemaker.train.model_trainer import ModelTrainer
from sagemaker.train.configs import (
    SourceCode, Compute, InputData, OutputDataConfig, StoppingCondition,
)
from sagemaker.core import image_uris

REGION = "us-east-1"
ROLE = "arn:aws:iam::YOUR_ACCOUNT_ID:role/SageMakerExecutionRole"
BUCKET = "your-ml-bucket"

training_image = image_uris.retrieve(
    framework="sklearn", region=REGION, version="1.2-1", py_version="py3",
    instance_type="ml.m5.xlarge", image_scope="training",
)

trainer = ModelTrainer(
    training_image=training_image,
    source_code=SourceCode(source_dir="src", entry_script="train.py"),
    compute=Compute(instance_type="ml.m5.xlarge", instance_count=1),
    hyperparameters={"n-estimators": 100},          # passed to train.py as command-line arguments
    base_job_name="rf-demo",
    role=ROLE,
    output_data_config=OutputDataConfig(s3_output_path=f"s3://{BUCKET}/output"),
    stopping_condition=StoppingCondition(max_runtime_in_seconds=900),   # hard cap on cost
)

# For real data, add: input_data_config=[InputData(channel_name="train", data_source=f"s3://{BUCKET}/train/")]
trainer.train(wait=True, logs=True)

I confirmed that the SDK accepts these configuration objects and resolves the image URI; creating the trainer calls AWS (it checks the role’s permissions), so I could not run it end to end without a real account. The call returns when the job finishes, with logs streaming from CloudWatch. The model artifact lands under the S3 output path (in a folder named after the job). SDK 3 is a major version: if you find older tutorials using sagemaker.sklearn.SKLearn(...).fit(...), they target SDK 2.

Optional: use managed Spot training (enable_managed_spot_training on the compute configuration, plus a maximum wait time) for large discounts on interruptible jobs, with checkpointing so interrupted jobs can resume.

Step 4: Queue Training Jobs With AWS Batch (Optional)

If several people or pipelines submit training jobs, add Batch. These commands follow AWS’s getting-started guide; the first creates a service environment that allows at most 5 training instances at once, the second a queue attached to it. The first time you create a service environment, Batch also creates a service-linked role so it can manage SageMaker Training jobs for you.

aws batch create-service-environment \
  --service-environment-name ml-training-env \
  --service-environment-type SAGEMAKER_TRAINING \
  --capacity-limits capacityUnit=NUM_INSTANCES,maxCapacity=5

aws batch create-job-queue \
  --job-queue-name ml-training-queue \
  --job-queue-type SAGEMAKER_TRAINING \
  --priority 1 \
  --service-environment-order order=1,serviceEnvironment=ml-training-env

A service job’s payload is a standard SageMaker CreateTrainingJob request. This is AWS’s minimal example (a container that just echoes text, so it costs almost nothing), which I validated against the current SageMaker API model; save it as training_job.json, filling in your account ID and bucket:

{
  "TrainingJobName": "my-simple-training-job",
  "RoleArn": "arn:aws:iam::YOUR_ACCOUNT_ID:role/SageMakerExecutionRole",
  "AlgorithmSpecification": {
    "TrainingInputMode": "File",
    "TrainingImage": "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.0.0-cpu-py310",
    "ContainerEntrypoint": ["echo", "hello world"]
  },
  "ResourceConfig": { "InstanceType": "ml.c5.xlarge", "InstanceCount": 1, "VolumeSizeInGB": 1 },
  "OutputDataConfig": { "S3OutputPath": "s3://your-s3-bucket/output" },
  "StoppingCondition": { "MaxRuntimeInSeconds": 30 }
}

aws batch submit-service-job \
  --job-queue ml-training-queue \
  --job-name my-batch-sm-job \
  --service-job-type SAGEMAKER_TRAINING \
  --retry-strategy attempts=1 \
  --timeout-config attemptDurationSeconds=60 \
  --service-request-payload file://training_job.json

aws batch list-service-jobs --job-queue ml-training-queue --job-status RUNNING
aws batch describe-service-job --job-id YOUR_JOB_ID

Note the image in the example is in us-west-2; use a container image from your own Region. For a real training job, the payload carries the same fields the SDK builds for you, plus input channels and (for framework containers) the script settings. A common pattern is to let your pipeline construct the payload and submit it through Batch so jobs queue by priority. When you’re done, delete the queue and service environment (aws batch delete-job-queue and delete-service-environment).

Step 5: Deploy the Model

SageMaker AI offers several ways to serve a model; per AWS’s deployment guide: real-time endpoints for low-latency interactive use, serverless endpoints for workloads with idle periods that can tolerate cold starts, asynchronous endpoints that queue requests (payloads up to 1 GB, processing up to an hour), and batch transform for scoring whole datasets without an endpoint. The SDK’s ModelBuilder class is AWS’s recommended code route for deploying your own model; the three-call boto3 version below is the most explicit way to see what is created, and I validated each request against the current boto3 service definitions.

import boto3
from sagemaker.core import image_uris

REGION = "us-east-1"
ROLE = "arn:aws:iam::YOUR_ACCOUNT_ID:role/SageMakerExecutionRole"
MODEL_DATA = "s3://your-ml-bucket/output/rf-demo-.../output/model.tar.gz"   # from the training job

image = image_uris.retrieve(
    framework="sklearn", region=REGION, version="1.2-1", py_version="py3",
    instance_type="ml.m5.large", image_scope="inference",
)

sm = boto3.client("sagemaker", region_name=REGION)

sm.create_model(
    ModelName="rf-demo",
    ExecutionRoleArn=ROLE,
    PrimaryContainer={
        "Image": image,
        "ModelDataUrl": MODEL_DATA,
        "Environment": {"SAGEMAKER_PROGRAM": "inference.py"},   # use code/inference.py inside the artifact
    },
)

# Serverless: pay per request, scales to zero, but has cold starts
sm.create_endpoint_config(
    EndpointConfigName="rf-demo-serverless",
    ProductionVariants=[{
        "VariantName": "AllTraffic",
        "ModelName": "rf-demo",
        "ServerlessConfig": {"MemorySizeInMB": 2048, "MaxConcurrency": 5},
    }],
)
# (Alternative: real-time, always-on instances)
# ProductionVariants=[{"VariantName": "AllTraffic", "ModelName": "rf-demo",
#                      "InstanceType": "ml.m5.large", "InitialInstanceCount": 1}]

sm.create_endpoint(EndpointName="rf-demo", EndpointConfigName="rf-demo-serverless")
sm.get_waiter("endpoint_in_service").wait(EndpointName="rf-demo")

Then call it with the runtime client, sending JSON in the format inference.py expects:

import json, boto3

runtime = boto3.client("sagemaker-runtime", region_name="us-east-1")
response = runtime.invoke_endpoint(
    EndpointName="rf-demo",
    ContentType="application/json",
    Body=json.dumps({"instances": [[0.0] * 20]}),
)
print(json.loads(response["Body"].read()))   # {"predictions": [1]}

Whether the container picks up inference.py from the artifact exactly as shown depends on the container version; if the endpoint fails to start, check the CloudWatch logs for the endpoint and consult the scikit-learn container documentation. Delete everything when you’re done:

aws sagemaker delete-endpoint --endpoint-name rf-demo
aws sagemaker delete-endpoint-config --endpoint-config-name rf-demo-serverless
aws sagemaker delete-model --model-name rf-demo

Monitoring and Operating the Model

• Operational metrics. Endpoints publish invocation counts, latency and error rates to CloudWatch; set alarms on error rates and latency, and add autoscaling for real-time endpoints so capacity follows traffic.
• Model quality. Accuracy in production decays as data changes (drift). SageMaker Model Monitor can compare live inputs and predictions with a baseline, but you still need ground truth labels to measure real accuracy, which often arrive late. Plan how you will get them.
• Safe rollouts. Deploy new models as separate production variants or endpoints and shift traffic gradually; keep the previous model available for rollback.
• Reproducibility. Record the code version, data version, hyperparameters and metrics for every trained model. Tools such as SageMaker Experiments, the Model Registry or MLflow exist for this.
• Security. Run training and endpoints in a VPC if data is sensitive, encrypt S3 and volumes with KMS keys, and use narrowly scoped execution roles.

If your goal is to use a pre-trained or foundation model rather than train your own, look at SageMaker JumpStart or Amazon Bedrock first; training from scratch is rarely the cheapest route. For related background see how large language models work and what is an agentic AI, and for AWS fundamentals deploying a web app on AWS.