Amazon SageMaker AI streams benchmark results to MLflow

Amazon SageMaker AI now streams benchmark and optimized inference recommendation results directly into a SageMaker MLflow App, so teams can track metrics, parameters and charts from multiple runs in one place in real time. The integration uses MlflowConfig in the job OutputConfig to publish runs to a chosen SageMaker MLflow App.

What did Amazon add and how does it work?

Amazon SageMaker AI added native MLflow integration for optimized inference recommendation jobs and benchmarking jobs, streaming metrics, parameters and charts into a SageMaker MLflow App as the job runs. Developers supply MlflowConfig in OutputConfig when creating a benchmark or recommendation job; SageMaker then streams results to the MLflow App ARN and the specified experiment and run name.

The integration consolidates multiple jobs under the same MlflowExperimentName so runs appear side by side in the MLflow UI. The announcement shows example code that sets REGION = "us-west-2", an MLFLOW_APP_ARN, an ENDPOINT_NAME and an S3 output bucket named using the pattern mlflow-sagemaker-{REGION}-{ACCOUNT_ID}. Example job creation calls include create_ai_benchmark_job and create_ai_workload_config with an OutputConfig that contains "MlflowConfig": {"MlflowResourceArn": MLFLOW_APP_ARN, "MlflowExperimentName": MLFLOW_EXPERIMENT, "MlflowRunName": BENCH_RUN_NAME}.

How do teams set it up and what must they grant?

Set up requires creating a SageMaker MLflow App in SageMaker Studio, adding the app ARN to your job execution role, and passing MlflowConfig when you submit benchmark or recommendation jobs. The post lists three setup steps: create an MLflow App in SageMaker Studio, grant the job’s execution role the permission "sagemaker-mlflow:*" on the MLflow App ARN, and pass MlflowConfig to the job.

The execution role must also allow SageMaker jobs to run, invoke endpoints, and write to S3. The required permissions include AmazonSageMakerFullAccess (or scoped equivalent), "sagemaker-mlflow:*", sagemaker:CallMlflowAppApi, sagemaker:DescribeMlflowApp, sagemaker:InvokeEndpoint and sagemaker:InvokeEndpointWithResponseStream, plus read/write to the configured S3 output bucket. The S3 output bucket must be in the same Region as the job.

The integration supports only SageMaker MLflow Apps and does not stream to self-hosted MLflow tracking servers. The tooling.version must be 0.8.0 or later to use MLflow nested run support.

What does a concrete example look like?

The post walks through a reproducible example using Qwen/Qwen2-0.5B-Instruct deployed to an ml.g6.12xlarge endpoint. A benchmark workload example sets tooling.version to "0.8.0", TOKENIZER = "Qwen/Qwen2-0.5B-Instruct", concurrency = 1 and request_count = 3, with prompt_input_tokens_mean = 32 and output_tokens_mean = 16. A recommendation workload example uses prompt_input_tokens_mean = 1600, output_tokens_mean = 600 and request_count = 100.

SageMaker AI benchmark jobs evaluate an existing endpoint; recommendation jobs provision their own endpoints internally during evaluation. The post also contrasts example instance comparisons such as qwen2-0.5b on ml.g4dn.12xlarge versus ml.p4d.24xlarge as use cases for side-by-side comparison in MLflow. The end-to-end notebook walkthrough typically takes 45–120 minutes depending on endpoint readiness, model size, workload configuration, instance availability and recommendation job search space.

Why it matters

This integration removes manual data consolidation for inference optimization by centralizing metrics, parameters and charts under a shared MLflow experiment. Teams can monitor long-running benchmark and recommendation jobs live, stop a run early if throughput is not meeting expectations, and retain a queryable audit trail of job parameters, timestamps and emitted artifacts for reproducibility and governance.

The change reduces friction when comparing instance types, model configurations, batch sizes or speculative decoding strategies, and it gives teams a single source of truth for optimization work where runs are comparable in the MLflow UI.

What to watch

Watch for wider adoption of the SageMaker MLflow App in teams that run many benchmark and recommendation jobs, and for notebook updates or tooling.version changes beyond 0.8.0 that expand nested-run support. Also check whether future releases add support for streaming to self-hosted MLflow tracking servers or additional MLflow features in the SageMaker UI.