# Accelerate generative AI development using managed MLflow on Amazon SageMaker AI Fully managed MLflow 3.0 on Amazon SageMaker AI enables you to accelerate generative AI by making it easier to track experiments and monitor performance of models and AI applications using a single tool. **Generative AI development with MLflow 3.0** As customers across industries accelerate their generative AI development, they require capabilities to track experiments, observe behavior, and evaluate performance of models and AI applications. Data scientists and developers lack tools for analyzing the performance of models and AI applications from experimentation to production, making it hard to root cause and resolve issues. Teams spend more time integrating tools than improving their models or generative AI applications. Training or fine-tuning generative AI and machine learning is an iterative process that requires experimenting with various combinations of data, algorithms, and parameters, while observing their impact on model accuracy. The iterative nature of experimentation results in numerous model training runs and versions, making it challenging to track the best performing models and their configurations. The complexity of managing and comparing iterative training runs increases with GenAI, where experimentation involves not only fine-tuning models but also exploring creative and diverse outputs. Researchers must adjust hyperparameters, select suitable model architectures, and curate diverse datasets to optimize both the quality and creativity of the generated content. Evaluating generative AI models requires both quantitative and qualitative metrics, adding another layer of complexity to the experimentation process. Experimentation tracking capabilities in MLflow 3.0 on Amazon SageMaker AI enables you to track, organize, view, analyze, and compare iterative ML experimentation to gain comparative insights and register and deploy your best performing models. Tracing capabilities in fully managed MLflow 3.0 enables you to record the inputs, outputs, and metadata at every step of a generative AI application, helping you to quickly identify the source of bugs or unexpected behaviors. By maintaining records of each model and application version, fully managed MLflow 3.0 offers traceability to connect AI responses to their source components, allowing you to quickly trace an issue directly to the specific code, data, or parameters that generated it. This dramatically reduces troubleshooting time and enables teams to focus more on innovation. ## MLflow integrations Use MLflow while training and evaluating models to find the best candidates for your use case. You can compare model performance, parameters, and metrics across experiments in the MLflow UI, keep track of your best models in the MLflow Model Registry, automatically register them as a SageMaker AI model, and deploy registered models to SageMaker AI endpoints. **Amazon SageMaker AI with MLflow** Use MLflow to track and manage the experimentation phase of the machine learning (ML) lifecycle with AWS integrations for model development, management, deployment, and tracking. **Amazon SageMaker Studio** Create and manage tracking servers, run notebooks to create experiments, and access the MLflow UI to view and compare experiment runs all through Studio. **SageMaker Model Registry** Manage model versions and catalog models for production by automatically registering models from MLflow Model Registry to SageMaker Model Registry. For more information, see [Automatically register SageMaker AI models with SageMaker Model Registry](mlflow-track-experiments-model-registration.md). **SageMaker AI Inference** Prepare your best models for deployment on a SageMaker AI endpoint using `ModelBuilder`. For more information, see [Deploy MLflow models with `ModelBuilder`](mlflow-track-experiments-model-deployment.md). **AWS Identity and Access Management** Configure access to MLflow using role-based access control (RBAC) with IAM. Write IAM identity policies to authorize the MLflow APIs that can be called by a client of an MLflow tracking server. All MLflow REST APIs are represented as IAM actions under the `sagemaker-mlflow` service prefix. For more information, see [Set up IAM permissions for MLflow](mlflow-create-tracking-server-iam.md). **AWS CloudTrail** View logs in AWS CloudTrail to help you enable operational and risk auditing, governance, and compliance of your AWS account. For more information, see [AWS CloudTrail logs](#mlflow-create-tracking-server-cloudtrail). **Amazon EventBridge** Automate the model review and deployment lifecycle using MLflow events captured by Amazon EventBridge. For more information, see [Amazon EventBridge events](#mlflow-create-tracking-server-eventbridge). ## Supported AWS Regions Amazon SageMaker AI with MLflow is generally available in all AWS commercial [Regions](https://docs.aws.amazon.com/sagemaker/latest/dg/regions-quotas.html) where Amazon SageMaker Studio is available, except the China Regions. SageMaker AI with MLflow is available using only the AWS CLI in the Europe (Zurich) Region, Asia Pacific (Hyderabad) Region, Asia Pacific (Melbourne) Region, and Canada West (Calgary) Region. Tracking servers are launched in a single availability zone within their specified Region. ## How it works An MLflow Tracking Server has three main components: compute, backend metadata storage, and artifact storage. The compute that hosts the tracking server and the backend metadata storage are securely hosted in the SageMaker AI service account. The artifact storage lives in an Amazon S3 bucket in your own AWS account. ![\[A diagram showing the compute and metadata store for an MLflow Tracking Server.\]](http://docs.aws.amazon.com/sagemaker/latest/dg/images/mlflow/mlflow-diagram.png) A tracking server has an ARN. You can use this ARN to connect the MLflow SDK to your Tracking Server and start logging your training runs to MLflow. Read on for more information about the following key concepts: + [Backend metadata storage](#mlflow-create-tracking-server-backend-store) + [Artifact storage](#mlflow-create-tracking-server-artifact-store) + [MLflow Tracking Server sizes](#mlflow-create-tracking-server-sizes) + [Tracking server versions](#mlflow-create-tracking-server-versions) + [AWS CloudTrail logs](#mlflow-create-tracking-server-cloudtrail) + [Amazon EventBridge events](#mlflow-create-tracking-server-eventbridge) ### Backend metadata storage When you create an MLflow Tracking Server, a [backend store](https://mlflow.org/docs/latest/tracking/backend-stores.html), which persists various metadata for each [Run](https://mlflow.org/docs/latest/tracking.html#runs), such as run ID, start and end times, parameters, and metrics, is automatically configured within the SageMaker AI service account and fully managed for you. ### Artifact storage To provide MLflow with persistent storage for metadata for each run, such as model weights, images, model files, and data files for your experiment runs, you must create an artifact store using Amazon S3. The artifact store must be set up within your AWS account and you must explicitly give MLflow access to Amazon S3 in order to access your artifact store. For more information, see [Artifact Stores](https://mlflow.org/docs/latest/tracking.html#artifact-stores) in the MLflow documentation. **Note** SageMaker AI MLflow has a 200 MB download size limit. ### MLflow app versions The following MLflow versions are available to use with SageMaker AI MLflow Apps: | MLflow version | Python version | | --- | --- | | [MLflow 3.4](https://mlflow.org/releases/3.4.0) (latest version) | [Python 3.9](https://www.python.org/downloads/release/python-390/) or later | The latest version of the MLflow App has the latest features, security patches, and bug fixes. When you create a new MLflow App it will be automatically updated to the latest supported version. For more information about creating an MLflow App, see [MLflow App Setup](mlflow-app-setup.md). MLflow Apps use semantic versioning. Versions are in the following format: `major-version.minor-version.patch-version`. ### MLflow Tracking Server sizes You can optionally specify the size of your tracking server in the Studio UI or with the AWS CLI parameter `--tracking-server-size`. You can 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. We recommend using a small tracking server for teams of up to 25 users, a medium tracking server for teams of up to 50 users, and a large tracking server for teams of up to 100 users. We assume that all users will make concurrent requests to your MLflow Tracking Server to make these recommendations. You should select the tracking server size based on your expected usage pattern and the TPS (Transactions Per Second) supported by each tracking server. **Note** The nature of your workload and the type of requests that you make to the tracking server dictate the TPS you see. | Tracking server size | Sustained TPS | Burst TPS | | --- | --- | --- | | Small | Up to 25 | Up to 50 | | Medium | Up to 50 | Up to 100 | | Large | Up to 100 | Up to 200 | ### Tracking server versions The following MLflow versions are available to use with SageMaker AI: | MLflow version | Python version | | --- | --- | | [MLflow 3.0](https://mlflow.org/releases/3) (latest version) | [Python 3.9](https://www.python.org/downloads/release/python-390/) or later | | [MLflow 2.16](https://mlflow.org/releases/2.16.0) | [Python 3.8](https://www.python.org/downloads/release/python-380/) or later | | [MLflow 2.13](https://mlflow.org/releases/2.13.0) | [Python 3.8](https://www.python.org/downloads/release/python-380/) or later | The latest version of the tracking server has the latest features, security patches, and bug fixes. When you create a new tracking server, we recommend using the latest version. For more information about creating a tracking server, see [MLflow Tracking Servers](mlflow-create-tracking-server.md). MLflow tracking servers use semantic versioning. Versions are in the following format: `major-version.minor-version.patch-version`. The latest features, such as new UI elements and API functionality, are in the minor-version. ### AWS CloudTrail logs AWS CloudTrail automatically logs activity related to your MLflow Tracking Server. The following control plane API calls are logged in CloudTrail: + CreateMlflowTrackingServer + DescribeMlflowTrackingServer + UpdateMlflowTrackingServer + DeleteMlflowTrackingServer + ListMlflowTrackingServers + CreatePresignedMlflowTrackingServer + StartMlflowTrackingServer + StopMlflowTrackingServer AWS CloudTrail also automatically logs activity related to your MLflow data plane. The following data plane API calls are logged in CloudTrail. For event names, add the prefix `Mlflow` (for example, `MlflowCreateExperiment`). + CreateExperiment + CreateModelVersion + CreateRegisteredModel + CreateRun + DeleteExperiment + DeleteModelVersion + DeleteModelVersionTag + DeleteRegisteredModel + DeleteRegisteredModelAlias + DeleteRegisteredModelTag + DeleteRun + DeleteTag + GetDownloadURIForModelVersionArtifacts + GetExperiment + GetExperimentByName + GetLatestModelVersions + GetMetricHistory + GetModelVersion + GetModelVersionByAlias + GetRegisteredModel + GetRun + ListArtifacts + LogBatch + LogInputs + LogMetric + LogModel + LogParam + RenameRegisteredModel + RestoreExperiment + RestoreRun + SearchExperiments + SearchModelVersions + SearchRegisteredModels + SearchRuns + SetExperimentTag + SetModelVersionTag + SetRegisteredModelAlias + SetRegisteredModelTag + SetTag + TransitionModelVersionStage + UpdateExperiment + UpdateModelVersion + UpdateRegisteredModel + UpdateRun + FinalizeLoggedModel + GetLoggedModel + DeleteLoggedModel + SearchLoggedModels + SetLoggedModelTags + DeleteLoggedModelTag + ListLoggedModelArtifacts + LogLoggedModelParams + LogOutputs For more information about CloudTrail, see the *[AWS CloudTrail User Guide](https://docs.aws.amazon.com/awscloudtrail/latest/userguide/cloudtrail-user-guide.html)*. ### Amazon EventBridge events Use EventBridge to route events from using MLflow with SageMaker AI to consumer applications across your organization. The following events are emitted to EventBridge: + "SageMaker Tracking Server Creating" + "SageMaker Tracking Server Created“ + "SageMaker Tracking Server Create Failed" + "SageMaker Tracking Server Updating" + "SageMaker Tracking Server Updated" + "SageMaker Tracking Server Update Failed" + "SageMaker Tracking Server Deleting" + "SageMaker Tracking Server Deleted" + "SageMaker Tracking Server Delete Failed" + "SageMaker Tracking Server Starting" + "SageMaker Tracking Server Started" + "SageMaker Tracking Server Start Failed" + "SageMaker Tracking Server Stopping" + "SageMaker Tracking Server Stopped" + "SageMaker Tracking Server Stop Failed" + "SageMaker Tracking Server Maintenance In Progress" + "SageMaker Tracking Server Maintenance Complete" + "SageMaker Tracking Server Maintenance Failed" + "SageMaker MLFlow Tracking Server Creating Run" + "SageMaker MLFlow Tracking Server Creating RegisteredModel" + "SageMaker MLFlow Tracking Server Creating ModelVersion" + "SageMaker MLFlow Tracking Server Transitioning ModelVersion Stage" + "SageMaker MLFlow Tracking Server Setting Registered Model Alias" For more information about EventBridge, see the *[Amazon EventBridge User Guide](https://docs.aws.amazon.com/eventbridge/latest/userguide/eb-what-is.html)*. **Topics** + [ ## MLflow integrations ](#mlflow-integrations) + [ ## Supported AWS Regions ](#mlflow-regions) + [ ## How it works ](#mlflow-create-tracking-server-how-it-works) + [ # MLflow App Setup ](mlflow-app-setup.md) + [ # MLflow Tracking Servers ](mlflow-create-tracking-server.md) + [ # Launch the MLflow UI using a presigned URL ](mlflow-launch-ui.md) + [ # Integrate MLflow with your environment ](mlflow-track-experiments.md) + [ # MLflow tutorials using example Jupyter notebooks ](mlflow-tutorials.md) + [ # Troubleshoot common setup issues ](mlflow-troubleshooting.md) + [ # Clean up MLflow resources ](mlflow-cleanup.md) + [ # Amazon SageMaker Experiments in Studio Classic ](experiments.md) # MLflow App Setup An [MLflow App](https://mlflow.org/docs/latest/tracking.html#mlflow-tracking-server-optional) is a stand-alone HTTP server that serves multiple REST API endpoints for tracking runs and experiments. An MLflow App is required to begin tracking your machine learning (ML) experiments with SageMaker AI and MLflow. You can create an MLflow App through the Studio UI, or through the AWS CLI for more granular security customization. You must have the correct IAM permissions configured to create an MLflow App. MLflow Apps are the latest managed MLflow offering on SageMaker and should be preferred over existing MLflow Tracking Servers. MLflow Apps offer additional features such as faster startup time, cross-account sharing, integrations with other SageMaker features, and other features beyond the existing MLflow Tracking Servers. **Topics** + [ # MLflow App Setup Prequisites ](mlflow-app-setup-prerequisites.md) + [ # Create MLflow App ](mlflow-app-setup-create-app.md) # MLflow App Setup Prequisites # Set up IAM permissions for MLflow Apps You must configure the necessary IAM service roles to get started with MLflow Apps in Amazon SageMaker AI. If you create a new Amazon SageMaker AI domain to access your experiments in Studio, you can configure the necessary IAM permissions during domain setup. For more information, see [Set up MLflow IAM permissions when creating a new domain](mlflow-create-tracking-server-iam.md#mlflow-create-tracking-server-iam-role-manager). To set up permissions using the IAM console, see [Create necessary IAM service roles in the IAM console](mlflow-create-tracking-server-iam.md#mlflow-create-tracking-server-iam-service-roles). You must configure authorization controls for `sagemaker-mlflow` actions. You can optionally define more granular authorization controls to govern action-specific MLflow permissions. For more information, see [Create action-specific authorization controls](#mlflow-create-app-update-iam-actions). ## Set up MLflow IAM permissions when creating a new domain When setting up a new Amazon SageMaker AI domain for your organization, you can configure IAM permissions for your domain service role through the **Users and ML Activities** settings. 1. Set up a new domain using the SageMaker AI console. On the **Set up SageMaker AI domain** page, choose **Set up for organizations**. For more information, see [Custom setup using the console](onboard-custom.md#onboard-custom-instructions-console). 1. When setting up **Users and ML Activities**, choose from the following ML activities for MLflow: **Use MLflow**, **Manage MLflow Apps**, and **Access required to AWS Services for MLflow**. For more information about these activities, see the explanations that follow this procedure. 1. Complete the setup and creation of your new domain. The following MLflow ML activities are available in Amazon SageMaker Role Manager: + **Use MLflow**: This ML activity grants the domain service role permission to call MLflow REST APIs in order to manage experiments, runs, and models in MLflow. + **Manage MLflow Apps**: This ML activity grants the domain service role permission to create, update, and delete MLflow Apps. + **Access required to AWS services for MLflow Apps**: This ML activity provides the domain service role permissions needed to access Amazon S3 and the SageMaker AI Model Registry. This allows you to use the domain service role as the tracking server service role. For more information about ML activities in Role Manager, see [ML activity reference](role-manager-ml-activities.md). ## Create necessary IAM service roles in the IAM console If you did not create or update your domain service role, you must instead create the following service roles in the IAM console in order to create and use an MLflow Apps: + An MLflow App IAM service role that the App can use to access SageMaker AI resources + A SageMaker AI IAM service role that SageMaker AI can use to create and manage MLflow resources ### IAM policies for the MLflow App IAM service role The MLflow App IAM service role is used by the app to access the resources it needs such as Amazon S3 and the SageMaker Model Registry. When creating the app IAM service role, use the following IAM trust policy: ------ #### [ JSON ] **** ``` { "Version":"2012-10-17", "Statement": [ { "Effect": "Allow", "Principal": { "Service": [ "sagemaker.amazonaws.com" ] }, "Action": "sts:AssumeRole" } ] } ``` ------ In the IAM console, add the following permissions policy to your app service role: ------ #### [ JSON ] **** ``` { "Version":"2012-10-17", "Statement": [ { "Effect": "Allow", "Action": [ "s3:Get*", "s3:Put*", "s3:List*", "sagemaker:AddTags", "sagemaker:CreateModelPackageGroup", "sagemaker:CreateModelPackage", "sagemaker:UpdateModelPackage", "sagemaker:DescribeModelPackageGroup" ], "Resource": "*" } ] } ``` ------ ### IAM policy for the SageMaker AI IAM service role The SageMaker AI service role is used by the client accessing the MLflow App and needs permissions to call MLflow REST APIs. The SageMaker AI service role also needs SageMaker API permissions to create, view update, and delete apps. You can create a new role or update an existing role. The SageMaker AI service role needs the following policy: ------ #### [ JSON ] **** ``` { "Version":"2012-10-17", "Statement": [ { "Effect": "Allow", "Action": [ "sagemaker-mlflow:*", "sagemaker:CreateMlflowTrackingServer", "sagemaker:ListMlflowTrackingServers", "sagemaker:UpdateMlflowTrackingServer", "sagemaker:DeleteMlflowTrackingServer", "sagemaker:StartMlflowTrackingServer", "sagemaker:StopMlflowTrackingServer", "sagemaker:CreatePresignedMlflowTrackingServerUrl" ], "Resource": "*" } ] } ``` ------ ## Create action-specific authorization controls You must set up authorization controls for `sagemaker-mlflow`, and can optionally configure action-specific authorization controls to govern more granular MLflow permissions that your users have on an MLflow Apps. **Note** The following steps assume that you have an ARN for an MLflow Apps already available. ### Data Plane IAM actions supported for MLflow Apps The following SageMaker AI MLflow actions are supported for authorization access control: + sagemaker:CallMlflowAppApi # Create MLflow App # Create an app using the AWS CLI You can create an app using the AWS CLI for more granular security customization. ## Prerequisites To create an app using the AWS CLI, you must have the following: + **Access to a terminal. **This can include local IDEs, an Amazon EC2 instance, or AWS CloudShell. + **Access to a development environment.** This can include local IDEs or a Jupyter notebook environment within Studio or Studio Classic. + **A configured AWS CLI installation**. For more information, see [Configure the AWS CLI](https://docs.aws.amazon.com/cli/latest/userguide/cli-chap-configure.html). + **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 App 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 App via the MLflow SDK permission to call MLflow APIs. For more information, see [Modifying a role permissions policy (console) ](https://docs.aws.amazon.com/IAM/latest/UserGuide/roles-managingrole-editing-console.html#roles-modify_permissions-policy). 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 AI 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 AWS CLI model Follow these command line steps within a terminal to set up the AWS CLI for Amazon SageMaker AI with MLflow. 1. Install an updated version of the AWS CLI. For more information, see [Install or update to the latest version of the AWS CLI](https://docs.aws.amazon.com/cli/latest/userguide/getting-started-install.html) in the *AWS CLI User Guide*. 1. Verify that the AWS CLI is installed using the following command: ``` aws sagemaker help ``` Press `q` to exit the prompt. For troubleshooting help, see [Troubleshoot common setup issues](mlflow-troubleshooting.md). ## Set up MLflow infrastructure The following section shows you how to set up an MLflow App along with the Amazon S3 bucket and IAM role needed for the app. ### Create an S3 bucket Within your terminal, use the following commands to create a general purpose Amazon S3 bucket: **Important** When you provide the Amazon S3 URI for your artifact store, ensure the Amazon S3 bucket is in the same AWS Region as your MLflow App. **Cross-region artifact storage is not supported**. ``` 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](https://docs.aws.amazon.com/IAM/latest/UserGuide/id_roles_terms-and-concepts.html) in the *AWS Identity and Access Management User Guide*. 1. Within your terminal, use the following command to create a file called `mlflow-trust-policy.json`. ``` cat < /tmp/mlflow-trust-policy.json { "Version": "2012-10-17", "Statement": [ { "Effect": "Allow", "Principal": { "Service": [ "sagemaker.amazonaws.com" ] }, "Action": "sts:AssumeRole" } ] } EOF ``` 1. Within your terminal, use the following command to create a file called `custom-policy.json`. ``` cat < /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 ``` 1. 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 app'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 App Within your terminal, use the `create-mlflow-app` API to create an app in the AWS Region of your choice. This step normally takes approximately 2-3 minutes. The following command creates a new app with automatic model registration enabled. To deactivate automatic model registration, specify `--no-automatic-model-registration`. After creating your app, you can launch the MLflow UI. For more information, see [Launch the MLflow UI using a presigned URL](mlflow-launch-ui.md). **Note** It may take up to 2-3 minutes to complete app creation. If the app takes over 3 minutes to create, check that you have the necessary IAM permissions. When you successfully create an app, it automatically starts. By default, the app that is created is the latest version and will be automatically updated. ``` app_name=app-name region=valid-region version=valid-version aws sagemaker create-mlflow-app \ --name $app_name \ --artifact-store-uri s3://$bucket_name \ --role-arn $role_arn \ --automatic-model-registration \ --region $region ``` The output should be similar to the following: ``` { "AppArn": "arn:aws:sagemaker:region:123456789012:mlflow-app/app-name" } ``` **Important** **Take note of the app ARN for later use.** You will also need the `$bucket_name` for clean up steps. # MLflow Tracking Servers An [MLflow Tracking Server](https://mlflow.org/docs/latest/tracking.html#mlflow-tracking-server-optional) is a stand-alone HTTP server that serves multiple REST API endpoints for tracking runs and experiments. A tracking server is required to begin tracking your machine learning (ML) experiments with SageMaker AI and MLflow. You can create a tracking server through the Studio UI, or through the AWS CLI for more granular security customization. You must have the correct IAM permissions configured to create an MLflow Tracking Server. **Topics** + [ # Set up IAM permissions for MLflow ](mlflow-create-tracking-server-iam.md) + [ # Create a tracking server using Studio ](mlflow-create-tracking-server-studio.md) + [ # Create a tracking server using the AWS CLI ](mlflow-create-tracking-server-cli.md) # Set up IAM permissions for MLflow You must configure the necessary IAM service roles to get started with MLflow in Amazon SageMaker AI. If you create a new Amazon SageMaker AI domain to access your experiments in Studio, you can configure the necessary IAM permissions during domain setup. For more information, see [Set up MLflow IAM permissions when creating a new domain](#mlflow-create-tracking-server-iam-role-manager). To set up permissions using the IAM console, see [Create necessary IAM service roles in the IAM console](#mlflow-create-tracking-server-iam-service-roles). You must configure authorization controls for `sagemaker-mlflow` actions. You can optionally define more granular authorization controls to govern action-specific MLflow permissions. For more information, see [Create action-specific authorization controls](#mlflow-create-tracking-server-update-iam-actions). ## Set up MLflow IAM permissions when creating a new domain When setting up a new Amazon SageMaker AI domain for your organization, you can configure IAM permissions for your domain service role through the **Users and ML Activities** settings. **To configure IAM permissions for using MLflow with SageMaker AI when setting up a new domain** 1. Set up a new domain using the SageMaker AI console. On the **Set up SageMaker AI domain** page, choose **Set up for organizations**. For more information, see [Custom setup using the console](onboard-custom.md#onboard-custom-instructions-console). 1. When setting up **Users and ML Activities**, choose from the following ML activities for MLflow: **Use MLflow**, **Manage MLflow Tracking Servers**, and **Access required to AWS Services for MLflow**. For more information about these activities, see the explanations that follow this procedure. 1. Complete the setup and creation of your new domain. The following MLflow ML activities are available in Amazon SageMaker Role Manager: + **Use MLflow**: This ML activity grants the domain service role permission to call MLflow REST APIs in order to manage experiments, runs, and models in MLflow. + **Manage MLflow Tracking Servers**: This ML activity grants the domain service role permission to create, update, start, stop, and delete tracking servers. + **Access required to AWS Services for MLflow**: This ML activity provides the domain service role permissions needed to access Amazon S3 and the SageMaker AI Model Registry. This allows you to use the domain service role as the tracking server service role. For more information about ML activities in Role Manager, see [ML activity reference](role-manager-ml-activities.md). ## Create necessary IAM service roles in the IAM console If you did not create or update your domain service role, you must instead create the following service roles in the IAM console in order to create and use an MLflow Tracking Server: + A tracking server IAM service role that the tracking server can use to access SageMaker AI resources + A SageMaker AI IAM service role that SageMaker AI can use to create and manage MLflow resources ### IAM policies for the tracking server IAM service role The tracking server IAM service role is used by the tracking server to access the resources it needs such as Amazon S3 and the SageMaker Model Registry. When creating the tracking server IAM service role, use the following IAM trust policy: ------ #### [ JSON ] **** ``` { "Version":"2012-10-17", "Statement": [ { "Effect": "Allow", "Principal": { "Service": [ "sagemaker.amazonaws.com" ] }, "Action": "sts:AssumeRole" } ] } ``` ------ In the IAM console, add the following permissions policy to your tracking server service role: ------ #### [ JSON ] **** ``` { "Version":"2012-10-17", "Statement": [ { "Effect": "Allow", "Action": [ "s3:Get*", "s3:Put*", "s3:List*", "sagemaker:AddTags", "sagemaker:CreateModelPackageGroup", "sagemaker:CreateModelPackage", "sagemaker:UpdateModelPackage", "sagemaker:DescribeModelPackageGroup" ], "Resource": "*" } ] } ``` ------ ### IAM policy for the SageMaker AI IAM service role The SageMaker AI service role is used by the client accessing the MLflow Tracking Server and needs permissions to call MLflow REST APIs. The SageMaker AI service role also needs SageMaker API permissions to create, view update, start, stop, and delete tracking servers. You can create a new role or update an existing role. The SageMaker AI service role needs the following policy: ------ #### [ JSON ] **** ``` { "Version":"2012-10-17", "Statement": [ { "Effect": "Allow", "Action": [ "sagemaker-mlflow:*", "sagemaker:CreateMlflowTrackingServer", "sagemaker:ListMlflowTrackingServers", "sagemaker:UpdateMlflowTrackingServer", "sagemaker:DeleteMlflowTrackingServer", "sagemaker:StartMlflowTrackingServer", "sagemaker:StopMlflowTrackingServer", "sagemaker:CreatePresignedMlflowTrackingServerUrl" ], "Resource": "*" } ] } ``` ------ ## Create action-specific authorization controls You must set up authorization controls for `sagemaker-mlflow`, and can optionally configure action-specific authorization controls to govern more granular MLflow permissions that your users have on an MLflow Tracking Server. **Note** The following steps assume that you have an ARN for an MLflow Tracking Server already available. To learn how to create a tracking server, see [Create a tracking server using Studio](mlflow-create-tracking-server-studio.md) or [Create a tracking server using the AWS CLI](mlflow-create-tracking-server-cli.md). The following command creates a file called `mlflow-policy.json` that provides your tracking server with IAM permissions for all available SageMaker AI MLflow actions. You can optionally limit the permissions a user has by choosing the specific actions you want that user to perform. For a list of available actions, see [IAM actions supported for MLflow](#mlflow-create-tracking-server-iam-actions). ``` # Replace "Resource":"*" with "Resource":"TrackingServerArn" # Replace "sagemaker-mlflow:*" with specific actions printf '{ "Version": "2012-10-17", "Statement": [ { "Effect": "Allow", "Action": "sagemaker-mlflow:*", "Resource": "*" } ] }' > mlflow-policy.json ``` Use the `mlflow-policy.json` file to create an IAM policy using the AWS CLI. ``` aws iam create-policy \ --policy-name MLflowPolicy \ --policy-document file://mlflow-policy.json ``` Retrieve your account ID and attach the policy to your IAM role. ``` # Get your account ID aws sts get-caller-identity # Attach the IAM policy using your exported role and account ID aws iam attach-role-policy \ --role-name $role_name \ --policy-arn arn:aws:iam::123456789012:policy/MLflowPolicy ``` ### IAM actions supported for MLflow The following SageMaker AI MLflow actions are supported for authorization access control: + sagemaker-mlflow:AccessUI + sagemaker-mlflow:CreateExperiment + sagemaker-mlflow:SearchExperiments + sagemaker-mlflow:GetExperiment + sagemaker-mlflow:GetExperimentByName + sagemaker-mlflow:DeleteExperiment + sagemaker-mlflow:RestoreExperiment + sagemaker-mlflow:UpdateExperiment + sagemaker-mlflow:CreateRun + sagemaker-mlflow:DeleteRun + sagemaker-mlflow:RestoreRun + sagemaker-mlflow:GetRun + sagemaker-mlflow:LogMetric + sagemaker-mlflow:LogBatch + sagemaker-mlflow:LogModel + sagemaker-mlflow:LogInputs + sagemaker-mlflow:SetExperimentTag + sagemaker-mlflow:SetTag + sagemaker-mlflow:DeleteTag + sagemaker-mlflow:LogParam + sagemaker-mlflow:GetMetricHistory + sagemaker-mlflow:SearchRuns + sagemaker-mlflow:ListArtifacts + sagemaker-mlflow:UpdateRun + sagemaker-mlflow:CreateRegisteredModel + sagemaker-mlflow:GetRegisteredModel + sagemaker-mlflow:RenameRegisteredModel + sagemaker-mlflow:UpdateRegisteredModel + sagemaker-mlflow:DeleteRegisteredModel + sagemaker-mlflow:GetLatestModelVersions + sagemaker-mlflow:CreateModelVersion + sagemaker-mlflow:GetModelVersion + sagemaker-mlflow:UpdateModelVersion + sagemaker-mlflow:DeleteModelVersion + sagemaker-mlflow:SearchModelVersions + sagemaker-mlflow:GetDownloadURIForModelVersionArtifacts + sagemaker-mlflow:TransitionModelVersionStage + sagemaker-mlflow:SearchRegisteredModels + sagemaker-mlflow:SetRegisteredModelTag + sagemaker-mlflow:DeleteRegisteredModelTag + sagemaker-mlflow:DeleteModelVersionTag + sagemaker-mlflow:DeleteRegisteredModelAlias + sagemaker-mlflow:SetRegisteredModelAlias + sagemaker-mlflow:GetModelVersionByAlias + sagemaker-mlflow:FinalizeLoggedModel + sagemaker-mlflow:GetLoggedModel + sagemaker-mlflow:DeleteLoggedModel + sagemaker-mlflow:SearchLoggedModels + sagemaker-mlflow:SetLoggedModelTags + sagemaker-mlflow:DeleteLoggedModelTag + sagemaker-mlflow:ListLoggedModelArtifacts + sagemaker-mlflow:LogLoggedModelParams + sagemaker-mlflow:LogOutputs # Create a tracking server using Studio You can create a tracking server from the SageMaker Studio MLflow UI. If you created your SageMaker Studio domain following the **Set up for organizations** workflow, the service role for your SageMaker Studio domain has sufficient permissions to serve as the SageMaker AI IAM service roles and the tracking server IAM service role. Create a tracking server from the SageMaker Studio MLflow UI with the following steps: 1. Navigate to Studio from the SageMaker AI console. Be sure that you are using the new Studio experience and have updated from Studio Classic. For more information, see [Migration from Amazon SageMaker Studio Classic](studio-updated-migrate.md). 1. Choose **MLflow** in the **Applications** pane of the Studio UI. 1. **(Optional)** If have not already created a Tracking Server or if you need to create a new one, you can choose **Create**. Then provide a unique tracking server name and S3 URI for artifact storage and create a tracking server. You can optionally choose **Configure** for more granular tracking server customization. 1. Choose **Create** in the **MLflow Tracking Servers** pane. The Studio domain IAM service role is used for the tracking server IAM service role. 1. Provide a unique name for your tracking server and an Amazon S3 URI for your tracking server artifact store. Your tracking server and the Amazon S3 bucket must be in the **same AWS Region**. **Important** When you provide the Amazon S3 URI for your artifact store, ensure the Amazon S3 bucket is in the same AWS Region as your tracking server. **Cross-region artifact storage is not supported**. 1. **(Optional)** Choose **Configure** to change default settings such as tracking server size, tags, and the IAM service role. 1. Choose **Create**. **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](mlflow-create-tracking-server-iam.md). When you successfully create a tracking server, it automatically starts. 1. After creating your tracking server, you can launch the MLflow UI. For more information, see [Launch the MLflow UI using a presigned URL](mlflow-launch-ui.md). ![\[The Create MLflow Tracking Server prompt in the Studio UI.\]](http://docs.aws.amazon.com/sagemaker/latest/dg/images/mlflow/mlflow-studio-create.png) # Create a tracking server using the AWS CLI You can create a tracking server using the AWS CLI for more granular security customization. ## Prerequisites To create a tracking server using the AWS CLI, you must have the following: + **Access to a terminal. **This can include local IDEs, an Amazon EC2 instance, or AWS CloudShell. + **Access to a development environment.** This can include local IDEs or a Jupyter notebook environment within Studio or Studio Classic. + **A configured AWS CLI installation**. For more information, see [Configure the AWS CLI](https://docs.aws.amazon.com/cli/latest/userguide/cli-chap-configure.html). + **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) ](https://docs.aws.amazon.com/IAM/latest/UserGuide/roles-managingrole-editing-console.html#roles-modify_permissions-policy). 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 AI 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 AWS CLI model Follow these command line steps within a terminal to set up the AWS CLI for Amazon SageMaker AI with MLflow. 1. Install an updated version of the AWS CLI. For more information, see [Install or update to the latest version of the AWS CLI](https://docs.aws.amazon.com/cli/latest/userguide/getting-started-install.html) in the *AWS CLI User Guide*. 1. Verify that the AWS CLI is installed using the following command: ``` aws sagemaker help ``` Press `q` to exit the prompt. For troubleshooting help, see [Troubleshoot common setup issues](mlflow-troubleshooting.md). ## 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: **Important** When you provide the Amazon S3 URI for your artifact store, ensure the Amazon S3 bucket is in the same AWS Region as your tracking server. **Cross-region artifact storage is not supported**. ``` 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](https://docs.aws.amazon.com/IAM/latest/UserGuide/id_roles_terms-and-concepts.html) in the *AWS Identity and Access Management User Guide*. 1. Within your terminal, use the following command to create a file called `mlflow-trust-policy.json`. ``` cat < /tmp/mlflow-trust-policy.json { "Version": "2012-10-17", "Statement": [ { "Effect": "Allow", "Principal": { "Service": [ "sagemaker.amazonaws.com" ] }, "Action": "sts:AssumeRole" } ] } EOF ``` 1. Within your terminal, use the following command to create a file called `custom-policy.json`. ``` cat < /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 ``` 1. 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 AWS 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](mlflow.md#mlflow-create-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](mlflow-launch-ui.md). **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](mlflow-create-tracking-server-iam.md). When you successfully create a tracking server, it automatically starts. When you create a tracking server, we recommend specifying the latest version. For information about the available versions, see [Tracking server versions](mlflow.md#mlflow-create-tracking-server-versions). By default, the tracking server that's created is the latest version. However, we recommend always specifying the latest version explicitly because the underlying MLflow APIs can change. ``` ts_name=tracking-server-name region=valid-region version=valid-version 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 \ --mlflow-version $version ``` 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. # Launch the MLflow UI using a presigned URL You can access the MLflow UI to view your experiments using a presigned URL. You can launch the MLflow UI either through Studio or using the AWS CLI in a terminal of your choice. ## Launch the MLflow UI using Studio After creating your tracking server, you can launch the MLflow UI directly from Studio. 1. Navigate to Studio from the SageMaker AI console. Be sure that you are using the new Studio experience and have updated from Studio Classic. For more information, see [Migration from Amazon SageMaker Studio Classic](studio-updated-migrate.md). 1. Choose **MLflow** in the **Applications** pane of the Studio UI. 1. **(Optional)** If have not already created a tracking server or if you need to create a new one, you can choose **Create**. Then provide a unique tracking server name and S3 URI for artifact storage and create a tracking server. You can optionally choose **Configure** for more granular tracking server customization. 1. Find the tracking server of your choice in the **MLflow Tracking Servers** pane. If the tracking server is **Off**, start the tracking server. 1. Choose the vertical menu icon in the right corner of the tracking server pane. Then, choose **Open MLflow**. This launches a presigned URL in a new tab in your current browser. ![\[The option to open a presigned URL through the MLflow Tracking Servers pane in the Studio UI.\]](http://docs.aws.amazon.com/sagemaker/latest/dg/images/mlflow/mlflow-studio-ui.png) ## Launch the MLflow UI using the AWS CLI You can access the MLflow UI to view your experiments using a presigned URL. Within your terminal, use the `create-presigned-mlflow-tracking-server-url` API to generate a presigned URL. ``` aws sagemaker create-presigned-mlflow-tracking-server-url \ --tracking-server-name $ts_name \ --session-expiration-duration-in-seconds 1800 \ --expires-in-seconds 300 \ --region $region ``` The output should look similar to the following: ``` { "AuthorizedUrl": "https://unique-key.us-west-2.experiments.sagemaker.aws.a2z.com/auth?authToken=example_token" } ``` Copy the entire presigned URL into the browser of your choice. You can use a new tab or a new private window. Press `q` to exit the prompt. The `--session-expiration-duration-in-seconds` parameter determines the length of time that your MLflow UI session remains valid. The session duration time is the amount of time that the MLflow UI can be loaded in the browser before a new presigned URL must be created. The minimum session duration is 30 minutes (1800 seconds) and the maximum session duration is 12 hours (43200 seconds). The default session duration is 12 hours if no other duration is specified. The `--expires-in-seconds parameter` determines the length of time that your presigned URL remains valid. The minimum URL expiration length is 5 seconds and the maximum URL expiration length is 5 minutes (300 seconds). The default URL expiration length is 300 seconds. The presigned URL can be used only once. The window should look similar to the following. ![\[The MLflow UI that launches after creating and using a presigned URL\]](http://docs.aws.amazon.com/sagemaker/latest/dg/images/mlflow/mlflow-ui.png) # Integrate MLflow with your environment The following page describes how to get started with the MLflow SDK and the AWS MLflow plugin within your development environment. This can include local IDEs or a Jupyter Notebook environment within Studio or Studio Classic. Amazon SageMaker AI uses an MLflow plugin to customize the behavior of the MLflow Python client and integrate AWS tooling. The AWS MLflow plugin authenticates API calls made with MLflow using [AWS Signature Version 4](https://docs.aws.amazon.com/AmazonS3/latest/API/sig-v4-authenticating-requests.html). The AWS MLflow plugin allows you to connect to your MLflow tracking server using the tracking server ARN. For more information about plugins, see [AWS MLflow plugin](https://pypi.org/project/sagemaker-mlflow/) and [MLflow plugins](https://mlflow.org/docs/latest/plugins.html). **Important** Your user IAM permissions within your development environment must have access to any relevant MLflow API actions to successfully run provided examples. For more information, see [Set up IAM permissions for MLflow](mlflow-create-tracking-server-iam.md). For more information about using the MLflow SDK, see [Python API](https://mlflow.org/docs/2.13.2/python_api/index.html) in the MLflow documentation. ## Install MLflow and the AWS MLflow plugin Within your development environment, install both MLflow and the AWS MLflow plugin. ``` pip install sagemaker-mlflow ``` To ensure compatibility between your MLflow client and tracking server, use the corresponding MLflow version based on your tracking server version: + For tracking server 2.13.x, use `mlflow==2.13.2` + For tracking server 2.16.x, use `mlflow==2.16.2` + For tracking server 3.0.x, use `mlflow==3.0.0` To see which versions of MLflow are available to use with SageMaker AI, see [Tracking server versions](mlflow.md#mlflow-create-tracking-server-versions). ## Connect to your MLflow Tracking Server Use `[mlflow.set\$1tracking\$1uri](https://mlflow.org/docs/2.13.2/python_api/mlflow.html#mlflow.set_tracking_uri)` to connect to a your tracking server from your development environment using its ARN: ``` import mlflow arn = "YOUR-TRACKING-SERVER-ARN" mlflow.set_tracking_uri(arn) ``` # Log metrics, parameters, and MLflow models during training After connecting to your MLflow Tracking Server, you can use the MLflow SDK to log metrics, parameters, and MLflow models. ## Log training metrics Use `mlflow.log_metric` within an MLflow training run to track metrics. For more information about logging metrics using MLflow, see `[mlflow.log\$1metric](https://mlflow.org/docs/2.13.2/python_api/mlflow.html#mlflow.log_metric)`. ``` with mlflow.start_run(): mlflow.log_metric("foo", 1) print(mlflow.search_runs()) ``` This script should create an experiment run and print out an output similar to the following: ``` run_id experiment_id status artifact_uri ... tags.mlflow.source.name tags.mlflow.user tags.mlflow.source.type tags.mlflow.runName 0 607eb5c558c148dea176d8929bd44869 0 FINISHED s3://dddd/0/607eb5c558c148dea176d8929bd44869/a... ... file.py user-id LOCAL experiment-code-name ``` Within the MLflow UI, this example should look similar to the following: ![\[An experiment shown in the top-level MLflow Experiments menu.\]](http://docs.aws.amazon.com/sagemaker/latest/dg/images/mlflow/mlflow-ui-experiments.png) Choose **Run Name** to see more run details. ![\[An experiment parameter shown on an experiment run page in the MLflow UI.\]](http://docs.aws.amazon.com/sagemaker/latest/dg/images/mlflow/mlflow-ui-foo.png) ## Log parameters and models **Note** The following example requires your environment to have `s3:PutObject` permissions. This permission should be associated with the IAM Role that the MLflow SDK user assumes when they log into or federate into their AWS account. For more information, see [User and role policy examples](https://docs.aws.amazon.com/AmazonS3/latest/userguide/example-policies-s3.html). The following example takes you through a basic model training workflow using SKLearn and shows you how to track that model in an MLflow experiment run. This example logs parameters, metrics, and model artifacts. ``` import mlflow from mlflow.models import infer_signature import pandas as pd from sklearn import datasets from sklearn.model_selection import train_test_split from sklearn.linear_model import LogisticRegression from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score # This is the ARN of the MLflow Tracking Server you created mlflow.set_tracking_uri(your-tracking-server-arn) mlflow.set_experiment("some-experiment") # Load the Iris dataset X, y = datasets.load_iris(return_X_y=True) # Split the data into training and test sets X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42) # Define the model hyperparameters params = {"solver": "lbfgs", "max_iter": 1000, "multi_class": "auto", "random_state": 8888} # Train the model lr = LogisticRegression(**params) lr.fit(X_train, y_train) # Predict on the test set y_pred = lr.predict(X_test) # Calculate accuracy as a target loss metric accuracy = accuracy_score(y_test, y_pred) # Start an MLflow run and log parameters, metrics, and model artifacts with mlflow.start_run(): # Log the hyperparameters mlflow.log_params(params) # Log the loss metric mlflow.log_metric("accuracy", accuracy) # Set a tag that we can use to remind ourselves what this run was for mlflow.set_tag("Training Info", "Basic LR model for iris data") # Infer the model signature signature = infer_signature(X_train, lr.predict(X_train)) # Log the model model_info = mlflow.sklearn.log_model( sk_model=lr, name="iris_model", # Changed from artifact_path to name for MLflow 3.0 signature=signature, input_example=X_train, registered_model_name="tracking-quickstart", ) ``` Within the MLflow UI, choose the experiment name in the left navigation pane to explore all associated runs. Choose the **Run Name** to see more information about each run. For this example, your experiment run page for this run should look similar to the following. ![\[Tracked parameters for an experiment run in the MLflow UI.\]](http://docs.aws.amazon.com/sagemaker/latest/dg/images/mlflow/mlflow-ui-parameters.png) This example logs the logistic regression model. Within the MLflow UI, you should also see the logged model artifacts. ![\[Tracked model artifacts for an experiment run in the MLflow UI.\]](http://docs.aws.amazon.com/sagemaker/latest/dg/images/mlflow/mlflow-ui-model-artifacts.png) # Automatically register SageMaker AI models with SageMaker Model Registry You can log MLflow models and automatically register them with SageMaker Model Registry using either the Python SDK or directly through the MLflow UI. **Note** Do not use spaces in a model name. While MLflow supports model names with spaces, SageMaker AI Model Package does not. The auto-registration process fails if you use spaces in your model name. ## Register models using the SageMaker Python SDK Use `create_registered_model` within your MLflow client to automatically create a model package group in SageMaker AI that corresponds to an existing MLflow model of your choice. ``` import mlflow from mlflow import MlflowClient mlflow.set_tracking_uri(arn) client = MlflowClient() mlflow_model_name = 'AutoRegisteredModel' client.create_registered_model(mlflow_model_name, tags={"key1": "value1"}) ``` Use `mlflow.register_model()` to automatically register a model with the SageMaker Model Registry during model training. When registering the MLflow model, a corresponding model package group and model package version are created in SageMaker AI. ``` import mlflow.sklearn from mlflow.models import infer_signature from sklearn.datasets import make_regression from sklearn.ensemble import RandomForestRegressor mlflow.set_tracking_uri(arn) params = {"n_estimators": 3, "random_state": 42} X, y = make_regression(n_features=4, n_informative=2, random_state=0, shuffle=False) # Log MLflow entities with mlflow.start_run() as run: rfr = RandomForestRegressor(**params).fit(X, y) signature = infer_signature(X, rfr.predict(X)) mlflow.log_params(params) mlflow.sklearn.log_model(rfr, artifact_path="sklearn-model", signature=signature) model_uri = f"runs:/{run.info.run_id}/sklearn-model" mv = mlflow.register_model(model_uri, "RandomForestRegressionModel") print(f"Name: {mv.name}") print(f"Version: {mv.version}") ``` ## Register models using the MLflow UI You can alternatively register a model with the SageMaker Model Registry directly in the MLflow UI. Within the **Models** menu in the MLflow UI, choose **Create Model**. Any models newly created in this way are added to the SageMaker Model Registry. ![\[Model registry creation within the MLflow UI.\]](http://docs.aws.amazon.com/sagemaker/latest/dg/images/mlflow/mlflow-ui-register-model.png) After logging a model during experiment tracking, navigate to the run page in the MLflow UI. Choose the **Artifacts** pane and choose **Register model** in the upper right corner to register the model version in both MLflow and SageMaker Model Registry. ![\[Model registry creation within the MLflow UI.\]](http://docs.aws.amazon.com/sagemaker/latest/dg/images/mlflow/mlflow-ui-register-model-2.png) ## View registered models in Studio Within the SageMaker Studio landing page, choose **Models** on the left navigation pane to view your registered models. For more information on getting started with Studio, see [Launch Amazon SageMaker Studio](https://docs.aws.amazon.com/sagemaker/latest/dg/studio-updated-launch.html). ![\[MLflow models registered in SageMaker Model Registry in the Studio UI.\]](http://docs.aws.amazon.com/sagemaker/latest/dg/images/mlflow/mlflow-studio-model-registry.png) # Deploy MLflow models with `ModelBuilder` You can deploy MLflow models to a SageMaker AI endpoint using Amazon SageMaker AI Model Builder. For more information about Amazon SageMaker AI Model Builder, see [Create a model in Amazon SageMaker AI with ModelBuilder](https://docs.aws.amazon.com/sagemaker/latest/dg/how-it-works-modelbuilder-creation.html). `ModelBuilder` is a Python class that takes a framework model or a user-specified inference specification and converts it to a deployable model. For more details about the `ModelBuilder` class, see [ModelBuilder](https://sagemaker.readthedocs.io/en/stable/api/inference/model_builder.html#sagemaker.serve.builder.model_builder.ModelBuilder). To deploy your MLflow model using `ModelBuilder`, provide a path to your MLflow artifacts in the `model_metadata["MLFLOW_MODEL_PATH"]` attribute. Read on for more information about valid model path input formats: **Note** If you provide your model artifact path in the form of an MLflow run ID or MLflow model registry path, then you must also specify your tracking server ARN through the `model_metadata["MLFLOW_TRACKING_ARN"]` attribute. + [Model paths that require an ARN in the `model_metadata`](#mlflow-track-experiments-model-deployment-with-arn) + [Model paths that do not require an ARN in the `model_metadata`](#mlflow-track-experiments-model-deployment-without-arn) ## Model paths that require an ARN in the `model_metadata` The following model paths do require that you specify an ARN in the `model_metadata` for deployment: + MLflow [run ID](https://mlflow.org/docs/latest/python_api/mlflow.entities.html?highlight=mlflow%20info#mlflow.entities.RunInfo.run_id): `runs:/aloy-run-id/run-relative/path/to/model` + MLflow [model registry path](https://mlflow.org/docs/latest/model-registry.html#find-registered-models): `models:/model-name/model-version` ## Model paths that do not require an ARN in the `model_metadata` The following model paths do not require that you specify an ARN in the `model_metadata` for deployment: + Local model path: `/Users/me/path/to/local/model` + Amazon S3 model path: `s3://amzn-s3-demo-bucket/path/to/model` + Model package ARN: `arn:aws:sagemaker:region:account-id:mlflow-tracking-server/tracking-server-name` For more information on how MLflow model deployment works with Amazon SageMaker AI, see [Deploy MLflow Model to Amazon SageMaker AI](https://mlflow.org/docs/latest/deployment/deploy-model-to-sagemaker.html) in the MLflow documentation. If using an Amazon S3 path, you can find the path of your registered model with the following commands: ``` registered_model = client.get_registered_model(name='AutoRegisteredModel') source_path = registered_model.latest_versions[0].source ``` The following sample is an overview of how to deploy your MLflow model using `ModelBuilder` and an MLflow model registry path. Because this sample provides the model artifact path in the form of an MLflow model registry path, the call to `ModelBuilder` must also specify a tracking server ARN through the `model_metadata["MLFLOW_TRACKING_ARN"]` attribute. **Important** You must use version [2.224.0](https://pypi.org/project/sagemaker/2.224.0/) or later of the SageMaker Python SDK to use `ModelBuilder`. **Note** Use the following code example for reference. For end-to-end examples that show you how to deploy registered MLflow models, see [MLflow tutorials using example Jupyter notebooks](mlflow-tutorials.md). ``` from sagemaker.serve import ModelBuilder from sagemaker.serve.mode.function_pointers import Mode from sagemaker.serve import SchemaBuilder my_schema = SchemaBuilder( sample_input=sample_input, sample_output=sample_output ) model_builder = ModelBuilder( mode=Mode.SAGEMAKER_ENDPOINT, schema_builder=my_schema, role_arn="Your-service-role-ARN", model_metadata={ # both model path and tracking server ARN are required if you use an mlflow run ID or mlflow model registry path as input "MLFLOW_MODEL_PATH": "models:/sklearn-model/1" "MLFLOW_TRACKING_ARN": "arn:aws:sagemaker:region:account-id:mlflow-tracking-server/tracking-server-name" } ) model = model_builder.build() predictor = model.deploy( initial_instance_count=1, instance_type="ml.c6i.xlarge" ) ``` To maintain [lineage tracking](https://docs.aws.amazon.com/sagemaker/latest/dg/lineage-tracking.html) for MLflow models deployed using `ModelBuilder`, you must have the following IAM permissions: + `sagemaker:CreateArtifact` + `sagemaker:ListArtifacts` + `sagemaker:AddAssociation` + `sagemaker:DescribeMLflowTrackingServer` **Important** Lineage tracking is optional. Deployment succeeds without the permissions related to lineage tracking. If you do not have the permissions configured, you will see a lineage tracking permissions error when calling `model.deploy()`. However, the endpoint deployment still succeeds and you can directly interact with your model endpoint. If the permissions above are configured, lineage tracking information is automatically created and stored. For more information and end-to-end examples, see [MLflow tutorials using example Jupyter notebooks](mlflow-tutorials.md). # MLflow tutorials using example Jupyter notebooks The following tutorials demonstrate how to integrate MLflow experiments into your training workflows. To clean up resources created by a notebook tutorial, see [Clean up MLflow resources](mlflow-cleanup.md). You can run SageMaker AI example notebooks using JupyterLab in Studio. For more information on JupyterLab, see [JupyterLab user guide](studio-updated-jl-user-guide.md). Explore the following example notebooks: + [SageMaker Training with MLflow](https://sagemaker-examples.readthedocs.io/en/latest/sagemaker-mlflow/sagemaker_training_mlflow.html) — Train and register a Scikit-Learn model using SageMaker AI in script mode. Learn how to integrate MLflow experiments into your training script. For more information on model training, see [Train a Model with Amazon SageMaker AI](https://docs.aws.amazon.com/sagemaker/latest/dg/how-it-works-training.html). + [SageMaker AI HPO with MLflow](https://sagemaker-examples.readthedocs.io/en/latest/sagemaker-mlflow/sagemaker_hpo_mlflow.html) — Learn how to track your ML experiment in MLflow with Amazon SageMaker AI automatic model tuning (AMT) and the SageMaker AI Python SDK. Each training iteration is logged as a run within the same experiment. For more information about hyperparameter optimization (HPO), see [Perform Automatic Model Tuning with Amazon SageMaker AI](https://docs.aws.amazon.com/sagemaker/latest/dg/automatic-model-tuning.html). + [SageMaker Pipelines with MLflow](https://sagemaker-examples.readthedocs.io/en/latest/sagemaker-mlflow/sagemaker_pipelines_mlflow.html) — Use Amazon SageMaker Pipelines and MLflow to train, evaluate and register a model. This notebook uses the `@step` decorator to build a SageMaker AI Pipeline. For more information on pipelines and the `@step` decorator, see [Create a pipeline with `@step`-decorated functions](https://docs.aws.amazon.com/sagemaker/latest/dg/pipelines-step-decorator-create-pipeline.html). + [Deploy an MLflow Model to SageMaker AI](https://sagemaker-examples.readthedocs.io/en/latest/sagemaker-mlflow/sagemaker_deployment_mlflow.html) — Train a decision tree model using SciKit-Learn. Then, use Amazon SageMaker AI `ModelBuilder` to deploy the model to a SageMaker AI endpoint and run inference using the deployed model. For more information about `ModelBuilder`, see [Deploy MLflow models with `ModelBuilder`](mlflow-track-experiments-model-deployment.md). # Troubleshoot common setup issues Explore common troubleshooting issues. ## Could not find executable named 'groff' When using the AWS CLI, you might encounter the following error: `Could not find executable named 'groff'`. If using a Mac, you can resolve this issue with the following command: ``` brew install groff ``` On a Linux machine, use the following commands: ``` sudo apt-get update -y sudo apt-get install groff -y ``` ## Command not found: jq When creating your AuthZ permission policy JSON file, you might encounter the following error: `jq: command not found`. If using a Mac, you can resolve this issue with the following command: ``` brew install jq ``` On a Linux machine, use the following commands: ``` sudo apt-get update -y sudo apt-get install jq -y ``` ## AWS MLflow plugin installation speeds Installing the AWS MLflow plugin can take several minutes when using a Mac Python environment. ## UnsupportedModelRegistryStoreURIException If you see the `UnsupportedModelRegistryStoreURIException`, do the following: 1. Restart your Jupyter notebook Kernel. 1. Reinstall the AWS MLflow plugin: ``` !pip install --force-reinstall sagemaker-mlflow ``` # Clean up MLflow resources We recommend deleting any resources when you no longer need them. You can delete tracking servers through Amazon SageMaker Studio or using the AWS CLI. You can delete additional resources such as Amazon S3 buckets, IAM roles, and IAM policies using the AWS CLI or directly in the AWS console. **Important** Don't delete the IAM role that you've used to create until you've deleted the tracking server itself. Otherwise, you'll lose access to the tracking server. ## Stop tracking servers We recommend stopping your tracking server when it is no longer in use. You can stop a tracking server in Studio or using the AWS CLI. ### Stop a tracking server using Studio To stop a tracking server in Studio: 1. Navigate to Studio. 1. Choose **MLflow** in the **Applications** pane of the Studio UI. 1. Find the tracking server of your choice in the **MLflow Tracking Servers** pane. Choose the **Stop** icon in the right corner of the tracking server pane. **Note** If your tracking server is **Off**, you see the **Start** icon. If the tracking server is **On**, you see the **Stop** icon. ### Stop a tracking server using the AWS CLI To stop the tracking server using the AWS CLI, use the following command: ``` aws sagemaker stop-mlflow-tracking-server \ --tracking-server-name $ts_name \ --region $region ``` To start the tracking server using the AWS CLI, 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 ``` ## Delete tracking servers You can fully delete a tracking server in Studio or using the AWS CLI. ### Delete a tracking server using Studio To delete a tracking server in Studio: 1. Navigate to Studio. 1. Choose **MLflow** in the **Applications** pane of the Studio UI. 1. Find the tracking server of your choice in the **MLflow Tracking Servers** pane. Choose the vertical menu icon in the right corner of the tracking server pane. Then, choose **Delete**. 1. Choose **Delete** to confirm deletion. ![\[The deletion option on a tracking server card in the MLflow Tracking Servers pane of the Studio UI.\]](http://docs.aws.amazon.com/sagemaker/latest/dg/images/mlflow/mlflow-studio-delete.png) ### Delete a tracking server using the AWS CLI Use the `DeleteMLflowTrackingServer` API to delete any tracking servers that you created. This may take some time. ``` aws sagemaker delete-mlflow-tracking-server \ --tracking-server-name $ts_name \ --region $region ``` To view the status of your tracking server, use the `DescribeMLflowTrackingServer` API and check the `TrackingServerStatus`. ``` aws sagemaker describe-mlflow-tracking-server \ --tracking-server-name $ts_name \ --region $region ``` ## Delete Amazon S3 buckets Delete any Amazon S3 bucket used as an artifact store for your tracking server using the following commands: ``` aws s3 rm s3://$bucket_name --recursive aws s3 rb s3://$bucket_name ``` You can alternatively delete an Amazon S3 bucket associated with your tracking server directly in the AWS console. For more information, see [Deleting a bucket](https://docs.aws.amazon.com/AmazonS3/latest/userguide/delete-bucket.html) in the *Amazon S3 User Guide*. ## Delete registered models You can delete any model groups and model versions created with MLflow directly in Studio. For more information, see [Delete a Model Group](https://docs.aws.amazon.com/sagemaker/latest/dg/model-registry-delete-model-group.html) and [Delete a Model Version](https://docs.aws.amazon.com/sagemaker/latest/dg/model-registry-delete-model-version.html). ## Delete experiments or runs You can use the MLflow SDK to delete experiments or runs. + [mlflow.delete\$1experiment](https://mlflow.org/docs/latest/python_api/mlflow.html?highlight=delete_experiment#mlflow.delete_experiment) + [mlflow.delete\$1run](https://mlflow.org/docs/latest/python_api/mlflow.html?highlight=delete_experiment#mlflow.delete_run) # Amazon SageMaker Experiments in Studio Classic **Important** Experiment tracking using the SageMaker Experiments Python SDK is only available in Studio Classic. We recommend using the new Studio experience and creating experiments using the latest SageMaker AI integrations with MLflow. There is no MLflow UI integration with Studio Classic. If you want to use MLflow with Studio, you must launch the MLflow UI using the AWS CLI. For more information, see [Launch the MLflow UI using the AWS CLI](mlflow-launch-ui.md#mlflow-launch-ui-cli). Amazon SageMaker Experiments Classic is a capability of Amazon SageMaker AI that lets you create, manage, analyze, and compare your machine learning experiments in Studio Classic. Use SageMaker Experiments to view, manage, analyze, and compare both custom experiments that you programmatically create and experiments automatically created from SageMaker AI jobs. Experiments Classic automatically tracks the inputs, parameters, configurations, and results of your iterations as *runs*. You can assign, group, and organize these runs into *experiments*. SageMaker Experiments is integrated with Amazon SageMaker Studio Classic, providing a visual interface to browse your active and past experiments, compare runs on key performance metrics, and identify the best performing models. SageMaker Experiments tracks all of the steps and artifacts that went into creating a model, and you can quickly revisit the origins of a model when you are troubleshooting issues in production, or auditing your models for compliance verifications. ## Migrate from Experiments Classic to Amazon SageMaker AI with MLflow Past experiments created using Experiments Classic are still available to view in Studio Classic. If you want to maintain and use past experiment code with MLflow, you must update your training code to use the MLflow SDK and run the training experiments again. For more information on getting started with the MLflow SDK and the AWS MLflow plugin, see [Integrate MLflow with your environment](mlflow-track-experiments.md). # Example notebooks for Experiments Classic The following example notebooks demonstrate how to track runs for various model training experiments. You can view the resulting experiments in Studio Classic after running the notebooks. For a tutorial that showcases additional features of Studio Classic, see [Amazon SageMaker Studio Classic Tour](gs-studio-end-to-end.md). ## Track experiments in a notebook environment To learn more about tracking experiments in a notebook environment, see the following example notebooks: + [Track an experiment while training a Keras model locally](https://sagemaker-examples.readthedocs.io/en/latest/sagemaker-experiments/local_experiment_tracking/keras_experiment.html) + [Track an experiment while training a Pytorch model locally or in your notebook](https://sagemaker-examples.readthedocs.io/en/latest/sagemaker-experiments/local_experiment_tracking/pytorch_experiment.html) ## Track bias and explainability for your experiments with SageMaker Clarify For a step-by-step guide on tracking bias and explainability for your experiments, see the following example notebook: + [ Fairness and Explainability with SageMaker Clarify ](https://sagemaker-examples.readthedocs.io/en/latest/sagemaker-experiments/sagemaker_clarify_integration/tracking_bias_explainability.html) ## Track experiments for SageMaker training jobs using script mode For more information about tracking experiments for SageMaker training jobs, see the following example notebooks: + [Run a SageMaker AI Experiment with Pytorch Distributed Data Parallel - MNIST Handwritten Digits Classification](https://sagemaker-examples.readthedocs.io/en/latest/sagemaker-experiments/sagemaker_job_tracking/pytorch_distributed_training_experiment.html) + [Track an experiment while training a Pytorch model with a SageMaker Training Job](https://sagemaker-examples.readthedocs.io/en/latest/sagemaker-experiments/sagemaker_job_tracking/pytorch_script_mode_training_job.html) + [Train a TensorFlow model with a SageMaker training job and track it using SageMaker Experiments](https://sagemaker-examples.readthedocs.io/en/latest/sagemaker-experiments/sagemaker_job_tracking/tensorflow_script_mode_training_job.html) # View experiments and runs Amazon SageMaker Studio Classic provides an experiments browser that you can use to view lists of experiments and runs. You can choose one of these entities to view detailed information about the entity or choose multiple entities for comparison. You can filter the list of experiments by entity name, type, and tags. **To view experiments and runs** 1. To view the experiment in Studio Classic, in the left sidebar, choose **Experiments**. Select the name of the experiment to view all associated runs. You can search experiments by typing directly into the **Search** bar or filtering for experiment type. You can also choose which columns to display in your experiment or run list. It might take a moment for the list to refresh and display a new experiment or experiment run. You can click **Refresh** to update the page. Your experiment list should look similar to the following: ![\[A list of experiments in the SageMaker Experiments UI\]](http://docs.aws.amazon.com/sagemaker/latest/dg/images/experiments-classic/experiments-overview.png) 1. In the experiments list, double-click an experiment to display a list of the runs in the experiment. **Note** Experiment runs that are automatically created by SageMaker AI jobs and containers are visible in the Experiments Studio Classic UI by default. To hide runs created by SageMaker AI jobs for a given experiment, choose the settings icon (![\[Black square icon representing a placeholder or empty image.\]](http://docs.aws.amazon.com/sagemaker/latest/dg/images/icons/Settings_squid.png)) and toggle **Show jobs**. ![\[A list of experiment runs in the SageMaker Experiments UI\]](http://docs.aws.amazon.com/sagemaker/latest/dg/images/experiments-classic/experiments-runs-overview.png) 1. Double-click a run to display information about a specific run. In the **Overview** pane, choose any of the following headings to see available information about each run: + **Metrics** – Metrics that are logged during a run. + **Charts** – Build your own charts to compare runs. + **Output artifacts** – Any resulting artifacts of the experiment run and the artifact locations in Amazon S3. + **Bias reports** – Pre-training or post-training bias reports generated using Clarify. + ** Explainability**– Explainability reports generated using Clarify. + **Debugs** – A list of debugger rules and any issues found.