Como a Snorkel AI obteve mais de 40% de economia de custos ao escalar workloads de machine learning usando o Amazon EKS

As startups de machine learning (ML) costumam ser grandes usuárias de computação, tendo em vista que treinam grandes modelos usando GPUs de ponta e os implantam em escala para inferência. O AWS Startups tem parceria com startups desde a concepção até o IPO, e já ajudou milhares de fundadores e inovadores de inteligência artificial (IA) a desenvolver seus negócios no Amazon Elastic Kubernetes Service (Amazon EKS). O Amazon EKS é uma escolha popular para criar e hospedar modelos de ML porque oferece a flexibilidade do Kubernetes com a segurança e a resiliência de ser um serviço gerenciado da AWS otimizado para criar workloads em contêineres altamente disponíveis.

A Snorkel AI é uma dessas empresas que se beneficia do Amazon EKS. A Snorkel AI equipa as empresas da Fortune 500, agências federais e inovadores de IA para criar, adaptar e destilar modelos de base (FMs) e grandes modelos de linguagem (LLMs) para um desempenho de alta precisão em conjuntos de dados específicos do domínio. Usando a abordagem centrada em dados da Snorkel para o desenvolvimento de IA, as organizações criaram serviços de IA prontos para produção destinados a casos de uso, incluindo processamento de sinistros de seguros, divulgação financeira, análise de estudos clínicos e aceleração do gerenciamento proativo de poços para perfuração offshore.

Nos últimos meses, a equipe da Snorkel trabalhou arduamente para lidar com os desafios exclusivos de projetar uma infraestrutura eficiente para dar suporte às workloads de desenvolvimento de ML sem aumentar as despesas com infraestrutura, diminuir a velocidade dos desenvolvedores ou prejudicar a experiência do usuário. Seu objetivo final era reduzir em mais de 40% os gastos com computação em cluster para o Snorkel Flow, sua plataforma de ML de ponta a ponta.

Uma visão geral do Snorkel Flow

A plataforma Snorkel Flow de desenvolvimento de dados de IA da Snorkel permite que as equipes de dados criem rapidamente aplicações de IA empregando um ciclo iterativo de rotulagem programática, treinamento rápido de modelos e análise de erros. Cada projeto começa quando os usuários criam um pequeno número de funções de rotulagem.

As funções de rotulagem empregam uma heurística simples, bancos de dados externos, modelos legados ou até mesmo chamadas para grandes modelos de linguagem para aplicar rótulos a faixas de dados não rotulados com base na intuição especializada codificada. O algoritmo de supervisão fraca da plataforma combina essas funções baseadas em regras para determinar o rótulo mais provável para cada registro. Em seguida, os usuários treinam um modelo simples com base nesses pontos de dados probabilísticos e avaliam o impacto de cada função de rotulagem. Na fase de análise, os usuários investigam as partes dos dados em que o modelo tem desempenho inferior. Em seguida, eles criam ou modificam as funções de rotulagem, treinam outro modelo rápido e continuam o ciclo. Quando os usuários estão satisfeitos com a qualidade de seus rótulos, eles criam um modelo final em uma arquitetura de model zoo,que varia de regressão logística a FMs, e o exportam para implantação.

Devido à natureza desse fluxo de trabalho, a infraestrutura do Snorkel Flow passa por períodos variados de alto uso de computação. Os custos operacionais aumentaram naturalmente à medida que a base de clientes e os recursos do produto de ML do Snorkel Flow aumentaram. Para obter um crescimento eficiente, a Snorkel buscou entender como aumentar as margens enquanto operava um software de ML de última geração. A Snorkel implementou as práticas a seguir para obter uma redução de mais de 40% nos custos de computação em cluster.

Soluções para a otimização de custos na nuvem 

As startups de software como serviço (SaaS) geralmente têm oportunidades de otimizar seus gastos com a nuvem. É essencial entender os fatores exclusivos que geram esses custos.

No caso da Snorkel, houve dois fatores significativos:

1. As workloads de desenvolvimento de ML geralmente exigem hardware especializado e caro, como GPUs. Normalmente, essas workloads são de natureza “intermitente”.
2. Empresas da Fortune 500 e grandes agências federais usam a Snorkel, incluindo grandes instituições financeiras com departamentos de TI sofisticados que têm requisitos específicos de implementação e privacidade de dados, usando uma plataforma em contêiner.

A equipe da Snorkel está interessada em criar sistemas que permitam um escalonamento eficiente sem um aumento linear nos custos de infraestrutura. Consequentemente, a Snorkel desenvolveu uma solução abrangente de escalonamento automático adaptada às suas workloads de ML no Amazon EKS para atender às preocupações com as despesas da nuvem. Essa solução não só agilizou as workloads que requerem computação intermitente, mas também atingiu suas metas de redução de custos.

Além da solução de escalonamento automático, as principais estratégias que contribuíram para a redução de mais de 40% nas despesas com a nuvem incluem:

• Colaboração com os líderes de engenharia e a equipe da AWS para adotar Savings Plans por meio de otimizações de configuração da nuvem.
• O dimensionamento adequado dos recursos por meio do monitoramento da utilização dos nós com o Prometheus e da consulta aos engenheiros de back-end para avaliar as necessidades dos componentes da plataforma.
• A mudança para tipos de máquinas virtuais (VMs) econômicas no Amazon EKS e a utilização de instâncias multi-GPU do Amazon Elastic Compute Cloud (Amazon EC2) para melhorar o desempenho de preço.
• Instituição de modificações no processo interno em que os engenheiros colaboravam com as equipes voltadas para o cliente para minimizar a computação ociosa.

Nesta publicação, a Snorkel compartilha o processo de abordagem desses desafios de escalonamento para ajudar a facilitar o projeto de uma infraestrutura melhor para sistemas de ML. Se você não conhece o Kubernetes, leia a postagem Introduction to Kubernetes da Snorkel, para saber mais sobre os conceitos básicos, e a postagem Machine learning on Kubernetes: wisdom learned at Snorkel, para saber mais sobre sua jornada com o Kubernetes até o momento.

Como funciona o Snorkel Flow na AWS?

Na prática, a interação do Snorkel Flow com a AWS segue a sequência descrita a seguir. Como a plataforma do Snorkel Flow depende muito de contêineres, a migração para a AWS foi praticamente perfeita.

• Os usuários acessam sua instância do Snorkel Flow por meio de um navegador da Web, que mapeia para uma regra no Amazon Route 53.
• O Route 53 encaminha a solicitação para um Application Load Balancer.
• O Application Load Balancer encaminha a solicitação para os pods do Snorkel Flow em execução em um cluster do EKS compartilhado. A Snorkel trocou o tipo de instância do EC2 de m5 para m6a para otimizar os custos, resultando em uma economia de 10% na computação, com impacto insignificante no desempenho com base no custo por hora para a mesma CPU e RAM.
• Além disso, a empresa fez o upgrade de uma única instância g4dn.8xlarge de GPU para uma instância g4dn.12xlarge de multi-GPU, o que permitiu atender quatro vezes mais pods de GPU.
• Cada instância do Snorkel Flow usa um volume do Amazon Elastic File System (Amazon EFS) para armazenar arquivos no disco.
• Uma fila auto-hospedada do Redis em um pod na instância do EC2 armazena os trabalhos recebidos, aguardando que os pods de trabalho os coletem.
• As métricas do EKS são enviadas para o Amazon CloudWatch, e os scripts personalizados monitoram os logs em busca de anomalias no desempenho do cluster.

Essa arquitetura produziu uma experiência estável e ágil para os usuários do Snorkel Flow.

Repensar como escalar

Antes da arquitetura descrita na Figura 1, as primeiras iterações da infraestrutura da Snorkel usavam recursos fixos. Os usuários da Snorkel compartilhavam que essas workloads intermitentes podiam levar muito tempo para serem concluídos e, portanto, afetavam negativamente sua experiência.

O escalonamento manual dos recursos de computação não era escalonável e era propenso a erros, levando a custos de nuvem que permaneciam elevados, mesmo durante períodos de baixa utilização. Era o pior dos dois mundos: baixa eficiência de custo da nuvem e desempenho mais lento do que o necessário.

Para lidar com esses desafios, a Snorkel implementou o escalonamento automático em vários níveis de sua infraestrutura, conforme discutido nas seções a seguir.

Projetar uma infraestrutura escalável levando em consideração a eficiência de custos

A distribuição do Kubernetes do Snorkel Flow envolve um conjunto de implantações executadas em um cluster do EKS que contém pods que executam vários componentes da plataforma.

Conforme mostrado na Figura 2, para lidar com os desafios exclusivos de trabalhar com workloads de computação intermitentes, a equipe da Snorkel introduziu um novo conceito para os pods do Kubernetes: categorizá-los semanticamente como “fixos” ou “flexíveis”.

• Os pods que são fixos não podem ser movidos com segurança de um nó para outro, seja porque perderão um estado importante na memória (como trabalhos de computação em andamento sem pontos de verificação), seja para minimizar o tempo de inatividade evitável dos componentes fundamentais da plataforma (por exemplo, o orquestrador do cluster Ray).
• Os pods que são flexíveis podem ser movidos com segurança para um novo nó. Essa distinção é significativa no contexto do escalonamento automático, porque os nós de redução de escala vertical envolvem a transferência de pods para fora dos nós subutilizados quando eles são encerrados.

Essa estrutura fixa/flexível oferece à Snorkel um meio específico de domínio para habilitar a redução de escala vertical automatizada do cluster, o que permite que eles ativem o escalonamento automático do cluster no Amazon EKS sem que o departamento financeiro envie mensagens a todo momento.

A abordagem inicial da Snorkel foi implantar podDisruptionBudgets no cluster do EKS para evitar que o escalonamento automático do cluster movesse pods flexíveis no decorrer do dia e não movesse pods fixos. Embora eficaz, essa abordagem deixou a equipe da Snorkel insatisfeita com o fato de a redução de escala vertical dos nós ser muito menor do que o teoricamente ideal.

Para resolver esse problema, a Snorkel colocou em camadas uma otimização de agendamento de pods que isolou os pods fixos em um pequeno grupo fixo de nós. Tal programação agendou pods flexíveis e de processamento (que são considerados pods fixos, mas são efêmeros devido ao escalonamento automático do nó de processamento) no grupo de nós flexíveis restantes.

Essas alterações permitiram que a Snorkel reduzisse a escala verticalmente dos nós flexíveis de forma eficiente à noite, quando se tornou seguro mover os pods flexíveis e reduzir a escala verticalmente da grande maioria dos pods de processamento.

Ao permitir a redução de escala vertical eficiente da grande maioria dos nós do cluster (ou seja, os nós flexíveis), a Snorkel conseguiu atingir sua meta de reduzir os custos da nuvem em mais de 40% para hospedar o Snorkel Flow.

Mais detalhes sobre a solução de escalonamento automático da Snorkel

A Snorkel divide a implementação da solução descrita na seção anterior em três esforços sequenciais:

1. Primeiro, a Snorkel implementou o “escalonamento automático de processamento”, um serviço personalizado de escalonamento automático baseado em Redis que permite que os pods de processamento aumentem e diminuam a escala verticalmente com base nos trabalhos nas filas de processamento.
2. Segundo, a empresa implementou o “escalonamento automático do cluster” reconfigurando suas implantações do Kubernetes para permitir que o escalonador automático do cluster do Kubernetes reduzisse a escala verticalmente dos nós, além de aumentar sua escala verticalmente.
3. Terceiro, a Snorkel implementou “otimizações de redução de escala vertical de nós” agrupando pods fixos em um pequeno grupo de nós fixos para evitar que pods fixos interferissem na redução de escala vertical dos nós restantes.

Escalonamento automático de processamento

A plataforma Snorkel Flow abstrai a computação em um paradigma em que os trabalhos aguardam nas filas do Redis e os processamentos são executados como processos em pods de processamento.

A Snorkel implementou uma solução de escalonamento automático de processamento (Figura 3) para pods de processamento executando uma função recorrente na API de back-end do Snorkel Flow. Em intervalos de poucos segundos, essa função verifica o cluster do Kubernetes e o Redis quanto à elegibilidade para redução e aumento da escala vertical.

Se houver trabalhos aguardando em uma ou mais filas relevantes baseadas no Redis, a função solicitará que a API do Kubernetes provisione pods de processamentor adicionais para processar esses trabalhos. Se a fila do Redis estiver vazia e não houver trabalhos em execução no registro de trabalhos, ela solicitará à API do Kubernetes que destrua os pods de processamento para liberar os recursos reservados de CPU e RAM.

Conforme mostrado na Figura 4, com essa implementação de escalonamento automático de processamento, os pods de processamento do Snorkel Flow tornaram-se efêmeros, aparecendo no cluster somente quando os trabalhos precisam ser processados.

Escalonamento automático do cluster

O recurso PodDisruptionBudget protege determinados pods contra interrupções (por exemplo, reinicializações voluntárias), permitindo a especificação do número máximo de réplicas de pod que podem estar indisponíveis em um determinado momento. Conforme mostrado na Figura 5, definir esse valor explicitamente como 0 para uma implantação garante que o escalonamento automático do cluster não reduza a escala verticalmente dos nós que executam os pods da implantação.

A implementação desse recurso em instâncias hospedadas do Snorkel Flow permitiu com segurança que o escalonamento automático do cluster reduzisse a escala verticalmente dos nós subutilizados. No entanto, a economia de custos obtida pela Snorkel foi marginal. A empresa ainda não conseguiu reduzir a escala verticalmente da maioria de seus nós porque todos os pods do Snorkel Flow estavam protegidos por um podDisruptionBudget associado.

Após uma análise mais detalhada, a equipe da Snorkel percebeu que essa proteção não precisa existir o tempo todo. As workloads são intermitentes, e a maior parte da interação do usuário com o Snorkel Flow ocorre durante o dia útil do cliente, o que significa que é seguro flexibilizar essa proteção fora do horário comercial. De forma semelhante ao escalonamento automático de processamento, a Snorkel implementou uma função recorrente que desativou os podDisruptionBudgets durante a noite para os pods flexíveis de uma instância, definindo o número máximo de réplicas de pods indisponíveis para 1, em vez de 0 (Figura 6). A solução anterior de escalonamento automático de processamento, combinada com o ClusterAutoscaler e o recurso PodDisruptionBudget, foi capaz de reduzir muito mais nós de processamento subutilizados do que antes. Os clientes que implantam o Snorkel Flow em sua nuvem podem configurar isso conforme necessário.

Otimizações de redução de escala vertical de nós

Mesmo com essas melhorias, a Snorkel percebeu que a maioria dos nós subutilizados não estava passando por uma redução de escala vertical.

Após uma investigação mais aprofundada, a Snorkel percebeu que o problema era decorrente do fato de pods fixos e flexíveis ocuparem o mesmo nó. Esse fato era problemático porque um pod fixo atribuído de forma pseudoaleatória a um nó que continha pods flexíveis “fixaria” esse nó e impediria que ele passasse por uma redução de escala vertical, mesmo quando estivesse subutilizado. Essa falta de controle sobre o agendamento de pods fixos levou a períodos em que a grande maioria dos nós do cluster não pôde ser reduzida, embora representassem muito mais capacidade de computação do que o necessário naquele momento.

A Snorkel aproveitou o recurso podAffinities do Kubernetes para resolver esse problema, o que permitiu restringir em quais nós um pod está qualificado para ser executado com base nos rótulos de outros pods já em execução em um determinado nó. A empresa adicionou rótulos aos pods para diferenciar entre pods fixos e flexíveis e adicionou uma divisão podAntiAffinity à configuração de suas implantações para garantir que os pods fixos não sejam agendados em nós que executam pods flexíveis, e vice-versa.

Essa implementação de podAffinities permitiu que a Snorkel AI dividisse os nós em dois grupos funcionais: o grupo de nós fixos contendo pods fixos, que nunca podem ser transferidos com segurança entre os nós (por exemplo, Redis devido ao cache), e o grupo de nós flexíveis contendo pods “flexíveis” que são efêmeros (como pods de processamento) ou seguros para serem transferidos fora do horário comercial (por exemplo, durante a noite).

Embora isso seja possível com a intervenção manual durante a manutenção da plataforma, a Snorkel não pode reduzir automaticamente a escala vertical dos nós fixos. Essa solução, no entanto, permite que a empresa reduza automaticamente a escala vertical dos nós flexíveis porque agora ela isolou os pods imóveis nos nós fixos.

Conclusão

As equipes da Snorkel e do AWS Startups esperam que o compartilhamento desse processo de pensamento e dessas soluções ajude outras startups a criar uma infraestrutura melhor para workloads de ML, que estão rapidamente se tornando mais importantes à medida que o ML, os grandes modelos de linguagem e outros FMs entram em produção para organizações em todo o mundo.

Você está participando do AWS re:Invent 2023? Essa implantação será destacada como parte da sessão da Snorkel Navigating the future of AI: Deploying generative models on Amazon EKS (sessão CON312). Não deixe de conferir.

Agradecemos a David Hao, Edmond Liu e Alec Xiang por ajudarem a tornar essa visão técnica uma realidade para a Snorkel. Agradecimentos especiais às pessoas mencionadas acima, bem como a Matt Casey, Henry Ehrenberg, Anthony Bishopric e a toda a equipe de engenharia de infraestrutura da Snorkel por seu feedback atencioso sobre este artigo.