7.2 Experiment Tracking with SageMaker Experiments
Key Takeaways
- SageMaker Experiments organizes ML workflows into a strict 3-tier hierarchy: Experiment (high-level business objective) -> Run Group / Trial (logical collection of related runs) -> Run / Trial Component (single execution step logging parameters, metrics, and artifacts).
- The modern SageMaker Python SDK provides the `sagemaker.experiments.Run` context manager to log parameters (`run.log_parameters`), custom metrics (`run.log_metric`), and output artifacts (`run.log_artifact`) seamlessly inside training scripts or notebooks.
- SageMaker Lineage Tracking constructs an automated, immutable Directed Acyclic Graph (DAG) connecting raw S3 data, preprocessing jobs, training jobs, models, and endpoints for regulatory compliance and auditability.
- SageMaker Studio integrates visual leaderboards, scatter plots, and parallel coordinates charts to compare hyperparameter sensitivity and model metrics across dozens or hundreds of experimental runs.
- Estimators automatically register runs with SageMaker Experiments by passing the `experiment_config={'ExperimentName': '...', 'RunName': '...'}` parameter to `estimator.fit()`.
Experiment Tracking with SageMaker Experiments
During iterative machine learning development, data scientists and ML engineers train dozens or hundreds of candidate models across different algorithm architectures, hyperparameter configurations, feature subsets, and data transformations. Without structured tracking, reproducing a high-performing model, auditing data provenance, or diagnosing why a model degraded becomes nearly impossible.
Amazon SageMaker Experiments is a managed capability that automatically tracks, organizes, compares, and evaluates machine learning iterations. It captures metadata across the entire ML workflow—including input datasets, Git commits, hyperparameters, training loss curves, evaluation matrices, and output model artifacts. Coupled with SageMaker Lineage Tracking, it provides an auditable history of how every production model was created.
1. The SageMaker Experiments Logical Hierarchy
SageMaker Experiments structures ML workflows into a clean, 3-tier hierarchical data model.
+-----------------------------------------------------------------------------------------+
| SAGEMAKER EXPERIMENTS HIERARCHY |
| |
| +---------------------------------------------------------------------------------+ |
| | EXPERIMENT: Top-Level Business Problem / Objective | |
| | (e.g., 'customer-churn-prediction-2026') | |
| +---------------------------------------------------------------------------------+ |
| | | |
| v v |
| +------------------------------------+ +------------------------------------+ |
| | RUN GROUP / TRIAL 1: | | RUN GROUP / TRIAL 2: | |
| | 'xgboost-feature-selection-v1' | | 'pytorch-transformer-baseline' | |
| +------------------------------------+ +------------------------------------+ |
| | | | |
| v v v |
| +---------------+ +---------------+ +---------------+ |
| | RUN 1: | | RUN 2: | | RUN 3: | |
| | Data Preproc | | Model Training| | Model Training| |
| | - S3 input | | - max_depth=6 | | - lr=1e-4 | |
| | - Clean code | | - val_auc=0.92| | - val_auc=0.95| |
| | - Output S3 | | - model.tar.gz| | - model.tar.gz| |
| +---------------+ +---------------+ +---------------+ |
+-----------------------------------------------------------------------------------------+
Core Hierarchy Components:
-
Experiment:
- The top-level logical container representing a specific ML project or business objective (e.g.,
fraud-detection-lightgbmorloan-default-v2). - Groups all iterations, regardless of whether they occurred days or months apart.
- The top-level logical container representing a specific ML project or business objective (e.g.,
-
Run Group / Trial:
- A logical grouping of related computational executions representing a specific strategy, hyperparameter tuning campaign, or pipeline execution (e.g.,
hyperband-search-round-1orpipeline-execution-20260816-01).
- A logical grouping of related computational executions representing a specific strategy, hyperparameter tuning campaign, or pipeline execution (e.g.,
-
Run / Trial Component:
- A single execution step within the workflow. Represents a distinct computational task: a data preprocessing job, a training job, or a model evaluation script.
- Directly captures inputs (datasets, configurations, hyperparameters), real-time logged metrics (loss, accuracy per epoch), and output artifacts (saved weights, confusion matrix charts).
2. Programmatic Tracking with the SageMaker Python SDK
The modern SageMaker Python SDK provides the sagemaker.experiments.Run context manager for lightweight, native tracking inside custom Python scripts, Jupyter notebooks, or SageMaker Training Jobs.
import sagemaker
from sagemaker.experiments import Run
from sagemaker.pytorch import PyTorch
session = sagemaker.Session()
role = sagemaker.get_execution_role()
# 1. Interactive Tracking in a Notebook / Local Script
with Run(
experiment_name='fraud-detection-model',
run_name='xgboost-baseline-run-01',
sagemaker_session=session
) as run:
# Log static configuration and hyperparameters
run.log_parameters({
'max_depth': 6,
'eta': 0.05,
'subsample': 0.8,
'objective': 'binary:logistic'
})
# Log input dataset artifact with S3 URI
run.log_artifact(name='training_data', local_path='s3://my-ml-bucket/train/v1.csv')
# Simulate training loop with step-by-step metric logging
for epoch in range(1, 11):
simulated_loss = 0.50 / epoch
simulated_auc = 0.80 + (0.015 * epoch)
run.log_metric(name='train_loss', value=simulated_loss, step=epoch)
run.log_metric(name='validation_auc', value=simulated_auc, step=epoch)
# Log output confusion matrix image
run.log_file(file_path='confusion_matrix.png', name='eval_confusion_matrix', is_output=True)
Passing Experiment Configuration to SageMaker Estimators
When launching managed training jobs via Estimators (estimator.fit()), you integrate directly with SageMaker Experiments using the experiment_config parameter:
estimator = PyTorch(
entry_point='train.py',
role=role,
instance_count=1,
instance_type='ml.g5.xlarge',
framework_version='2.1.0',
py_version='py310',
hyperparameters={'epochs': 20, 'lr': 0.001}
)
# Fit estimator while linking execution to Experiment and Run
estimator.fit(
inputs={'train': 's3://my-ml-bucket/train/'},
experiment_config={
'ExperimentName': 'fraud-detection-model',
'RunName': 'pytorch-resnet-run-02'
}
)
3. Data Lineage Tracking & Provenance DAGs
SageMaker Lineage Tracking automatically records the end-to-end lineage of machine learning workflows. It creates a formal Directed Acyclic Graph (DAG) capturing the precise relationships between raw datasets, processing containers, training runs, model artifacts, and deployed inference endpoints.
+-----------------------------------------------------------------------------------------+
| SAGEMAKER LINEAGE PROVENANCE GRAPH |
| |
| [Raw S3 Dataset] ----(ContributedTo)---> [Processing Job: Feature Engineering] |
| (Artifact) (Action) |
| | |
| (Produced) |
| v |
| [Git Commit / Image] -(UsedIn)-> [Processed S3 Data: Train/Val] (Artifact) |
| | |
| (ContributedTo) |
| v |
| [Training Job Run] (Action) |
| | |
| (Produced) |
| v |
| [model.tar.gz] (Artifact) |
| | |
| (ContributedTo) |
| v |
| [Model Package Registry] (Context) |
| | |
| (Deployed) |
| v |
| [Production Endpoint] (Action) |
+-----------------------------------------------------------------------------------------+
Lineage Entity Types:
- Artifact: Represents data or physical files (e.g., S3 raw CSVs, S3 processed Parquet files, Docker container images in ECR,
model.tar.gzweights). - Action: Represents a computational step that consumes or produces artifacts (e.g., a Processing Job, Training Job, Transform Job, or Endpoint deployment).
- Context: Represents a logical grouping or environment (e.g., an Experiment Run, a Pipeline execution, or a Model Package Group).
- Association: Defines the directional relationship between entities (
ContributedTo,Produced,DerivedFrom,AssociatedWith).
Querying Lineage for Regulatory Compliance
Regulated industries (e.g., finance, healthcare) require proving that a deployed production model was trained strictly on verified, compliant data. SageMaker Lineage APIs enable tracing the complete provenance graph backwards from an active endpoint:
from sagemaker.lineage.visualizer import LineageVisualizer
# Query and visualize full upstream lineage from Model Package ARN
lineage_viz = LineageVisualizer(sagemaker_session=session)
lineage_viz.render(model_package_arn='arn:aws:sagemaker:us-east-1:123456789012:model-package/fraud-model/2')
4. SageMaker Studio Visual Analysis & Leaderboards
SageMaker Studio provides an interactive visual dashboard for comparing metrics and parameters across experimental runs without writing manual visualization code.
+-----------------------------------------------------------------------------------------+
| SAGEMAKER STUDIO EXPERIMENTS INTERFACE |
| |
| Leaderboard Table: |
| +-------------------+------------+-----------+-----------+---------------+--------+ |
| | Run Name | Algorithm | max_depth | lr | val_pr_auc | Status | |
| +-------------------+------------+-----------+-----------+---------------+--------+ |
| | run-xgb-depth8-01 | XGBoost | 8 | 0.05 | 0.9482 (Best) | Done | |
| | run-xgb-depth6-02 | XGBoost | 6 | 0.05 | 0.9310 | Done | |
| | run-nn-mlp-01 | PyTorch | N/A | 1e-4 | 0.9125 | Done | |
| +-------------------+------------+-----------+-----------+---------------+--------+ |
| |
| Visualization Charts: |
| 1. Parallel Coordinates Plot: Maps multi-dimensional hyperparameters against metrics. |
| 2. Scatter Plots: Evaluates correlation between learning rate and validation loss. |
| 3. Metric Time-Series Overlay: Overlays training loss per epoch across 20 runs. |
+-----------------------------------------------------------------------------------------+
Key Studio Analytical Capabilities:
- Experiment Leaderboards: Automatically rank and sort runs based on any logged objective metric (e.g., highest
val_pr_aucor lowesteval_loss). - Parallel Coordinates Chart: Visualizes complex interactions between multiple hyperparameters (e.g.,
learning_rate,batch_size,num_layers) and the resulting target metric, identifying optimal parameter corridors. - Scatter & Metric Plots: Compares training curves in real time as training jobs execute, enabling engineers to spot divergence or overfitting immediately.
A machine learning team is designing an automated MLOps pipeline on AWS. The pipeline includes a data preprocessing step using SageMaker Processing, followed by model training using SageMaker Training, and a model evaluation step. The team wants to use SageMaker Experiments to track parameters, execution metrics, and artifacts across all three steps for every pipeline execution. How should the team structure the SageMaker Experiments hierarchy?
A financial services institution is undergoing a regulatory compliance audit for a credit approval model deployed on a SageMaker real-time endpoint. The auditor requires full proof of the exact S3 training dataset version, Git commit hash, hyperparameter settings, and Docker container image used to create the specific model artifact running on the endpoint. Which SageMaker feature directly provides this end-to-end auditable Directed Acyclic Graph (DAG)?
A deep learning practitioner is writing a custom PyTorch training script to execute on Amazon SageMaker. The practitioner wants to log training loss and validation accuracy at the end of each epoch so that the learning curves can be visualized and compared in real time in the SageMaker Studio Experiments UI. Which code pattern should the practitioner use in the training script?
A data science team ran 50 hyperparameter tuning training jobs with varying learning rates, batch sizes, tree depths, and regularization penalties. The lead engineer wants to quickly identify which hyperparameter combinations resulted in the highest validation PR-AUC and evaluate the correlation between multiple hyperparameters simultaneously. Which SageMaker Studio capability should the engineer use?