diff --git a/agent-evaluation/references/setup-guide.md b/agent-evaluation/references/setup-guide.md index 321949c..be942d9 100644 --- a/agent-evaluation/references/setup-guide.md +++ b/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:::mlflow-tracking-server/" +# or: arn:aws:sagemaker:::mlflow-app/ +export AWS_DEFAULT_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" diff --git a/agent-evaluation/scripts/setup_mlflow.py b/agent-evaluation/scripts/setup_mlflow.py index 948134d..7e63880 100644 --- a/agent-evaluation/scripts/setup_mlflow.py +++ b/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:::mlflow-app/ or .../mlflow-tracking-server/), + which only works when the `sagemaker-mlflow` plugin is installed (it registers the + `arn` tracking scheme). Without it, set_tracking_uri() 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