mcpField
S
sagemaker-mlflow

sagemaker-mlflow

github.com/mlflowCloud & DevOps82394 installsUpdated 2d ago
Repository

About sagemaker-mlflow

sagemaker-mlflow is a skill for connecting MLflow to SageMaker-managed tracking resources. It checks Python, the plugin, AWS credentials, and the resource ARN, then verifies the connection before handing work to another MLflow skill.

Key features

  1. 01Check Python version and plugin
  2. 02Validate AWS credentials
  3. 03Select or discover resource ARN
  4. 04Verify tracking connection
  5. 05Pass connection to other skills

Use cases

  1. 01Set up an MLflow tracking URI
  2. 02Find a SageMaker resource ARN
  3. 03Diagnose connection prerequisites

Install

Copy the command and run it in a terminal.

$npx skills add https://github.com/mlflow/skills --skill sagemaker-mlflow

After installing, type/sagemaker-mlflowin Claude Code to confirm it loaded.

Bundled filesscripts/discover_arns.pyscripts/verify_connection.py

Details

From the project README

Connect to SageMaker MLflow

Connect using the sagemaker-mlflow plugin, an ARN tracking URI and AWS credentials. Other MLflow skills inherit the connection through 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: Check the plugin

If missing, stop and give the user this installation command:

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 — see Step 2, then repeat verification.
  • 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.

Comments

Markdown supported