We are excited to inform that a fully managed MLflow capability on Amazon SageMaker is now generally available. A popular open-source tool called MLflow is essential for assisting machine learning (ML) teams in handling the full ML lifecycle. Customers can now quickly and easily set up and operate MLflow Tracking Servers with only a few clicks thanks to this new introduction, which streamlines the procedure and increases productivity.
With MLflow, data scientists and machine learning developers may monitor several iterations of model training as runs within experiments, assess models, compare runs with visualisations, and add the top models to a Model Registry. With Amazon SageMaker, ML administrators can quickly and easily create secure, scalable MLflow environments on AWS by doing away with the monotonous heavy lifting involved in MLflow setup and management.
Essential elements of SageMaker’s managed MLflow
Three essential elements form the foundation of SageMaker’s fully managed MLflow capability:
- MLflow Tracking Server: Using the SageMaker Studio UI, you can quickly establish an MLflow Tracking Server. You may start effectively monitoring your machine learning experiments by using this standalone HTTP server, which provides numerous REST API endpoints for tracking runs and experiments. An additional option for finer-grained security customisation is to use the AWS Command Line Interface (AWS CLI).
- MLflow backend metadata store: This essential component of the MLflow Tracking Server is where all experiment, run, and artifact-related metadata is stored. This ensures thorough recording and management of your machine learning experiments and contains information such as experiment titles, run IDs, parameter values, metrics, tags, and artefact locations.
- The MLflow artefact store component offers a place to store all of the artefacts produced by machine learning experiments, including datasets, trained models, logs, and graphs. It provides a customer-managed AWS account for safely and effectively storing these artefacts using an Amazon Simple Storage Service (Amazon S3) bucket.
Amazon SageMaker’s advantages with MLflow
Your machine learning workflows can be improved and streamlined by integrating Amazon SageMaker with MLflow:
- Complete Experiment Monitoring: Monitor MLflow experiments in managed IDEs in SageMaker Studio, local IDEs, training jobs, processing tasks, and pipelines.
- Complete MLflow Functionality: Utilise all of MLflow’s experimental features, including MLflow Model Registry, MLflow Evaluations, and Tracking, to quickly compare and assess the outcomes of training iterations.
- Unified Model Governance: You may deploy MLflow models to SageMaker inference without the need to create special containers thanks to the unified model governance experience provided by models registered in MLflow, which automatically appear in the SageMaker Model Registry.
- Effective Server Management: Use the SageMaker Studio UI or APIs to provision, delete, and upgrade MLflow Tracking Servers as needed. Customers don’t need to handle the underlying infrastructure because SageMaker takes care of your tracking servers’ scaling, patching, and continuous maintenance.
- Enhanced Security: AWS Identity and Access Management (IAM) provides safe access to MLflow Tracking Servers. To provide strong security for your machine learning environments, create IAM policies that allow or prohibit access to particular MLflow APIs.
- Sufficient Oversight and Management: Amazon EventBridge and AWS CloudTrail can be used to monitor an MLflow Tracking Server’s behaviour in order to facilitate efficient Tracking Server governance.
Prerequisites for the MLflow Tracking Server (environment setup)
Make a domain for SageMaker Studio
With the new SageMaker Studio experience, you may establish a SageMaker Studio domain.
Set up the role for IAM execution
To register models in SageMaker and read and publish artefacts to Amazon S3, the MLflow Tracking Server requires an IAM execution role. The Tracking Server execution role can be either created as a separate role or it can be used as the Studio domain execution role. For further information on the IAM position, see the SageMaker Developer Guide if you decide to build a new role for this. For information on what IAM policy the Studio domain execution role requires, see the SageMaker Developer Guide if you decide to upgrade it.
Set up the tracking server for MLflow
Throughout the guide, you create an MLflow Tracking Server using the default options, which include the Tracking Server execution role (Studio domain execution role), the Tracking Server version (2.13.2), and the Tracking Server size (Small). We advise utilising a Small Tracking Server for teams of up to 25 users, as the size of the tracking server affects how much utilisation it can accommodate.
To begin, pick MLflow under Applications in your SageMaker Studio domain that was created during the environment setup previously described, then click Create.
Next, give the Tracking Server’s Name and Artefact store location (S3 URI).
An MLflow Tracking Server can take up to twenty-five minutes to set up.
Monitor and contrast training runs
You’ll need your Tracking Server ARN, which was assigned during the creation process, and a Jupyter Notebook in order to begin recording metrics, parameters, and artefacts to MLflow. Training runs can be tracked using the MLflow SDK and compared with the MLflow UI.
The sagemaker-mlflow plugin is required to authenticate all MLflow API queries made by the MLflow SDK using AWS Signature V4 in order to register models from the MLflow Model Registry to the SageMaker Model Registry.
Enumerate potential models
The model whose metrics best fit your needs can be registered in the MLflow Model Registry once you’ve compared the several runs as explained in Step 4. A model’s registration suggests that it might be suitable for production deployment; additional testing will be required to confirm this suitability. You may deploy MLflow models to SageMaker inference by having them automatically appear in the SageMaker Model Registry after they have been registered in MLflow. This provides a consistent model governance experience. This makes it possible for data scientists who use MLflow primarily for testing to transfer their models to ML engineers, who utilise the SageMaker Model Registry to oversee and manage model deployments in production.
Pricing
An MLflow Tracking Server will cost you money once it is built until you stop it or delete it. The amount of data logged to the tracking servers, the size chosen, and the length of time the servers have been operational all factor into the cost of the tracking servers. When not in use, Tracking Servers can be terminated using the SageMaker Studio UI or API to reduce expenses. See the Amazon SageMaker pricing for additional information on costs.
Currently accessible
With the exception of China and US GovCloud Regions, SageMaker with MLflow is typically accessible in all AWS Regions where SageMaker Studio is available. AWS cordially encourage you to investigate this new feature and see how it might improve the control and efficiency of your machine learning endeavours.