Connect to SageMaker MLflow

Establishes the connection to SageMaker Managed MLflow (the sagemaker-mlflow plugin + an ARN tracking URI + AWS credentials), then hands off to the requested MLflow skill (instrumenting-with-mlflow-tracing, agent-evaluation, retrieving-mlflow-traces, querying-mlflow-metrics, etc.), which inherit the connection via MLFLOW_TRACKING_URI.

Requires Python >= 3.10 (mlflow >= 3.8). Validate and report the fix — never auto-install the plugin or acquire credentials (mirrors the repo's Databricks pattern).

Step 1: Preconditions

• python --version -> must be 3.10+. If older, STOP and tell the user.
• Confirm AWS credentials resolve (env vars, aws configure / SSO, or an IAM role). If they error (ExpiredToken / no creds), STOP and ask the user to configure them. (scripts/verify_connection.py also checks this, via boto3.)

Step 2: Install the plugin

pip install sagemaker-mlflow

Required even if MLFLOW_TRACKING_URI is already an ARN — without it, mlflow.set_tracking_uri() fails with UnsupportedModelRegistryStoreURIException.

Step 3: Select the resource ARN

• If the user already provided an ARN (or MLFLOW_TRACKING_URI is set to arn:aws:sagemaker:...), use it as-is.
• Otherwise, discover it in the user's region:

  python scripts/discover_arns.py
  

  If exactly one is returned, use it; if several, list name/status/ARN and ask the user which to use. Both mlflow-app and mlflow-tracking-server ARNs work; the ARN is also shown on the resource's detail page in the SageMaker Studio console.

Then export it (the ARN is region-bound — a cross-region ARN will not authenticate):

export MLFLOW_TRACKING_URI="<arn>"    # mlflow-app/<id> or mlflow-tracking-server/<name>
export AWS_DEFAULT_REGION="<region>"  # AWS credentials resolved via the default chain (SigV4)

Step 4: Verify

python scripts/verify_connection.py "$MLFLOW_TRACKING_URI"

Checks all three prerequisites (plugin / credentials / ARN connectivity) and exits non-zero with a remediation hint on the first failure. Do not proceed until it passes.

After the connection verifies, hand off to the requested MLflow skill — it inherits the connected environment via MLFLOW_TRACKING_URI.

Troubleshooting

• UnsupportedModelRegistryStoreURIException — the sagemaker-mlflow plugin is not installed; pip install sagemaker-mlflow, then re-run scripts/verify_connection.py.
• AccessDenied / 403 — credentials are missing/expired, or the principal lacks MLflow permissions on the ARN. Check that credentials resolve, grant sagemaker-mlflow / sagemaker permissions, and confirm the ARN region matches AWS_DEFAULT_REGION.