Rastreamento Completo de Modelos de Machine Learning com DVC e Amazon SageMaker AI MLflow Apps

Desafios no rastreamento de modelos de machine learning

Equipes que trabalham com machine learning (ML) enfrentam dificuldades para garantir a rastreabilidade completa de seus modelos, desde os dados e códigos que os treinaram até as versões exatas dos conjuntos de dados utilizados e as métricas dos experimentos que justificaram seu uso em produção. Sem essa rastreabilidade, perguntas como “qual conjunto de dados treinou o modelo atualmente em produção?” ou “é possível reproduzir o modelo implantado há seis meses?” geram investigações demoradas em logs dispersos, notebooks e buckets do Amazon S3.

Essa lacuna é especialmente crítica em setores regulados, como saúde, serviços financeiros e veículos autônomos, onde auditorias exigem vincular modelos implantados aos dados precisos usados no treinamento, além de permitir a exclusão de registros individuais mediante solicitação.

Combinação de ferramentas para rastreamento end-to-end

Para superar esses desafios, uma arquitetura integrada foi proposta, combinando três ferramentas:

• DVC (Data Version Control): para versionamento de datasets e vinculação a commits Git;
• Amazon SageMaker AI: para processamento escalável, treinamento e implantação;
• Amazon SageMaker AI MLflow Apps: para rastreamento de experimentos, registro de modelos e linhagem.

Essa integração cria uma cadeia de rastreabilidade completa, onde cada modelo pode ser vinculado ao conjunto de dados exato usado em seu treinamento.

Arquitetura e fluxo de dados

A arquitetura funciona em quatro etapas principais:

1. Processamento: Um job de processamento do SageMaker AI pré-processa os dados brutos, versiona o dataset processado com DVC, armazenando os dados no Amazon S3 e metadados no repositório Git.
2. Treinamento: Um job de treinamento clona o repositório DVC em um commit específico, recupera o dataset versionado e treina o modelo, registrando métricas e parâmetros no MLflow, incluindo o data_git_commit_id que referencia o commit DVC do dataset.
3. Registro: O modelo treinado é registrado no MLflow Model Registry e pode ser implantado em um endpoint SageMaker AI.
4. Implantação: A partir do registro, o modelo é disponibilizado para inferência em tempo real.

Essa cadeia permite responder perguntas como “qual versão do dataset treinou este modelo?” e “posso reproduzir o treinamento deste modelo?” com comandos simples para recuperar os dados e o código exatos.

Pré-requisitos para replicar a solução

• Conta AWS com permissões para Amazon SageMaker (Processing, Training, MLflow Apps, Endpoints), Amazon S3, AWS CodeCommit e IAM.
• Python 3.11 ou 3.12.
• SageMaker Python SDK v3.4.0 ou superior.
• Repositório companion disponível em GitHub com notebooks e dependências.
• Configuração do IAM para permitir que SageMaker assuma o papel necessário.

Embora o exemplo utilize AWS CodeCommit como backend Git para DVC, é possível adaptar para outros provedores como GitHub ou GitLab, desde que o SageMaker tenha acesso e as credenciais estejam configuradas.

Como DVC e SageMaker AI MLflow Apps se complementam

DVC é uma ferramenta open source que estende o Git para versionamento eficiente de grandes datasets e artefatos ML, armazenando dados pesados no Amazon S3 e versões leves no Git. Ele usa armazenamento endereçado por conteúdo para evitar duplicação desnecessária e suporta pipelines reprodutíveis e gerenciamento de experimentos.

Por outro lado, SageMaker AI MLflow App é um serviço gerenciado que oferece rastreamento de experimentos, registro de modelos com versionamento e gerenciamento do ciclo de vida, além de integração com implantações.

Na arquitetura apresentada, DVC gerencia a linhagem dos dados até o treinamento e MLflow gerencia a linhagem do treinamento até a implantação, conectados pelo commit Git que referencia o dataset exato.

Padrão 1: Linhagem a nível de dataset

O padrão básico demonstra como DVC e MLflow juntos permitem rastrear modelos até a versão exata do dataset usada no treinamento. Usando o dataset CIFAR-10, são realizados dois experimentos:

• v1.0: treinamento com 5% dos dados (~2.250 imagens);
• v2.0: treinamento com 10% dos dados (~4.500 imagens).

O pipeline executa duas etapas em cada versão:

1. Processamento: job SageMaker que baixa, amostra, divide e formata os dados, versiona com DVC, envia os dados para S3 e os metadados para o Git com uma tag única.
2. Treinamento: job SageMaker que clona o repositório DVC na tag especificada, recupera os dados, treina um modelo MobileNetV3-Small pré-treinado e registra parâmetros, métricas e modelo no MLflow.

O data_git_commit_id é registrado em MLflow para vincular o modelo ao dataset exato. A interface MLflow permite comparar métricas, hiperparâmetros e versões de dados entre as execuções, além de acessar o registro dos modelos com histórico e links para o treinamento.

O modelo recomendado (v2.0) pode ser implantado em um endpoint SageMaker para inferência em tempo real.

Limitações do padrão 1

Esse padrão responde perguntas sobre versões de datasets e reprodução de modelos, mas não permite identificar se um registro individual específico estava no dataset usado no treinamento. Para isso, é necessário o padrão 2.

Padrão 2: Linhagem a nível de registro individual para compliance em saúde

Voltado a ambientes regulados, este padrão adiciona rastreabilidade a nível de registros individuais (ex: pacientes) usando manifestos e registros de consentimento.

O manifesto é um arquivo CSV versionado pelo DVC que lista cada registro presente no dataset, contendo campos como patient_id, scan_id, caminho do arquivo, divisão e rótulo.

O registro de consentimento é outro CSV que indica o status de consentimento de cada paciente, usado para filtrar quais registros devem ser incluídos no dataset processado. Atualizações nesse registro (como revogação de consentimento) geram novos datasets versionados e novos treinamentos, garantindo que dados de pacientes que optaram por sair sejam excluídos de modelos futuros.

Fluxo de opt-out demonstrado

1. v1.0: treinamento com todos os pacientes consentidos.
2. Evento de opt-out: paciente PAT-00023 revoga consentimento, registro atualizado no S3.
3. v2.0: processamento e treinamento com dataset limpo, excluindo registros do paciente que optou por sair.
4. Verificação de auditoria: consultas no MLflow confirmam que o paciente aparece apenas no modelo v1.0 e está ausente nos posteriores.

Consultas de auditoria

O módulo utils/audit_queries.py oferece funções para buscar modelos que usaram dados de um paciente específico, verificar exclusão após uma data e listar pacientes por modelo, facilitando a conformidade e auditoria.

Por que essa solução importa no mundo real

Empresas que operam em setores regulados precisam garantir auditabilidade rigorosa e conformidade com legislações de privacidade e proteção de dados. Essa arquitetura integrada oferece uma solução prática e escalável para rastrear modelos de ML desde os dados até a implantação, facilitando auditorias, reprodução de experimentos e respeito a direitos de exclusão de dados.

Links úteis

• Post original no AWS ML Blog
• Repositório companion com notebooks
• Documentação DVC sobre registro de dados
• Amazon SageMaker AI MLflow Apps
• SageMaker AI Pipelines
• AWS CodeConnections