Streamline external access to Amazon SageMaker MLflow using a REST API proxy

Machine learning (ML) teams use MLflow to manage their ML lifecycle effectively. Amazon SageMaker MLflow provides comprehensive ML experiment tracking and model management capabilities. However, many enterprises have existing infrastructure requirements that need HTTPS-based integrations rather than direct SDK usage.

Many organizations need to integrate Amazon SageMaker MLflow with their established systems while maintaining their security and infrastructure patterns. This integration challenge affects teams who can’t use the SDK directly because of corporate security policies, network restrictions, or legacy system constraints.

In this post, we demonstrate how to build a secure Flask-based MLflow proxy service that provides HTTPS access to Amazon SageMaker MLflow without requiring the MLflow SDK. This solution is for organizations undergoing cloud transformation who want to preserve their existing ML workflows while adopting cloud-native services.

Implementing the MLflow proxy service for MLflow HTTPS requests.

Configuring AWS Identity and Access Management (IAM) authentication for secure access.

Managing URL pre-signing and request transformation.

Access SageMaker MLflow securely through standard HTTPS endpoints.

Maintain compliance with your organization’s security requirements.

Integrate MLflow with existing enterprise systems.

Reduce implementation complexity and maintenance overhead.

A lightweight Flask-based MLflow proxy architecture provides secure integration between enterprise systems and Amazon SageMaker MLflow through three key components.

Traffic distribution for MLflow UI and REST API requests.

Initial request handling and routing.

Support for custom domain names and SSL termination.

Note: This implementation uses ALB, but you can alternatively use other routing solutions such as Nginx based on your requirements.

Intercepting and processing incoming HTTPS requests.

Managing AWS authentication and request signing.

Transforming URLs for secure MLflow endpoint access.

Handling response routing back to clients.

Support for two MLflow deployment modes: MLflow Tracking Server – managed MLflow tracking server. MLflowApp – serverless MLflow application.

MLflow Tracking Server – managed MLflow tracking server.

MLflowApp – serverless MLflow application.

Backend metadata store for tracking information.

Storage for model files and data.

This architecture provides secure communication while maintaining compatibility with existing enterprise systems. The proxy service acts as a bridge, transforming standard HTTPS requests into authenticated AWS API calls that can interact with SageMaker MLflow.

The following diagram shows how the Flask proxy service provides secure communication between external clients and Amazon SageMaker MLflow.

An ALB that handles incoming traffic.

A Flask proxy service that manages authentication and request transformation.

Amazon SageMaker MLflow that processes ML operations.

Let’s explore how requests flow through this architecture to provide secure MLflow access.

When a client initiates an HTTPS request, it first reaches the ALB, which acts as the entry point for all incoming traffic. The ALB then routes these requests to the MLflow proxy service.

Handles authentication through AWS IAM integration.

Transforms URLs and pre-signs them for secure access.

Processes the MLflow REST API endpoints as needed.

The MLflow proxy service transforms the incoming request into an authenticated AWS request before making the API call to SageMaker MLflow REST endpoints. After SageMaker MLflow processes the request, it returns a response which the MLflow proxy service processes and routes back to the original client.

This workflow maintains security while providing integration between enterprise systems and SageMaker MLflow.

A workstation with the following tools installed: AWS Command Line Interface (AWS CLI) configured with permissions to create: Amazon Virtual Private Cloud (Amazon VPC) and associated networking components. Amazon Elastic Compute Cloud (Amazon EC2) instances. Amazon SageMaker AI resources. Amazon Simple Storage Service (Amazon S3) buckets. AWS Identity and Access Management (IAM) roles and policies. AWS CloudFormation stacks. AWS Application Load Balancers. Node.js version 18.0.0 or later. NPM. AWS Cloud Development Kit (AWS CDK) CLI version 2.100.0 or later. Python 3.x with pip or pip3.

AWS Command Line Interface (AWS CLI) configured with permissions to create: Amazon Virtual Private Cloud (Amazon VPC) and associated networking components. Amazon Elastic Compute Cloud (Amazon EC2) instances. Amazon SageMaker AI resources. Amazon Simple Storage Service (Amazon S3) buckets. AWS Identity and Access Management (IAM) roles and policies. AWS CloudFormation stacks. AWS Application Load Balancers.

Amazon Virtual Private Cloud (Amazon VPC) and associated networking components.

Amazon Elastic Compute Cloud (Amazon EC2) instances.

Amazon SageMaker AI resources.

Amazon Simple Storage Service (Amazon S3) buckets.

AWS Identity and Access Management (IAM) roles and policies.

AWS CloudFormation stacks.

AWS Application Load Balancers.

Node.js version 18.0.0 or later.

AWS Cloud Development Kit (AWS CDK) CLI version 2.100.0 or later.

Python 3.x with pip or pip3.

Required knowledge: Basic understanding of AWS services and IAM permissions. Familiarity with Python and Flask applications. Understanding of MLflow concepts and operations.

Basic understanding of AWS services and IAM permissions.

Familiarity with Python and Flask applications.

Understanding of MLflow concepts and operations.

Cost considerations: This solution creates AWS resources that might incur costs. Key cost-driving resources include: Amazon EC2 instances. Application Load Balancer. Amazon SageMaker AI resources. Amazon S3 storage.

This solution creates AWS resources that might incur costs.

Key cost-driving resources include: Amazon EC2 instances. Application Load Balancer. Amazon SageMaker AI resources. Amazon S3 storage.

Amazon EC2 instances.

Application Load Balancer.

Amazon SageMaker AI resources.

Amazon S3 storage.

For information about AWS service pricing, see AWS Pricing Calculator.

This section walks you through deploying the solution in your AWS account and validating it. The deployment process takes approximately 40 minutes.

Download the solution code and install dependencies: # Clone the repository git clone https://github.com/aws-samples/sample-sagemaker-mlflow-rest-apis.git # Navigate to project directory and install dependencies cd sample-sagemaker-mlflow-rest-apis npm ci

Bootstrap your environment for AWS CDK. Skip this step if your AWS account and Region are already bootstrapped for AWS CDK.Bootstrap the AWS account and Region for CDK: npx cdk bootstrap aws://<ACCOUNT_ID>/<REGION>

Deploy the required resources on your AWS account.The solution consists of four CDK stacks: Networking stack — creates the VPC and networking components. SageMaker AI domain stack — sets up the SageMaker domain. SageMaker MLflow stack — deploys the MLflow tracking server or MLflow serverless app. Flask application stack — deploys the MLflow proxy service. Deploy all the stacks with one of the following commands. For tracking server based deployment: npx cdk deploy --all --require-approval=never -c mlflowType=tracking For serverless app based deployment: npx cdk deploy --all --require-approval=never -c mlflowType=serverless

Networking stack — creates the VPC and networking components.

SageMaker AI domain stack — sets up the SageMaker domain.

SageMaker MLflow stack — deploys the MLflow tracking server or MLflow serverless app.

Flask application stack — deploys the MLflow proxy service.

Deploy all the stacks with one of the following commands.

Connect to the EC2 instance: Note the Amazon EC2 instance ID from the CDK output or from the sagemaker-infra-flaskapp-{mlflowType} AWS CloudFormation stack output section. Use AWS Systems Manager Session Manager to connect. Follow the Session Manager connection guide.

Note the Amazon EC2 instance ID from the CDK output or from the sagemaker-infra-flaskapp-{mlflowType} AWS CloudFormation stack output section.

Use AWS Systems Manager Session Manager to connect. Follow the Session Manager connection guide.

Install Python 3.13 and dependencies.Install Python packages: # Switch to root user sudo su - cd /root # Install Python and dependencies chmod +x install_python13.sh ./install_python13.sh Note: This script is designed for Ubuntu-based systems. For other Linux distributions, install Python 3.12+, PIP3, and Virtualenv using your system’s package manager.

Note: This script is designed for Ubuntu-based systems. For other Linux distributions, install Python 3.12+, PIP3, and Virtualenv using your system’s package manager.

Install and start the MLflow proxy service: chmod +x setup_mlflow_proxy_app.sh ./setup_mlflow_proxy_app.sh

This section demonstrates how to interact with MLflow REST APIs through the ALB.

Note: These examples use the HTTP (unsecured) protocol. For production environments, we recommend HTTPS. We use curl to make the API requests in this post, but you can use any tool you prefer. The provided curl commands work identically for both tracking server and serverless modes; the proxy service handles the differences transparently.

Get your ALB DNS name by running the following command on your workstation: aws cloudformation describe-stacks --stack-name sagemaker-infra-flaskapp-{mlflowType} --query 'Stacks[0].Outputs[?OutputKey==`ALBUrl`].OutputValue' --output text

Test MLflow API endpoints by running the following commands on your workstation. Replace <ALB DNS>, <EXP ID>, <RUN ID>, and <RUN NAME> with appropriate values. Create an experiment: curl -X POST http://<ALB DNS>/ajax-api/2.0/mlflow/experiments/create -H "Content-Type: application/json" -d '{"name": "mlflow-experiment"}' Search experiments: curl -X POST http://<ALB DNS>/ajax-api/2.0/mlflow/experiments/search -H "Content-Type: application/json" -d '{"max_results": 5}' Get an experiment: curl -X GET 'http://<ALB DNS>/ajax-api/2.0/mlflow/experiments/get?experiment_id=0' Create a run inside an experiment: curl -X POST http://<ALB DNS>/ajax-api/2.0/mlflow/runs/create -H "Content-Type: application/json" -d '{"experiment_id": <EXP ID>, "run_name": "<RUN NAME>"}' List artifacts from a run: curl -X GET "http://<ALB DNS>/ajax-api/2.0/mlflow/artifacts/list?run_id=<RUN ID>" Set a tag on a run: curl -X POST "http://<ALB DNS>/ajax-api/2.0/mlflow/runs/set-tag" -H "Content-Type: application/json" -d '{"run_id": "<RUN ID>", "key": "model_type","value": "api-test"}' Delete a run: curl -X POST http://<ALB DNS>/ajax-api/2.0/mlflow/runs/delete -H "Content-Type: application/json" -d '{"run_id": "<RUN ID>"}' Note: You can also open the MLflow UI and view the changes you make using the preceding curl commands. For instructions on launching the MLflow UI, see Launch the MLflow UI using a presigned URL.

Create an experiment: curl -X POST http://<ALB DNS>/ajax-api/2.0/mlflow/experiments/create -H "Content-Type: application/json" -d '{"name": "mlflow-experiment"}'

Search experiments: curl -X POST http://<ALB DNS>/ajax-api/2.0/mlflow/experiments/search -H "Content-Type: application/json" -d '{"max_results": 5}'

Get an experiment: curl -X GET 'http://<ALB DNS>/ajax-api/2.0/mlflow/experiments/get?experiment_id=0'

Create a run inside an experiment: curl -X POST http://<ALB DNS>/ajax-api/2.0/mlflow/runs/create -H "Content-Type: application/json" -d '{"experiment_id": <EXP ID>, "run_name": "<RUN NAME>"}'

List artifacts from a run: curl -X GET "http://<ALB DNS>/ajax-api/2.0/mlflow/artifacts/list?run_id=<RUN ID>"

Set a tag on a run: curl -X POST "http://<ALB DNS>/ajax-api/2.0/mlflow/runs/set-tag" -H "Content-Type: application/json" -d '{"run_id": "<RUN ID>", "key": "model_type","value": "api-test"}'

Delete a run: curl -X POST http://<ALB DNS>/ajax-api/2.0/mlflow/runs/delete -H "Content-Type: application/json" -d '{"run_id": "<RUN ID>"}'

Note: You can also open the MLflow UI and view the changes you make using the preceding curl commands. For instructions on launching the MLflow UI, see Launch the MLflow UI using a presigned URL.

Delete CDK-managed resources.Navigate to the root directory of the cloned repository on your workstation and run the following.For tracking server based deployment: npx cdk destroy --all -c mlflowType=tracking For serverless app based deployment: npx cdk destroy --all -c mlflowType=serverless Note: The networking and SageMaker domain stacks are shared across both deployment modes. AWS CDK only deletes them when the last MLflow or Flask app stack pair is removed.

Note: The networking and SageMaker domain stacks are shared across both deployment modes. AWS CDK only deletes them when the last MLflow or Flask app stack pair is removed.

Manual resource cleanup. Some resources might require manual deletion because of retention policies or dependencies: Amazon S3 buckets: Navigate to the Amazon S3 console. Identify the buckets created by this solution. Empty each bucket and delete it. Amazon CloudWatch log groups: In the CloudWatch console, find the log groups associated with this solution. Delete these log groups.

Amazon S3 buckets: Navigate to the Amazon S3 console. Identify the buckets created by this solution. Empty each bucket and delete it.

Navigate to the Amazon S3 console.

Identify the buckets created by this solution.

Empty each bucket and delete it.

Amazon CloudWatch log groups: In the CloudWatch console, find the log groups associated with this solution. Delete these log groups.

In the CloudWatch console, find the log groups associated with this solution.

Delete these log groups.

Configure Amazon CloudWatch monitoring for the Flask-based proxy service to track application health, detect anomalies, and set up alerts for suspicious activities.

Implement rate limiting for the Flask-based proxy service to protect against potential denial-of-service (DoS) attacks and control the number of requests from individual clients. You can use AWS WAF (web application firewall) with the ALB to implement rate-based rules.

Deploy an internal (non-internet-facing) ALB to restrict proxy access to your private network. This setup makes sure that only traffic from within your VPC or connected networks can reach the service. Connect through VPC peering or AWS Transit Gateway.

Enable HTTPS termination at the ALB level for secure communication between clients and your application. You can use AWS Certificate Manager (ACM) to provision and manage SSL/TLS certificates for your application. For instructions on configuring HTTPS listeners, see the Application Load Balancer HTTPS listeners documentation.

These security measures help protect the Flask application against common web vulnerabilities and provide secure communication between components.

In this post, we showed how to build a secure Flask-based proxy service that provides HTTPS access to Amazon SageMaker MLflow. This solution helps organizations bridge their existing infrastructure with AWS managed MLflow capabilities while maintaining enterprise security requirements.

Integration with existing enterprise security controls.

Minimal changes to existing ML workflows.

Reduced deployment complexity.

REST API integration.

Compatibility with enterprise proxy services.

Explore the Amazon SageMaker MLflow documentation.

Read about MLflow tracking servers and MLflow apps.

Try this solution in your own environment and let us know your experience in the comments.