SageMaker
Estimator
Model Training
Machine Learning
AWS

How to use SageMaker Estimator for model training and saving

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

In Amazon SageMaker, an Estimator defines how a training job runs and where its artifacts end up. The critical rule is that your training script saves the finished model into SM_MODEL_DIR, and SageMaker packages that directory and uploads it to the S3 output_path configured on the Estimator.

Create the Estimator

An Estimator bundles the training image, entry script, compute settings, and output location. For framework-based training jobs, the SDK classes such as SKLearn, PyTorch, or TensorFlow are the easiest place to start.

This example uses SKLearn:

python
1import sagemaker
2from sagemaker.sklearn.estimator import SKLearn
3
4session = sagemaker.Session()
5role = "arn:aws:iam::123456789012:role/SageMakerExecutionRole"
6bucket = session.default_bucket()
7
8estimator = SKLearn(
9    entry_point="train.py",
10    source_dir="src",
11    role=role,
12    framework_version="1.2-1",
13    py_version="py3",
14    instance_count=1,
15    instance_type="ml.m5.large",
16    output_path=f"s3://{bucket}/training-output/demo/",
17    hyperparameters={
18        "n-estimators": 200,
19        "max-depth": 6,
20    },
21)
22
23estimator.fit(
24    inputs={
25        "train": f"s3://{bucket}/datasets/demo/"
26    }
27)

When you call fit, SageMaker uploads the code, starts the training container, mounts the input channels, and streams logs to CloudWatch. After the job completes, estimator.model_data points at the final model.tar.gz file in S3.

Save the Model in SM_MODEL_DIR

The most important part of the training script is where you write the finished model. SageMaker expects the final deployable artifact in /opt/ml/model, which is exposed to your code through the SM_MODEL_DIR environment variable.

Here is a minimal training script:

python
1import argparse
2import os
3
4import joblib
5import pandas as pd
6from sklearn.ensemble import RandomForestClassifier
7
8
9def parse_args():
10    parser = argparse.ArgumentParser()
11    parser.add_argument("--model-dir", type=str, default=os.environ["SM_MODEL_DIR"])
12    parser.add_argument("--train", type=str, default=os.environ["SM_CHANNEL_TRAIN"])
13    parser.add_argument("--n-estimators", type=int, default=100)
14    parser.add_argument("--max-depth", type=int, default=5)
15    return parser.parse_args()
16
17
18if __name__ == "__main__":
19    args = parse_args()
20
21    data = pd.read_csv(os.path.join(args.train, "train.csv"))
22    features = data.drop(columns=["label"])
23    labels = data["label"]
24
25    model = RandomForestClassifier(
26        n_estimators=args.n_estimators,
27        max_depth=args.max_depth,
28        random_state=42,
29    )
30    model.fit(features, labels)
31
32    joblib.dump(model, os.path.join(args.model_dir, "model.joblib"))

SageMaker compresses everything inside args.model_dir into model.tar.gz after the training process exits successfully. That archive becomes the model artifact you deploy later.

Understand the Data and Hyperparameter Flow

SageMaker maps training inputs to channels. In the example above, the train channel becomes the SM_CHANNEL_TRAIN directory inside the container. Hyperparameters passed to the Estimator become command-line arguments, which is why argparse is the normal pattern inside train.py.

This separation is useful because it keeps the script portable. You can run the same script locally by passing --train and --model-dir, or hand it to SageMaker with different instances and S3 prefixes without changing the code.

If you have helper modules, place them under source_dir. SageMaker uploads that directory together with the entry script.

Use the Saved Artifact After Training

After the job finishes, the SDK exposes the model artifact location:

python
print(estimator.model_data)

A typical value looks like s3://bucket/training-output/demo/job-name/output/model.tar.gz. That S3 path can then be used for deployment:

python
1from sagemaker.sklearn.model import SKLearnModel
2
3model = SKLearnModel(
4    model_data=estimator.model_data,
5    role=role,
6    framework_version="1.2-1",
7    py_version="py3",
8    entry_point="inference.py",
9)

The important point is that you do not manually upload the final model from train.py. Your script writes locally to SM_MODEL_DIR, and SageMaker handles packaging and S3 upload.

Common Pitfalls

  • Saving the model somewhere other than SM_MODEL_DIR. The training job may succeed, but the final artifact will not be included in the uploaded archive.
  • Treating SM_CHANNEL_TRAIN as a file instead of a directory. Input channels are mounted folders, so your code usually joins a filename onto that path.
  • Confusing output_path with the in-container model directory. output_path is an S3 destination, not the local place where your script writes files.
  • Forgetting to include dependencies in source_dir or the container image. Training can fail even when the Estimator configuration itself is correct.
  • Writing checkpoints and final models into the same place without intent. Temporary checkpoints and the finished deployable model should be managed separately.

Summary

  • An Estimator defines how SageMaker runs training and where the final artifact is uploaded.
  • Save the finished model into SM_MODEL_DIR, not an arbitrary temporary path.
  • Read training data from mounted channel directories such as SM_CHANNEL_TRAIN.
  • Let SageMaker package /opt/ml/model and upload it to the configured output_path.
  • Use estimator.model_data as the canonical S3 location for deployment and later reuse.

Related reading
Course
Beginner
27 lessons
10 hours
System Design Fundamentals

Build a strong foundation in designing scalable, reliable distributed systems.

View the course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.