6.3 Experiment Environments, Tracking & Lineage with Experiments and ML Metadata
Key Takeaways
- Experiments on Agent Platform tracks the steps, inputs such as parameters and datasets, and outputs such as models and metrics of each experiment run.
- Autologging in the Agent Platform SDK uses MLflow autologging and captures parameters and metrics from frameworks including Keras, scikit-learn, XGBoost, LightGBM, and PyTorch Lightning.
- Time-series metrics logged with log_time_series_metrics are stored in a managed TensorBoard instance tied to the experiment or run.
- ML Metadata records lineage as a graph in which artifacts and executions are nodes, events are edges, and contexts group related nodes such as a pipeline run.
- Experiment runs cost nothing extra; you pay only for the resources used during the experiment.
Section 2.3 of the exam guide covers choosing the appropriate Google Cloud environment for development and experimentation (Experiments on Agent Platform, Agent Platform Pipelines, and Kubeflow Pipelines) and tracking and comparing model artifacts, versions, and lineage (Experiments and ML Metadata). Evaluation metrics themselves are covered in Chapter 7.
Choosing the Experimentation Environment
| Environment | Best for | Trade-offs |
|---|---|---|
| Notebook + Experiments SDK (Colab Enterprise or Workbench) | Early, interactive exploration: trying features, algorithms, and prompts | Easy to lose reproducibility if code and data versions aren't logged |
| Agent Platform Pipelines (serverless KFP) | Repeatable multi-step experiments: data prep → train → evaluate, with automatic metadata and lineage. Runs can be grouped into an experiment | Needs code packaged as components. Best once the workflow stabilizes |
| Kubeflow Pipelines on GKE (self-managed) | Organizations standardized on Kubeflow, needing portability across on-premises and clouds, or deep Kubernetes customization | You operate the cluster, upgrades, and security |
| Custom training jobs with Experiments | Full-scale or distributed training runs that must be compared | Longer iteration loop than notebooks |
Pipelines built with the Kubeflow Pipelines (KFP) SDK can run on Agent Platform Pipelines, so teams can prototype locally or on self-managed Kubeflow and move to the managed service. Given the exam's preference for managed services, Agent Platform Pipelines is usually the right answer unless the scenario explicitly requires self-managed Kubeflow.
Framework matters too. Experiments works with any Python framework through the SDK and integrates deeply with TensorFlow. Autologging, built on MLflow autologging, captures parameters and metrics from Fastai, Gluon, Keras, LightGBM, PyTorch Lightning, scikit-learn, Spark, Statsmodels, and XGBoost.
Experiments on Agent Platform
An experiment groups experiment runs. Each run records:
- Steps (preprocessing, training)
- Inputs (algorithm, parameters, dataset)
- Outputs (models, checkpoints, metrics)
Experiment runs cost nothing extra. You pay only for compute and storage used.
from google.cloud import aiplatform
aiplatform.init(project=PROJECT, location=REGION,
experiment="churn-xgb", experiment_tensorboard=TB_RESOURCE)
aiplatform.start_run("xgb-depth6-lr0p05")
aiplatform.log_params({"max_depth": 6, "learning_rate": 0.05, "data_version": "v12"})
# ... train ...
aiplatform.log_metrics({"auc_pr": 0.61, "recall_at_p90": 0.44})
aiplatform.end_run()
| API | Records |
|---|---|
log_params | Hyperparameters and configuration |
log_metrics | Summary metrics for the run |
log_time_series_metrics | Per-step metrics such as loss curves, stored in a managed TensorBoard instance (documented as Vertex AI TensorBoard) |
autolog() | Automatic parameter and metric capture for supported frameworks |
PipelineJob.submit(experiment=...) | Adds a pipeline run to an experiment so pipeline runs can be compared |
The console shows runs side by side, and the SDK returns runs as DataFrames for analysis. TensorBoard adds loss curves, histograms, computation graphs, embedding projections, and a shareable link, and the TensorBoard profiler helps diagnose slow training.
ML Metadata and Lineage
Agent Platform ML Metadata builds on the open-source ML Metadata (MLMD) library and stores metadata as a graph inside a regional MetadataStore (usually one per project).
| Concept | Meaning | Example |
|---|---|---|
| Artifact | Data produced or consumed by the workflow | Dataset, model, evaluation metrics |
| Execution | A workflow step with runtime parameters | Training step, validation step |
| Event | Edge linking an artifact to an execution as input or output | "Training execution consumed dataset v12" |
| Context | Groups artifacts and executions | One pipeline run, or one experiment run |
| MetadataSchema | Type definition (system schemas such as system.Model, or custom schemas) | Validates metadata fields |
Pipelines record this graph automatically. Custom code can write it with the SDK. Lineage answers audit questions such as:
- Which dataset and hyperparameters produced the model now in production?
- Which models were trained on a dataset that turned out to have a labeling bug?
- Which model version produced a specific batch prediction?
Versions and the Model Registry
Experiments and metadata track how a model was made. Model Registry tracks which versions exist and which is deployed. Log the winning run's model to the registry as a new version, with a link back to its experiment run. Use aliases such as default or champion to show which version is promoted (Chapter 13).
Worked Scenario
A regulated lender must show auditors, for any production credit model, the training data snapshot, preprocessing code version, hyperparameters, evaluation results, and approver.
- Training runs as an Agent Platform Pipeline submitted with
experiment="credit-risk", so every run is comparable and ML Metadata records the lineage automatically. - The pipeline takes the managed dataset version and the container image digest as inputs, so both show up in lineage.
- Evaluation metrics are logged to the run. The approved model goes to Model Registry with labels for the approver and ticket ID.
- Auditors follow lineage from the deployed model version back to its dataset and run, with no spreadsheets involved.
What to Log in Every Run
A run is only reproducible if you can rebuild it. Log at least:
| Category | Examples |
|---|---|
| Data | Managed dataset version or BigQuery snapshot time, row counts, label distribution |
| Code | Git commit, container image digest, preprocessing version |
| Configuration | Hyperparameters, random seeds, hardware (machine type, accelerator count) |
| Results | Summary metrics, per-slice metrics, loss curves, evaluation artifacts |
| Context | Owner, purpose ("test focal loss for class imbalance"), links to tickets |
Exam Traps
- Tracking experiments in spreadsheets or notebook markdown when the scenario asks for comparable, auditable runs.
- Choosing self-managed Kubeflow on GKE when nothing requires it. Agent Platform Pipelines removes cluster operations.
- Confusing Experiments (compare runs) with Model Registry (manage versions for deployment) and ML Metadata (lineage graph). Production systems use all three together.
A team wants loss curves from every training step in each experiment run to appear in Experiments on Agent Platform for comparison. What must they configure?
In Agent Platform ML Metadata, what represents the link recording that a training step consumed a specific dataset?
A company with no Kubernetes expertise wants repeatable, comparable experiments that chain preprocessing, training, and evaluation with automatic lineage tracking. Which environment should it choose?