Como simplificar o acesso externo ao Amazon SageMaker MLflow com um proxy REST API

O problema: quando o SDK não é uma opção

Equipes de aprendizado de máquina (ML) usam o MLflow para gerenciar o ciclo de vida dos seus modelos de forma eficiente. O Amazon SageMaker MLflow oferece capacidades abrangentes de rastreamento de experimentos e gerenciamento de modelos. No entanto, muitas empresas possuem requisitos de infraestrutura que exigem integrações baseadas em HTTPS em vez do uso direto do SDK.

O desafio é comum: políticas de segurança corporativa, restrições de rede ou limitações de sistemas legados impedem que as equipes usem o SDK do MLflow diretamente. A AWS publicou uma solução que endereça exatamente esse cenário — um serviço proxy Flask que atua como intermediário seguro entre os sistemas existentes e o SageMaker MLflow.

Visão geral da solução

A arquitetura proposta é composta por três componentes principais que trabalham em conjunto para garantir integração segura e compatível com ambientes corporativos.

Componente 1: Application Load Balancer (ALB)

Um AWS Application Load Balancer atua como roteador de entrada, responsável por:

  • Distribuir o tráfego entre requisições da interface do MLflow e chamadas REST API.
  • Realizar o tratamento inicial das requisições e o roteamento.
  • Suportar domínios customizados e encerramento SSL.

Vale destacar que o ALB é a opção utilizada nessa implementação, mas outras soluções de roteamento, como Nginx, também podem ser utilizadas conforme a necessidade.

Componente 2: Flask MLflow Proxy Service

No centro da arquitetura, uma aplicação Python baseada em Flask é responsável por:

  • Interceptar e processar as requisições HTTPS recebidas.
  • Gerenciar a autenticação AWS e a assinatura das requisições.
  • Transformar URLs para acesso seguro aos endpoints do MLflow.
  • Encaminhar as respostas de volta aos clientes.

Componente 3: Amazon SageMaker MLflow

O serviço gerenciado da AWS suporta dois modos de implantação do MLflow:

  • MLflow Tracking Server — servidor de rastreamento gerenciado.
  • MLflowApp — aplicação MLflow serverless.

Em ambos os casos, o serviço fornece o armazenamento de metadados de experimentos e os arquivos de modelos.

Fluxo de requisições na arquitetura

O fluxo de uma requisição nessa arquitetura segue uma sequência bem definida. Quando um cliente inicia uma chamada HTTPS, ela primeiro chega ao ALB, que a encaminha para o serviço proxy Flask. A partir daí, o proxy realiza uma série de operações críticas:

  • Autenticação via integração com o Gerenciamento de Identidade e Acesso da AWS (IAM).
  • Transformação e pré-assinatura de URLs para acesso seguro.
  • Processamento dos endpoints REST do MLflow conforme necessário.

Após essas transformações, a requisição autenticada é enviada ao SageMaker MLflow, que a processa e retorna a resposta. O proxy então encaminha essa resposta de volta ao cliente original. Toda essa cadeia mantém a segurança sem exigir que o cliente conheça os detalhes internos da autenticação AWS.

Pré-requisitos para implantação

Antes de começar, é necessário ter os seguintes recursos disponíveis:

Do ponto de vista de conhecimento, é esperado que o profissional tenha familiaridade básica com serviços AWS e permissões IAM, com aplicações Python e Flask, e com os conceitos e operações do MLflow.

Em relação aos custos, a solução cria recursos AWS que podem gerar cobranças, principalmente instâncias Amazon EC2, o Application Load Balancer, recursos Amazon SageMaker AI e armazenamento Amazon S3. Para estimar os valores, a AWS disponibiliza a Calculadora de Preços AWS.

Implantando a solução

O processo de implantação leva aproximadamente 40 minutos e está dividido em três etapas.

Etapa 1: Implantação da infraestrutura com AWS CDK

O primeiro passo é baixar o código da solução e instalar as dependências:

# 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

Em seguida, é necessário fazer o bootstrap do ambiente para o AWS CDK. Esse passo pode ser ignorado se a conta e a região AWS já estiverem configuradas para o CDK:

npx cdk bootstrap aws://<ACCOUNT_ID>/<REGION>

A solução é composta por quatro stacks CDK:

  • Networking stack — cria a VPC e os componentes de rede.
  • SageMaker AI domain stack — configura o domínio SageMaker.
  • SageMaker MLflow stack — implanta o servidor de rastreamento ou o app serverless do MLflow.
  • Flask application stack — implanta o serviço proxy MLflow.

Para implantar todas as stacks de uma vez, utilize um dos comandos abaixo conforme o modo desejado.

Para implantação baseada em tracking server:

npx cdk deploy --all --require-approval=never -c mlflowType=tracking

Para implantação baseada em app serverless:

npx cdk deploy --all --require-approval=never -c mlflowType=serverless

Etapa 2: Instalação e configuração do Flask MLflow Proxy Service

Após a implantação da infraestrutura, é necessário conectar à instância EC2. O ID da instância pode ser obtido na saída do CDK ou na seção de outputs da stack sagemaker-infra-flaskapp-{mlflowType} no AWS CloudFormation. A conexão deve ser feita pelo AWS Systems Manager Session Manager, seguindo o guia de conexão do Session Manager.

Com a sessão aberta, instale o Python 3.13 e as dependências:

# Switch to root user
sudo su -
cd /root

# Install Python and dependencies
chmod +x install_python13.sh
./install_python13.sh

O script foi desenvolvido para sistemas baseados em Ubuntu. Para outras distribuições Linux, é necessário instalar manualmente o Python 3.12+, PIP3 e Virtualenv usando o gerenciador de pacotes do sistema.

Em seguida, instale e inicie o serviço proxy MLflow:

chmod +x setup_mlflow_proxy_app.sh
./setup_mlflow_proxy_app.sh

Para verificar o status do serviço:

systemctl status mlflowproxy

Caso o serviço não esteja em execução, os logs podem ser consultados com o seguinte comando:

journalctl -u mlflowproxy

Etapa 3: Validação do acesso à REST API do MLflow

Com a solução implantada, é possível interagir com as REST APIs do MLflow por meio do ALB. Os exemplos abaixo utilizam o protocolo HTTP (sem segurança). Para ambientes de produção, a AWS recomenda o uso de HTTPS. As chamadas podem ser feitas com curl ou qualquer outra ferramenta de preferência — os comandos funcionam da mesma forma tanto para o modo tracking server quanto para o modo serverless.

Primeiro, obtenha o nome DNS do ALB com o seguinte comando na sua estação de trabalho:

aws cloudformation describe-stacks --stack-name sagemaker-infra-flaskapp-{mlflowType} --query 'Stacks[0].Outputs[?OutputKey==`ALBUrl`].OutputValue' --output text

Com o DNS em mãos, substitua <ALB DNS>, <EXP ID>, <RUN ID> e <RUN NAME> pelos valores correspondentes nos comandos abaixo.

Criar um experimento:

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

Buscar experimentos:

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

Obter um experimento:

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

Criar uma execução dentro de um experimento:

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>"}'

Listar artefatos de uma execução:

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

Definir uma tag em uma execução:

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"}'

Excluir uma execução:

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

Também é possível abrir a interface do MLflow para visualizar as alterações realizadas pelos comandos acima. Para instruções sobre como acessar a interface, consulte como iniciar a interface do MLflow usando uma URL pré-assinada.

Limpeza dos recursos

Para evitar cobranças contínuas, os recursos criados pela solução devem ser removidos. Para destruir as stacks gerenciadas pelo CDK, execute o comando adequado ao modo de implantação utilizado.

Para tracking server:

npx cdk destroy --all -c mlflowType=tracking

Para app serverless:

npx cdk destroy --all -c mlflowType=serverless

As stacks de rede e do domínio SageMaker são compartilhadas entre os modos de implantação e só serão removidas quando o último par de stacks MLflow/Flask for excluído. Além disso, alguns recursos podem exigir exclusão manual por conta de políticas de retenção ou dependências:

  • Buckets Amazon S3: acesse o console do S3, identifique os buckets criados pela solução, esvazie-os e exclua-os.
  • Grupos de logs do Amazon CloudWatch: no console do CloudWatch, localize e exclua os grupos de logs associados à solução.

Considerações de segurança para produção

Ao levar essa solução para um ambiente produtivo, a AWS recomenda uma série de medidas adicionais de segurança:

  • Configure o monitoramento com Amazon CloudWatch para acompanhar a saúde da aplicação proxy, detectar anomalias e configurar alertas para atividades suspeitas.
  • Implemente limitação de taxa (rate limiting) no serviço proxy Flask para proteção contra ataques de negação de serviço (DoS). O AWS WAF (firewall de aplicação web) pode ser usado junto ao ALB para implementar regras baseadas em taxa.
  • Implante um ALB interno (não voltado para a internet) para restringir o acesso ao proxy apenas à rede privada, permitindo tráfego somente de dentro da VPC ou de redes conectadas via VPC Peering ou AWS Transit Gateway.
  • Habilite o encerramento HTTPS no nível do ALB para comunicação segura entre clientes e a aplicação. O AWS Certificate Manager (ACM) pode ser usado para provisionar e gerenciar certificados SSL/TLS. Para detalhes sobre a configuração, consulte a documentação de listeners HTTPS do Application Load Balancer.

Conclusão

A solução apresentada pela AWS oferece uma forma prática de integrar o Amazon SageMaker MLflow com sistemas corporativos existentes, sem abrir mão dos padrões de segurança e infraestrutura já estabelecidos. O serviço proxy Flask atua como uma ponte entre o mundo HTTPS das aplicações empresariais e as APIs autenticadas da AWS, reduzindo a complexidade de implementação e permitindo que as equipes de ML continuem operando com seus fluxos de trabalho atuais.

Os principais benefícios da abordagem incluem integração com controles de segurança corporativos já existentes, mínimas alterações nos fluxos de trabalho de ML, menor complexidade de implantação, integração via REST API e compatibilidade com serviços de proxy corporativos.

Para aprofundar o conhecimento sobre o tema, a AWS disponibiliza a documentação do Amazon SageMaker MLflow, além de materiais sobre servidores de rastreamento MLflow e MLflow apps.

Fonte

Streamline external access to Amazon SageMaker MLflow using a REST API proxy (https://aws.amazon.com/blogs/machine-learning/streamline-external-access-to-amazon-sagemaker-mlflow-using-a-rest-api-proxy/)

Comments

Leave a Reply

Your email address will not be published. Required fields are marked *