Experiment Tracking That Survives a Team Change: Managed MLflow on SageMaker AI

Every SageMaker AI team eventually hits the same wall: a dozen training runs, three people, and no reliable answer to "which model is in production and what data produced it?" SageMaker AI now solves this with a fully managed MLflow tracking server you create as an AWS resource — no EC2 instance to patch, no self-hosted backend store, and native integration with the SageMaker Model Registry.

This tutorial sets one up end to end: create the tracking server, log a training run from a SageMaker training job, compare runs, register the winner to the Model Registry, and shut the server down so it stops billing.

Assumes the SageMaker Python SDK v2 (pip install "sagemaker>=2.215" "mlflow>=2.16" sagemaker-mlflow) and a role that can create SageMaker resources and read/write an S3 artifact bucket.

1. Create the managed tracking server

The tracking server is a first-class SageMaker resource. It needs an S3 location for artifacts and an IAM role it can assume to write there.

import boto3

sm = boto3.client("sagemaker")

sm.create_mlflow_tracking_server(
    TrackingServerName="na-mlflow-prod",
    ArtifactStoreUri="s3://my-ml-artifacts/mlflow/",
    TrackingServerSize="Small",
    RoleArn="arn:aws:iam::123456789012:role/service-role/AmazonSageMaker-ExecutionRole",
    AutomaticModelRegistration=True,
    WeeklyMaintenanceWindowStart="Sun:03:00",
)

Notes that matter:

• TrackingServerSize is Small, Medium, or Large. Small handles a handful of concurrent users and thousands of runs; size up only when the UI gets slow. You can update the size later without losing data.
• AutomaticModelRegistration=True means any model logged with mlflow.<flavor>.log_model(..., registered_model_name=...) also lands in the SageMaker Model Registry as a model package. This is the feature that makes managed MLflow worth using over a self-hosted server.
• Creation takes roughly 25 minutes. Poll it:

import time

while True:
    d = sm.describe_mlflow_tracking_server(TrackingServerName="na-mlflow-prod")
    print(d["TrackingServerStatus"])
    if d["TrackingServerStatus"] in ("Created", "CreateFailed"):
        break
    time.sleep(60)

TRACKING_ARN = d["TrackingServerArn"]
TRACKING_URL = d["TrackingServerUrl"]

The role needs s3:GetObject, s3:PutObject, and s3:ListBucket on the artifact prefix, plus sagemaker:CreateModelPackage if you use automatic registration. The managed MLflow documentation lists the full policy.

2. Grant access to your team

Access is IAM, not usernames and passwords. Attach a policy like this to the users or roles that should be able to log and browse runs:

{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Effect": "Allow",
      "Action": [
        "sagemaker-mlflow:AccessUI",
        "sagemaker-mlflow:*"
      ],
      "Resource": "arn:aws:sagemaker:us-east-1:123456789012:mlflow-tracking-server/na-mlflow-prod"
    }
  ]
}

For read-only analysts, narrow the actions to the Get* and Search* verbs. To open the UI, generate a presigned URL rather than exposing an endpoint:

print(sm.create_presigned_mlflow_tracking_server_url(
    TrackingServerName="na-mlflow-prod",
    ExpiresInSeconds=1800,
)["AuthorizedUrl"])

3. Log a training run

The sagemaker-mlflow plugin lets the MLflow client authenticate to the server with SigV4 when you point it at the tracking server ARN. Here is a training script that works both locally and inside a SageMaker training job:

# train.py
import argparse, os
import mlflow, mlflow.sklearn
import pandas as pd
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.metrics import f1_score, roc_auc_score
from sklearn.model_selection import train_test_split

p = argparse.ArgumentParser()
p.add_argument("--tracking-arn", default=os.environ.get("MLFLOW_TRACKING_ARN"))
p.add_argument("--n-estimators", type=int, default=200)
p.add_argument("--learning-rate", type=float, default=0.1)
p.add_argument("--max-depth", type=int, default=3)
args = p.parse_args()

mlflow.set_tracking_uri(args.tracking_arn)
mlflow.set_experiment("churn-classifier")

df = pd.read_csv(f"{os.environ.get('SM_CHANNEL_TRAIN', '.')}/train.csv")
X, y = df.drop(columns=["label"]), df["label"]
X_tr, X_te, y_tr, y_te = train_test_split(X, y, test_size=0.2, random_state=42)

with mlflow.start_run() as run:
    mlflow.log_params({
        "n_estimators": args.n_estimators,
        "learning_rate": args.learning_rate,
        "max_depth": args.max_depth,
        "rows": len(df),
    })

    model = GradientBoostingClassifier(
        n_estimators=args.n_estimators,
        learning_rate=args.learning_rate,
        max_depth=args.max_depth,
    ).fit(X_tr, y_tr)

    preds = model.predict(X_te)
    probs = model.predict_proba(X_te)[:, 1]
    mlflow.log_metrics({
        "f1": f1_score(y_te, preds),
        "roc_auc": roc_auc_score(y_te, probs),
    })

    mlflow.sklearn.log_model(
        model,
        name="model",
        input_example=X_te.head(5),
    )
    print("run_id:", run.info.run_id)

mlflow.autolog() covers most of the parameter and metric logging automatically for scikit-learn, XGBoost, PyTorch Lightning, and Transformers — but log the things autolog cannot know: dataset version, feature set name, git SHA, and the S3 URI of the training data. Those are what make a run reproducible six months later.

Launch it as a SageMaker training job:

from sagemaker.sklearn.estimator import SKLearn

est = SKLearn(
    entry_point="train.py",
    source_dir=".",
    role=role,
    framework_version="1.2-1",
    instance_type="ml.m5.xlarge",
    instance_count=1,
    environment={"MLFLOW_TRACKING_ARN": TRACKING_ARN},
    hyperparameters={"n_estimators": 300, "learning_rate": 0.05},
)
est.fit({"train": "s3://my-data/churn/train/"})

Because the tracking ARN is passed as an environment variable, the same script runs unchanged in a notebook, in a training job, and inside a SageMaker Pipelines step.

4. Compare runs and pick a winner

The UI is the standard MLflow experiment view — sort by metric, select several runs, and use the parallel-coordinates plot to see which hyperparameter actually moved the number. From code:

import mlflow

mlflow.set_tracking_uri(TRACKING_ARN)
runs = mlflow.search_runs(
    experiment_names=["churn-classifier"],
    filter_string="metrics.roc_auc > 0.85 and params.max_depth = '3'",
    order_by=["metrics.roc_auc DESC"],
    max_results=10,
)
print(runs[["run_id", "params.n_estimators", "metrics.f1", "metrics.roc_auc"]])
best_run_id = runs.iloc[0]["run_id"]

Two habits worth enforcing on a team:

• Tag runs with intent. mlflow.set_tag("purpose", "candidate") versus "exploration" keeps a 400-run experiment readable.
• Log the evaluation dataset as an artifact or a dataset reference, not just the metric. A metric without its dataset is a rumour.

5. Register the winner to the SageMaker Model Registry

With AutomaticModelRegistration=True, registering in MLflow creates a SageMaker model package group and version:

result = mlflow.register_model(
    model_uri=f"runs:/{best_run_id}/model",
    name="churn-classifier",
)
print(result.name, result.version)

Check the SageMaker side:

pkgs = sm.list_model_packages(ModelPackageGroupName="churn-classifier")
arn = pkgs["ModelPackageSummaryList"][0]["ModelPackageArn"]

sm.update_model_package(
    ModelPackageArn=arn,
    ModelApprovalStatus="Approved",
    ApprovalDescription="ROC-AUC 0.891 on 2026-05 holdout; approved by ML lead.",
)

Approval is the hand-off point: an EventBridge rule on ModelPackageApprovalStatus changing to Approved can trigger the deployment pipeline, so nobody deploys a model straight out of a notebook. That deployment side is covered in A Minimal MLOps Loop: SageMaker Pipelines + Model Registry.

You can also deploy directly from the MLflow model URI:

from sagemaker.serve import ModelBuilder, SchemaBuilder

builder = ModelBuilder(
    mode="SageMakerEndpoint",
    role_arn=role,
    model_metadata={"MLFLOW_MODEL_PATH": f"runs:/{best_run_id}/model"},
    schema_builder=SchemaBuilder(X_te.head(1), y_te.head(1)),
)
predictor = builder.build().deploy(
    initial_instance_count=1, instance_type="ml.m5.xlarge"
)

6. Tracking GenAI work, not just classifiers

The same server handles LLM experiments, which is where most 2026 teams actually need discipline. mlflow.evaluate() scores a prompt-and-model combination against a labelled set and logs the results as a run:

import mlflow, pandas as pd

eval_df = pd.DataFrame({
    "inputs": ["Summarise this ticket: ...", "Summarise this ticket: ..."],
    "ground_truth": ["Card declined, retry issued.", "Refund requested."],
})

with mlflow.start_run(run_name="summariser-v3-prompt-b"):
    mlflow.log_param("base_model", "meta-llama/Llama-3.1-8B-Instruct")
    mlflow.log_param("prompt_version", "b")
    mlflow.log_param("temperature", 0.2)
    mlflow.evaluate(
        model=my_endpoint_fn,          # any callable taking a DataFrame
        data=eval_df,
        targets="ground_truth",
        model_type="text-summarization",
    )

Enabling MLflow tracing (mlflow.<library>.autolog() for supported LLM frameworks) captures the full span tree for a RAG or agent call — retrieval, prompt, tokens, latency — so a regression can be traced to the retrieval step instead of blamed on "the model". If you are still choosing between Bedrock and SageMaker AI for the generation layer, see our 2026 decision guide.

7. Cost and lifecycle

A managed tracking server bills per hour while it is in the Created state, plus S3 for artifacts. It does not scale to zero. Practical hygiene:

# Pause overnight or between projects
sm.stop_mlflow_tracking_server(TrackingServerName="na-mlflow-prod")
sm.start_mlflow_tracking_server(TrackingServerName="na-mlflow-prod")

# Delete when the project ends (artifacts stay in S3)
sm.delete_mlflow_tracking_server(TrackingServerName="na-mlflow-prod")

• Run one server per team or environment, not one per project — experiments namespace runs perfectly well.
• Stop dev servers on a schedule with an EventBridge rule; a stopped server keeps its metadata.
• Set an S3 lifecycle policy on the artifact prefix so five-gigabyte checkpoints from abandoned experiments do not accumulate for years.
• Check the current per-hour rate for each server size on the SageMaker AI pricing page before choosing Medium.

Common failure modes

Symptom	Cause	Fix
AccessDeniedException on mlflow.set_tracking_uri(arn)	Caller lacks sagemaker-mlflow:* on the server ARN	Attach the policy in step 2
Runs log but artifacts fail	Tracking server role cannot write the S3 prefix	Fix the role's bucket policy, not the caller's
MLFLOW_TRACKING_URI ignored in a training job	Plugin missing in the container	Add sagemaker-mlflow to requirements.txt in source_dir
Models log but never appear in the Model Registry	AutomaticModelRegistration disabled	update_mlflow_tracking_server(..., AutomaticModelRegistration=True)
UI link expires mid-session	Presigned URL TTL	Regenerate with a longer ExpiresInSeconds

Where this fits

Managed MLflow closes the gap between "someone trained a good model" and "we can prove which model, which data, and who approved it". Combine it with Pipelines for repeatable training, the Model Registry for approval gates, and Model Monitor for what happens after deployment, and you have an audit trail that survives a team change.

Need a hand standing up experiment tracking and a governed release path on SageMaker AI? Talk to us.