Como o Amazon Search executa projetos de machine learning resilientes e em grande escala com o Amazon SageMaker

Se você pesquisou um item para comprar no amazon.com, você usou os serviços de pesquisa da Amazon. Na Amazon Search, somos responsáveis ​​pela experiência de pesquisa e descoberta de nossos clientes em todo o mundo. Em segundo plano, indexamos nosso catálogo mundial de produtos, implantamos frotas altamente escaláveis ​​da AWS e usamos machine learning (ML) avançado para corresponder produtos relevantes e interessantes às consultas de cada cliente.

Nossos cientistas treinam regularmente milhares de modelos de ML para melhorar a qualidade dos resultados da pesquisa. Apoiar a experimentação em larga escala apresenta seus próprios desafios, especialmente quando se trata de melhorar a produtividade dos cientistas que treinam esses modelos de ML.

Neste post, compartilhamos como construímos um sistema de gestão em torno de Amazon Sage Maker trabalhos de treinamento, permitindo que nossos cientistas disparem e esqueçam milhares de experimentos e sejam notificados quando necessário. Eles agora podem se concentrar em tarefas de alto valor e resolver erros algorítmicos, economizando 60% do tempo.

O desafio

Na Amazon Search, nossos cientistas resolvem problemas de recuperação de informações experimentando e executando vários trabalhos de treinamento de modelo de ML no SageMaker. Para acompanhar a inovação de nossa equipe, a complexidade de nossos modelos e o número de trabalhos de treinamento aumentaram ao longo do tempo. Os trabalhos de treinamento do SageMaker nos permitem reduzir o tempo e o custo para treinar e ajustar esses modelos em escala, sem a necessidade de gerenciar a infraestrutura.

Como tudo em projetos de ML de grande escala, os trabalhos de treinamento podem falhar devido a vários fatores. Esta postagem se concentra em escassez de capacidade e falhas devido a erros de algoritmo.

Projetamos uma arquitetura com um sistema de gerenciamento de tarefas para tolerar e reduzir a probabilidade de falha de uma tarefa devido à indisponibilidade de capacidade ou erros de algoritmo. Ele permite que os cientistas demitam e esqueçam milhares de trabalhos de treinamento, tentem novamente automaticamente em falhas transitórias e sejam notificados sobre sucesso ou falha, se necessário.

Visão geral da solução

No diagrama de solução a seguir, usamos os trabalhos de treinamento do SageMaker como a unidade básica de nossa solução. Ou seja, um trabalho representa o treinamento de ponta a ponta de um modelo de ML.

O fluxo de trabalho de alto nível desta solução é o seguinte:

1. Os cientistas invocam uma API para enviar um novo trabalho ao sistema.
2. O trabalho é registrado no New status em um repositório de metadados.
3. Um agendador de tarefas recupera de forma assíncrona New jobs do repositório de metadados, analisa sua entrada e tenta iniciar jobs de treinamento do SageMaker para cada um. Seu status muda para Launched or Failed dependendo do sucesso.
4. Um monitor verifica o progresso dos trabalhos em intervalos regulares e relata suas Completed, Failedou InProgress estado no repositório de metadados.
5. Um notificador é acionado para relatar Completed e Failed empregos para os cientistas.

A persistência do histórico de trabalhos no repositório de metadados também permite que nossa equipe realize análises de tendências e monitore o progresso do projeto.

Esta solução de agendamento de tarefas usa componentes serverless fracamente acoplados com base em AWS Lambda, Amazon DynamoDB, Serviço de notificação simples da Amazon (Amazon SNS), e Amazon Event Bridge. Isso garante escalabilidade horizontal, permitindo que nossos cientistas lancem milhares de trabalhos com o mínimo de esforço operacional. O diagrama a seguir ilustra a arquitetura sem servidor.

Nas seções a seguir, entraremos em mais detalhes sobre cada serviço e seus componentes.

DynamoDB como repositório de metadados para execuções de trabalho

A facilidade de uso e a escalabilidade do DynamoDB tornaram uma escolha natural manter os metadados de trabalhos em uma tabela do DynamoDB. Essa solução armazena vários atributos de trabalhos enviados por cientistas, ajudando assim no acompanhamento do progresso e na orquestração do fluxo de trabalho. Os atributos mais importantes são os seguintes:

• ID de trabalho – Um ID de trabalho exclusivo. Isso pode ser gerado automaticamente ou fornecido pelo cientista.
• Status de trabalho – O status do trabalho.
• TrabalhoArgs – Outros argumentos necessários para criar um trabalho de treinamento, como o caminho de entrada no Amazon S3, o URI da imagem de treinamento e muito mais. Para obter uma lista completa dos parâmetros necessários para criar um trabalho de treinamento, consulte Criar trabalho de treinamento.

Lambda para a lógica principal

Nós usamos três baseado em contêiner Funções do Lambda para orquestrar o fluxo de trabalho do trabalho:

• Enviar trabalho – Esta função é invocada por cientistas quando eles precisam lançar novos trabalhos. Ele atua como uma API para simplificar. Você também pode enfrentá-lo com Gateway de API da Amazon, se necessário. Essa função registra os trabalhos na tabela do DynamoDB.
• Iniciar trabalhos – Esta função recupera periodicamente New jobs da tabela do DynamoDB e os inicia usando o SageMaker Criar trabalho de treinamento comando. Ele tenta novamente em falhas transitórias, como ResourceLimitExceeded e CapacityError, para instrumentar a resiliência no sistema. Em seguida, ele atualiza o status do trabalho como Launched or Failed dependendo do sucesso.
• Monitorar trabalhos – Esta função monitora periodicamente o progresso do trabalho usando o DescreverTrainingJob comando e atualiza a tabela do DynamoDB de acordo. Ele sonda Failed jobs dos metadados e avalia se eles devem ser reenviados ou marcados como com falha terminal. Ele também publica mensagens de notificação para os cientistas quando seus trabalhos atingem um estado terminal.

EventBridge para agendamento

Usamos o EventBridge para executar as funções Launch Jobs e Monitor Jobs Lambda em uma programação. Para obter mais informações, consulte Tutorial: Agendar funções do AWS Lambda usando EventBridge.

Alternativamente, você pode usar Streams do Amazon DynamoDB para os gatilhos. Para mais informações, veja Streams do DynamoDB e gatilhos do AWS Lambda.

Notificações com Amazon SNS

Nossos cientistas estão notificado por e-mail usando o Amazon SNS quando seus trabalhos atingem um estado terminal (Failed após um número máximo de tentativas), Completedou Stopped.

Conclusão

Nesta postagem, compartilhamos como o Amazon Search adiciona resiliência às cargas de trabalho de treinamento do modelo de ML, agendando-as e repetindo-as em caso de falta de capacidade ou erros de algoritmo. Usamos as funções do Lambda em conjunto com uma tabela do DynamoDB como um repositório de metadados central para orquestrar todo o fluxo de trabalho.

Esse sistema de agendamento permite que os cientistas enviem seus trabalhos e os esqueçam. Isso economiza tempo e permite que eles se concentrem em escrever modelos melhores.

Para ir mais longe em seus aprendizados, você pode visitar Incrível Criador de Sábios e encontre em um único lugar todos os recursos relevantes e atualizados necessários para trabalhar com o SageMaker.

---

Sobre os autores

Luochao Wang é engenheiro de software na Amazon Search. Ele se concentra em sistemas distribuídos escaláveis ​​e ferramentas de automação na nuvem para acelerar o ritmo da inovação científica para aplicativos de aprendizado de máquina.

Ishan Bhatt é engenheiro de software na equipe Amazon Prime Video. Ele trabalha principalmente no espaço MLOps e tem experiência na criação de produtos MLOps nos últimos 4 anos usando o Amazon SageMaker.

Abhinandan Patni é engenheiro de software sênior na Amazon Search. Ele se concentra na construção de sistemas e ferramentas para treinamento de aprendizado profundo distribuído e escalável e inferência em tempo real.

Eiman Elnahrawy é engenheiro de software principal da Amazon Search, liderando os esforços de aceleração, dimensionamento e automação de aprendizado de máquina. Sua experiência abrange várias áreas, incluindo Machine Learning, Sistemas Distribuídos e Personalização.

Sofian Hamiti é um arquiteto de soluções especialista em IA/ML na AWS. Ele ajuda clientes de todos os setores a acelerar sua jornada de IA/ML, ajudando-os a criar e operacionalizar soluções de aprendizado de máquina de ponta a ponta.

Dra. Romi Datta  é gerente sênior de gerenciamento de produtos na equipe do Amazon SageMaker responsável pelo treinamento, processamento e armazenamento de recursos. Ele está na AWS há mais de 4 anos, ocupando vários cargos de liderança em gerenciamento de produtos no SageMaker, S3 e IoT. Antes da AWS, ele trabalhou em várias funções de gerenciamento de produtos, engenharia e liderança operacional na IBM, Texas Instruments e Nvidia. Ele tem um MS e Ph.D. em Engenharia Elétrica e de Computação pela Universidade do Texas em Austin e MBA pela Booth School of Business da Universidade de Chicago.

RJ é um engenheiro da equipe Search M5 que lidera os esforços para construir sistemas de aprendizado profundo em larga escala para treinamento e inferência. Fora do trabalho explora diferentes cozinhas de comida e pratica desportos de raquete.