Create a tracking server using the Amazon CLI - Amazon SageMaker
PrerequisitesSet up Amazon CLI modelSet up MLflow infrastructureCreate MLflow tracking server
Services or capabilities described in Amazon Web Services documentation might vary by Region. To see the differences applicable to the China Regions, see Getting Started with Amazon Web Services in China (PDF).

Create a tracking server using the Amazon CLI

You can create a tracking server using the Amazon CLI for more granular security customization.

Prerequisites

To create a tracking server using the Amazon CLI, you must have the following:

  • Access to a terminal. This can include local IDEs, an Amazon EC2 instance, or Amazon CloudShell.

  • Access to a development environment. This can include local IDEs or a Jupyter notebook environment within Studio or Studio Classic.

  • A configured Amazon CLI installation. For more information, see Configure the Amazon CLI.

  • An IAM role with appropriate permissions. The following steps require your environment to have iam:CreateRole, iam:CreatePolicy, iam:AttachRolePolicy, and iam:ListPolicies permissions. These permissions are needed on the role that is being used to run the steps in this user guide. The instructions in this guide create an IAM role that is used as the execution role of the MLflow Tracking Server so that it can access data in your Amazon S3 buckets. Additionally, a policy is created to give the IAM role of the user that is interacting with the Tracking Server via the MLflow SDK permission to call MLflow APIs. For more information, see Modifying a role permissions policy (console) .

    If using a SageMaker Studio Notebook, update the service role for your Studio user profile with these IAM permissions. To update the service role, navigate to the SageMaker console and select the domain you are using. Then, under the domain, select the user profile you are using. You will see the service role listed there. Navigate to the IAM console, search for the service role under Roles, and update your role with a policy that allows the iam:CreateRole, iam:CreatePolicy, iam:AttachRolePolicy, and iam:ListPolicies actions.

Set up Amazon CLI model

Follow these command line steps within a terminal to set up the Amazon CLI for Amazon SageMaker with MLflow.

  1. Install an updated version of the Amazon CLI. For more information, see Install or update to the latest version of the Amazon CLI in the Amazon CLI User Guide.

  2. Verify that the Amazon CLI is installed using the following command: 

    aws sagemaker help

    Press q to exit the prompt.

    For troubleshooting help, see Troubleshoot common setup issues.

Set up MLflow infrastructure

The following section shows you how to set up an MLflow Tracking Server along with the Amazon S3 bucket and IAM role needed for the tracking server.

Create an S3 bucket

Within your terminal, use the following commands to create a general purpose Amazon S3 bucket:

Note

The Amazon S3 bucket used for your artifact store must be in the same Amazon Web Services Region as your tracking server.

bucket_name=bucket-name
region=valid-region

aws s3api create-bucket \
  --bucket $bucket_name \
  --region $region \
  --create-bucket-configuration LocationConstraint=$region

The output should look similar to the following:

{
    "Location": "/bucket-name"
}

Set up IAM trust policies

Use the following steps to create an IAM trust policy. For more information about roles and trust policies, see Roles terms and concepts in the Amazon Identity and Access Management User Guide.

  1. Within your terminal, use the following command to create a file called mlflow-trust-policy.json.

    cat <<EOF > /tmp/mlflow-trust-policy.json
{
     "Version": "2012-10-17",
     "Statement": [
         {
             "Effect": "Allow",
             "Principal": {
                 "Service": [                     
                      "sagemaker.amazonaws.com"
                 ]
             },
             "Action": "sts:AssumeRole"
         }
     ]
 }
EOF

  2. Within your terminal, use the following command to create a file called custom-policy.json.

    cat <<EOF > /tmp/custom-policy.json
{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Action": [
                "s3:Get*",
                "s3:Put*",
                "sagemaker:AddTags",
                "sagemaker:CreateModelPackageGroup",
                "sagemaker:CreateModelPackage",
                "sagemaker:DescribeModelPackageGroup",
                "sagemaker:UpdateModelPackage",
                "s3:List*"
            ],
            "Resource": "*"
        }
    ]
}
EOF

  3. Use the trust policy file to create a role. Then, attach IAM role policies that allow MLflow to access Amazon S3 and SageMaker Model Registry within your account. MLflow must have access to Amazon S3 for your tracking server's artifact store and SageMaker Model Registry for automatic model registration.

    Note

    If you are updating an existing role, use the following command instead: aws iam update-assume-role-policy --role-name $role_name --policy-document file:///tmp/mlflow-trust-policy.json.

    role_name=role-name

aws iam  create-role \
  --role-name $role_name \
  --assume-role-policy-document file:///tmp/mlflow-trust-policy.json

aws iam put-role-policy \
  --role-name $role_name \
  --policy-name custom-policy \
  --policy-document file:///tmp/custom-policy.json

role_arn=$(aws iam get-role --role-name  $role_name --query 'Role.Arn' --output text)

Create MLflow tracking server

Within your terminal, use the create-mlflow-tracking-server API to create a tracking server in the Amazon Web Services Region of your choice. This step can take up to 25 minutes.

You can optionally specify the size of your tracking server with the parameter --tracking-server-config. Choose between "Small", "Medium", and "Large". The default MLflow Tracking Server configuration size is "Small". You can choose a size depending on the projected use of the tracking server such as the volume of data logged, number of users, and frequency of use. For more information, see MLflow Tracking Server sizes.

The following command creates a new tracking server with automatic model registration enabled. To deactivate automatic model registration, specify --no-automatic-model-registration.

After creating your tracking server, you can launch the MLflow UI. For more information, see Launch the MLflow UI using a presigned URL.

Note

It may take up to 25 minutes to complete tracking server creation. If the tracking server takes over 25 minutes to create, check that you have the necessary IAM permissions. For more information on IAM permissions, see Set up IAM permissions for MLflow. When you successfully create a tracking server, it automatically starts.

ts_name=tracking-server-name
region=valid-region

aws sagemaker create-mlflow-tracking-server \
 --tracking-server-name $ts_name \
 --artifact-store-uri s3://$bucket_name \
 --role-arn $role_arn \
 --automatic-model-registration \
 --region $region

The output should be similar to the following: 

{
    "TrackingServerArn": "arn:aws:sagemaker:region:123456789012:mlflow-tracking-server/tracking-server-name"
}
Important

Take note of the tracking server ARN for later use. You will also need the $bucket_name for clean up steps.

Describe MLflow Tracking Server

Check the status of your MLflow Tracking Server creation with the describe-mlflow-tracking-server API.

aws sagemaker describe-mlflow-tracking-server \
  --tracking-server-name $ts_name \
  --region $region

When the MLflow Tracking Server creation is still in progress, the TrackingServerStatus is "Creating". When the MLflow tracking server creation is complete, the TrackingServerStatus is "Created". The output should be similar to the following: 

{
    "TrackingServerArn": "arn:aws:sagemaker:region:123456789012:mlflow-tracking-server/tracking-server-name",
    "MlflowTrackingServerName": "tracking-server-name",
    "CreationTime": "2024-03-15T19:41:43+00:00",
    "LastModifiedTime": "2024-03-15T19:41:43+00:00",
    "CreatedBy": {},
    "LastModifiedBy": {},
    "ArtifactStoreUri": "s3://bucket-name",
    "TrackingServerConfig": "Small",
    "MlflowVersion": "v2.11.3",
    "TrackingServerStatus": "Created"
}

List MLflow Tracking Server

List MLflow Tracking Servers with the list-mlflow-tracking-servers API.

aws sagemaker list-mlflow-tracking-servers \
  --region $region

Your output should look similar to the following:

{
    "TrackingServerSummaries": [
        {
            "TrackingServerArn": "arn:aws:sagemaker:region:123456789012:mlflow-tracking-server/tracking-server-name",
            "MlflowTrackingServerName": "tracking-server-name",
            "CreationTime": "2024-04-11T16:58:27+00:00",
            "LastModifiedTime": "2024-04-11T16:58:27+00:00",
            "TrackingServerStatus": "CreatePending",
            "MlflowVersion": "v2.11.3"
        }
    ]
}

By default, tracking servers are listed in descending order by creation time. To change the list order, you can optionally specify --sort-order to be Ascending.

You can optionally filter the listed tracking servers by --tracking-server-status, such as "Creating" or "Created".

Use the --created-after filter to only list tracking servers created after a specific date and time. Listed tracking servers are shown with a date and time such as "2024-03-16T01:46:56+00:00". The --created-after parameter takes in a Unix timestamp. To convert a date and time into a Unix timestamp, see EpochConverter.

aws sagemaker list-mlflow-tracking-servers \
  --region $region \
  --sort-order Ascending \
  --tracking-server-status Created \
  --created-after 1712852168

If you have more than one tracking server in the same Amazon Web Services Region, you can use the --next-token parameter to iterate through your tracking servers. 

# List one tracking server in a specified Amazon Web Services Region to get a NextToken
aws sagemaker list-mlflow-tracking-servers \
  --max-results 1 \
  --region $region

# Save the NextToken for this listed tracking server in a variable
next_token=$(aws experiments-beta list-mlflow-tracking-servers \
  --max-results 1 \
  --region $region | jq -r .NextToken)

# Use the NextToken to list the next tracking server and get a new NextToken
aws sagemaker list-mlflow-tracking-servers \
  --max-results 1 \
  --region $region \
  --next-token $next_token

To see all possible list options, run the following command:

aws sagemaker list-mlflow-tracking-servers help

Stop or start the MLflow Tracking Server

To stop the tracking server, use the following command: 

aws sagemaker stop-mlflow-tracking-server \
  --tracking-server-name $ts_name \
  --region $region

To start the tracking server, use the following command:

Note

It may take up to 25 minutes to start your tracking server.

aws sagemaker start-mlflow-tracking-server \
  --tracking-server-name $ts_name \
  --region $region

Update the MLflow Tracking Server

You can update the artifact storage Amazon S3 bucket, tracking server size, automatic model registration configuration, or weekly maintenance window at any time. A tracking server must be stopped to be updated.

To update the tracking server and change the artifact store URI, use the following command: 

aws sagemaker update-mlflow-tracking-server \
  --tracking-server-name $ts_name \
  --artifact-store-uri $updated-artifact-store-uri \ 
  --region $region