Inferência de Machine Learning com Criptografia Ponta a Ponta no Amazon SageMaker AI e FHE

O problema: como processar dados sensíveis na nuvem sem expô-los?

Inferência de aprendizado de máquina (ML) frequentemente envolve informações altamente sensíveis — prontuários médicos, dados corporativos proprietários ou comunicações pessoais. A grande questão é: como aproveitar a escalabilidade da nuvem sem entregar esses dados a terceiros?

A AWS publicou uma abordagem técnica que responde exatamente a isso: usar o Amazon SageMaker AI combinado com Criptografia Totalmente Homomórfica (FHE — Fully Homomorphic Encryption) para executar inferência de ML com os dados permanecendo criptografados durante todo o processo — inclusive enquanto o modelo está processando a consulta. Nem o próprio SageMaker AI tem acesso ao conteúdo em texto simples.

O que é Criptografia Totalmente Homomórfica (FHE)?

FHE é uma forma de criptografia que permite realizar operações matemáticas diretamente sobre dados criptografados, sem precisar descriptografá-los. No contexto de ML, isso significa que um modelo pode receber uma consulta cifrada, processá-la e devolver uma predição também cifrada — sem jamais ver o conteúdo original. Apenas o cliente, que possui a chave privada, consegue descriptografar o resultado.

Essa abordagem é diferente — e complementar — a ambientes de computação confidencial como o Nitro System da AWS no Amazon EC2. Com o AWS Nitro Enclaves, os dados são descriptografados e processados em ambientes isolados de hardware. Com FHE, os dados nunca são descriptografados: a segurança é garantida pela matemática, não pelo hardware.

Casos de uso práticos

A AWS ilustra três cenários onde essa capacidade é especialmente valiosa:

Saúde: uma operadora de plano de saúde quer oferecer a médicos um modelo que prediz resultados de procedimentos com base em dados diagnósticos. Publicar o modelo na nuvem simplifica a operação, mas regulações de privacidade impedem expor dados de pacientes a terceiros.
Setor de energia: uma empresa de petróleo e gás usa ML para avaliar fotos de satélite de possíveis locais de perfuração. Hospedar o modelo na nuvem reduz custos, mas as imagens envolvem locais politicamente sensíveis.
Telecomunicações: uma operadora quer processar e-mails de clientes para detectar spam e phishing. A escala exige ML na nuvem, mas regulações de proteção de dados exigem que as mensagens permaneçam criptografadas em terceiros.

A biblioteca concrete-ml: uma abordagem de alto nível

Um post anterior no blog da AWS já havia abordado FHE para inferência no SageMaker, mostrando como implementar um algoritmo de regressão linear do zero usando a biblioteca de baixo nível SEAL. Desta vez, a abordagem é mais flexível: usa a concrete-ml, uma biblioteca de alto nível desenvolvida especificamente para inferência com FHE, compatível com a API do scikit-learn e que suporta vários tipos de modelos prontos para uso.

Vale destacar: a concrete-ml é disponibilizada pela Zama para prototipagem ou uso não comercial sem licença paga. Para uso comercial, pode ser necessária uma licença comercial.

Visão geral da solução

O fluxo completo da solução funciona da seguinte forma:

O dono do modelo prepara e normaliza os dados de treinamento (por exemplo, para a escala [-1, 1]).
O modelo é treinado em uma versão habilitada para FHE — projetada para operar sobre dados criptografados.
O modelo é hospedado no SageMaker AI.
Clientes criptografam suas consultas usando o esquema FHE compatível com o modelo.
As consultas criptografadas são enviadas ao endpoint, que as processa sem descriptografar nenhum valor.
O endpoint retorna a predição criptografada ao cliente, que a descriptografa localmente.

Pré-requisitos

Para implementar essa solução, são necessários:

Ambiente local com Python 3.12, pip e Docker instalados. O uso de ambientes virtuais é recomendado.
Uma conta AWS com: repositórios no Amazon ECR para as imagens de container; buckets no Amazon S3 para o modelo, código de treinamento, chaves e textos cifrados (seguindo as boas práticas de segurança do S3); e roles no IAM para o criador do modelo, criador do endpoint, o próprio endpoint e os clientes.

Políticas de IAM para essas roles, junto com um exemplo funcional usando o corpus MNIST de dígitos manuscritos, estão disponíveis no repositório de código de exemplo.

Treinamento

Container de treinamento

O container de treinamento é construído a partir de uma imagem python:3.12, com as dependências sagemaker_training, concrete-ml, concrete-python e torch instaladas. A atenção à paridade de versões entre Python, concrete-ml e concrete-python é obrigatória — a biblioteca exige consistência em todo o sistema.

O Dockerfile.training tem o seguinte conteúdo:

FROM python:3.12
RUN apt-get update && apt-get upgrade -y && apt-get clean
RUN apt-get -y install --no-install-recommends cmake
RUN pip install sagemaker_training==5.1.1 concrete-ml==1.9.0 concrete-python==2.10.0 torch==2.3.1

Após construir a imagem localmente, ela deve ser enviada ao Amazon ECR com os comandos de autenticação, tag e push padrão.

Script de treinamento

O treinamento com concrete-ml não difere de qualquer outro framework de ML — e o container de treinamento funciona como qualquer outro container de treinamento customizado. O treinamento ocorre sobre dados em texto simples; a concrete-ml não exige pré-processamento além da normalização. O training_script.py segue este template:

import argparse
import os
import numpy
from concrete.ml.sklearn import <Model class to train>
from concrete.ml.deployment import FHEModelDev

def do_training(model_dir, train):
    # Load your data from the train directory
    # Train your model instance, then save it
    # with the following line.
    FHEModelDev(model_dir, model).save()

def model_fn(model_dir):
    # SageMaker AI requires this function exist but doesn't use it
    raise NotImplementedError

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--model-dir', type=str, default=os.environ['SM_MODEL_DIR'])
    parser.add_argument('--train', type=str, default=os.environ['SM_CHANNEL_TRAINING'])
    args = parser.parse_args()
    do_training(args.model_dir, args.train)

Framework customizado e execução do job

Para integrar o container ao SageMaker AI, a AWS recomenda criar um framework customizado. O framework.py encapsula as configurações do container e simplifica o lançamento do job de treinamento.

O job é iniciado com o start_training.py, que usa a instância ml.m5.xlarge para modelos pequenos ou ml.m5.4xlarge para modelos maiores. Após a conclusão, dois arquivos são salvos no bucket S3: server.zip (usado pelo endpoint de inferência) e client.zip (usado pelos clientes para criptografar consultas).

Inferência

Desafios técnicos específicos do FHE

A inferência baseada em FHE traz algumas restrições técnicas novas em relação à inferência padrão:

Clientes precisam do client.zip para gerar chaves criptográficas.
Textos cifrados FHE podem ultrapassar os limites de tamanho de requisição do SageMaker AI — portanto, cliente e serviço precisam trocá-los via Amazon S3.
A avaliação FHE pode demorar mais do que os timeouts padrão do SageMaker AI, exigindo o uso de inferência assíncrona.
O endpoint precisa de uma chave de avaliação (um tipo de chave pública) fornecida pelo cliente para executar a avaliação FHE.

Predictor do endpoint

O predictor.py é um servidor Flask que recebe um JSON com dois URIs do S3 — um para a consulta criptografada e outro para a chave de avaliação. Ele baixa esses arquivos, executa o modelo FHE sobre eles e retorna a predição criptografada. O container de inferência inclui ainda os arquivos de boilerplate (nginx.conf, serve e wsgi.py) necessários para containers de inferência customizados no SageMaker. É importante aumentar o valor de timeout no nginx.conf para permitir que a avaliação FHE seja concluída.

O Dockerfile.inference instala nginx, Flask, gunicorn, sagemaker, concrete-ml e concrete-python sobre a imagem base python:3.12.

Deploy do endpoint

O endpoint é implantado com configuração de inferência assíncrona. O start_inference_endpoint.py configura o modelo com a URI da imagem ECR, o caminho do modelo no S3 e a role IAM do endpoint. A instância mínima recomendada é ml.m5.xlarge; para melhor desempenho, ml.m5.24xlarge. Atenção: endpoints geram cobranças contínuas até serem excluídos. Consulte a tabela de preços do Amazon SageMaker AI e lembre-se de deletar o endpoint ao final dos testes.

Cliente de inferência

O client.py abstrai todos os detalhes de FHE do usuário final. Seu fluxo é:

Baixar o client.zip do S3.
Gerar as chaves privada e de avaliação.
Criptografar a consulta e fazer upload da consulta cifrada e da chave de avaliação para o S3.
Enviar os URIs ao endpoint e aguardar a resposta assíncrona.
Baixar a predição criptografada e descriptografá-la localmente.

Uma observação de design: cliente e endpoint tratam consultas e respostas de forma diferente. A resposta criptografada é uma sequência de bytes que o SageMaker AI gerencia naturalmente. Já a consulta do cliente é uma estrutura JSON que precisa conter o URI da chave de avaliação — embutir a consulta criptografada nesse JSON exigiria codificação (como Base64), adicionando processamento desnecessário. Por isso, o código de exemplo gerencia o upload da consulta criptografada diretamente no S3.

Desempenho e considerações de segurança

FHE oferece proteção criptográfica robusta, mas com um custo de desempenho significativo. A sobrecarga pode chegar a 100.000x em relação à inferência em texto simples. Algumas formas de reduzir esse impacto:

Aumentar o número de vCPUs da instância.
Quantização: técnica de ML que reduz a precisão numérica usada na inferência. Como o tempo de execução do concrete-ml cresce com a precisão, a quantização tem impacto ainda maior aqui do que em ML convencional. A quantização pode reduzir a acurácia do modelo, mas não é afetada pela conversão para FHE. Nos testes descritos, a quantização no código do modelo reduziu a sobrecarga para 2800x (de 67ms para 187s em uma instância ml.m5.xlarge) sem perda observável de acurácia. Aumentando os vCPUs, é possível chegar a 500x (46s em uma ml.m5.24xlarge).

Por isso, FHE ainda não é adequado para aplicações interativas sensíveis à latência. No entanto, é prático para cargas de trabalho assíncronas ou em lote, onde os requisitos de privacidade superam as preocupações com latência — exatamente o perfil dos três casos de uso apresentados.

Limitações de segurança importantes

É fundamental que os clientes mantenham em segredo as consultas descriptografadas e as predições, pois a combinação de um texto cifrado com seu texto simples pode revelar informações sobre a chave secreta.
Esta solução não protege o sigilo do modelo. As consultas e respostas ficam opacas ao SageMaker AI, mas o modelo em si pode ser visível ao serviço. O modelo também pode ser vulnerável a ataques de “roubo de modelo” por quem tiver acesso a consultas e respostas em texto simples.
O concrete-ml não fornece privacidade de circuito: informações sobre o modelo podem ser reveladas pelos textos cifrados.

A segurança é uma responsabilidade compartilhada entre a AWS e cada cliente. As boas práticas incluem: aplicar o princípio do menor privilégio nas roles IAM; habilitar criptografia padrão nos buckets S3 para valores que não sejam textos cifrados FHE; e restringir permissões dos buckets ao mínimo necessário.

Limpeza de recursos

Para evitar cobranças desnecessárias após os testes, é necessário excluir: o endpoint de inferência, a configuração do endpoint, o modelo no SageMaker AI, os artefatos no S3 (modelo, consultas cifradas, respostas cifradas e chaves de avaliação) e as imagens de container no ECR.

Problemas comuns

TimeoutError durante inferência: aumente o max_attempts no WaiterConfig ou use um tipo de instância maior.
Erros de AccessDenied: verifique se as roles IAM têm as permissões corretas para S3 e SageMaker AI.
Falhas na construção do container: verifique se o Docker tem memória suficiente (acima de 8 GB).
Erros de servidor durante inferência: verifique a paridade de versões dos pacotes concrete-ml.

Conclusão

A combinação de SageMaker AI com FHE via concrete-ml permite realizar inferência sobre dados que permanecem criptografados do início ao fim do processo. Essa abordagem une a agilidade, escala e infraestrutura gerenciada do SageMaker AI com proteção criptográfica que vai da consulta até a resposta. Para cenários onde a privacidade dos dados é inegociável — saúde, energia, telecomunicações — essa arquitetura oferece um caminho concreto e implementável.

Para aprofundamento, a AWS disponibiliza recursos adicionais:

Fonte

End-to-end encrypted ML inference with Amazon SageMaker AI and FHE (https://aws.amazon.com/blogs/machine-learning/end-to-end-encrypted-ml-inference-with-amazon-sagemaker-ai-and-fhe/)