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.
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:
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:
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:
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:
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_TRAINas a file instead of a directory. Input channels are mounted folders, so your code usually joins a filename onto that path. - Confusing
output_pathwith the in-container model directory.output_pathis an S3 destination, not the local place where your script writes files. - Forgetting to include dependencies in
source_diror 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/modeland upload it to the configuredoutput_path. - Use
estimator.model_dataas the canonical S3 location for deployment and later reuse.
Related reading
- How to use sample weights with tensorflow datasets?
- How to use sample weights with tensorflow datasets?
- how to use scipy.optimize.linear_sum_assignment in tensorflow or keras?
- How to use several summary collections in Tensorflow?
- How to use scikit-learn PCA for features reduction and know which features are discarded
- How to use SGD for time series analysis
- How to use spot instance with amazon elastic beanstalk?
- How to use Sub and GetAtt functions at the same time in CloudFormation template?

System Design Fundamentals
Build a strong foundation in designing scalable, reliable distributed systems.
View the courseTrack 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.