docs/setup: support SageMaker Managed MLflow as a tracking backend

agent-evaluation/references/setup-guide.md

@@ -37,6 +37,13 @@ If not installed or version too old:
 uv pip install mlflow>=3.8.0
 ```
 
+**SageMaker Managed MLflow:** also install the auth plugin (it registers the `arn://` tracking
+scheme; without it an ARN tracking URI fails with `UnsupportedModelRegistryStoreURIException`):
+
+```bash
+uv pip install sagemaker-mlflow
+```
+
 ## Step 2: Configure Environment
 
 ### Quick Setup (Recommended - 90% of cases)
@@ -109,6 +116,12 @@ Choose your tracking server type:
 # For Databricks
 export MLFLOW_TRACKING_URI="databricks"
 
+# For SageMaker Managed MLflow (requires `pip install sagemaker-mlflow` from Step 1)
+# Use the resource ARN as the tracking URI (Aloy tracking-server or Mercury mlflow-app):
+export MLFLOW_TRACKING_URI="arn:aws:sagemaker:<region>:<account>:mlflow-tracking-server/<name>"
+# or: arn:aws:sagemaker:<region>:<account>:mlflow-app/<id>
+export AWS_DEFAULT_REGION="<region>"   # AWS creds resolved via the default chain (SigV4)
+
 # For local server (start server first: mlflow server --host 127.0.0.1 --port 5000 &)
 export MLFLOW_TRACKING_URI="http://127.0.0.1:5000"
 

agent-evaluation/scripts/setup_mlflow.py

@@ -48,6 +48,26 @@ def check_mlflow_installed() -> bool:
         return False
 
 
+def ensure_backend_plugin(tracking_uri: str) -> None:
+    """Ensure the backend plugin for the chosen tracking URI is installed.
+
+    SageMaker Managed MLflow uses an ARN tracking URI
+    (arn:aws:sagemaker:<region>:<acct>:mlflow-app/<id> or .../mlflow-tracking-server/<name>),
+    which only works when the `sagemaker-mlflow` plugin is installed (it registers the
+    `arn` tracking scheme). Without it, set_tracking_uri(<arn>) raises
+    UnsupportedModelRegistryStoreURIException.
+    """
+    if tracking_uri.startswith("arn:aws:sagemaker:"):
+        try:
+            import sagemaker_mlflow  # noqa: F401
+
+            print("✓ sagemaker-mlflow plugin installed (SageMaker Managed MLflow)")
+        except ImportError:
+            print("✗ SageMaker Managed MLflow detected but sagemaker-mlflow is not installed")
+            print("  Install with: pip install sagemaker-mlflow")
+            sys.exit(1)
+
+
 def detect_databricks_profiles() -> list[str]:
     """Detect available and valid Databricks profiles.
 
@@ -280,6 +300,9 @@ def main():
     # Configure tracking URI (auto-detects if not provided)
     tracking_uri = configure_tracking_uri(args.tracking_uri)
 
+    # Ensure backend plugin (e.g. sagemaker-mlflow for SageMaker ARN URIs) is installed
+    ensure_backend_plugin(tracking_uri)
+
     # Configure experiment ID (auto-detects if not provided)
     experiment_id = configure_experiment_id(
         tracking_uri, args.experiment_id, args.experiment_name, args.create