Depurador interativo de código aberto do LinkedIn para pipelines de IA K8s

O Kubernetes é cada vez mais popular como plataforma para a construção de projetos de aprendizado de máquina, mas a experiência do desenvolvedor no Kubernetes continua desafiadora, muitas vezes exigindo mais conhecimento em infraestrutura do que muitos programadores estão interessados ​​em adquirir.

E apesar da promessa de que os contêineres encapsulam aplicativos e suas dependências significam um ambiente portátil e consistente durante todo o ciclo de desenvolvimento, isso simplesmente não é prático para os modelos maiores, como aqueles usados ​​em IA generativa, onde nem o conjunto de dados nem o hardware da GPU estão disponíveis para desenvolvedores que trabalham localmente. .

Para melhorar a experiência do desenvolvedor, o LinkedIn criou o FlyteInteractive, que fornece ao usuário “um ambiente de desenvolvimento interativo que permite executar diretamente o código, conectar-se ao Microsoft VSCode Server dentro de pods Kubernetes com acesso a recursos e dados em grande escala na grade. É como fazer SSH na porta da GPU e desenvolver diretamente a partir daí. Tudo é exatamente igual”, explicou Jason Zhu, engenheiro de aprendizado de máquina do LinkedIn que ajudou a criar o software.

Em vez de escrever um conjunto de dados simulado para usar com seu modelo, os desenvolvedores têm acesso ao conjunto de dados real que está no cluster usando o suporte de desenvolvimento remoto do VSCode, o que evita perda de tempo em um modelo que não consegue lidar com o conjunto de dados em tamanho real. “À medida que avançamos em direção a arquiteturas maiores e mais complexas. É quase impossível desenvolver o código localmente e testá-lo”, explicou.

“Os recursos disponíveis para o desenvolvimento local não incluem as mesmas GPUs sofisticadas e caras usadas na produção, a mesma quantidade de memória – ou as complexidades de um sistema distribuído. Você pode comprometer o tamanho e a complexidade do modelo (para executá-lo localmente), mas isso também comprometerá a chance de que, uma vez carregado o modelo para dados de produção reais, ele seja bem-sucedido.”

Voo antecipado

Como o nome sugere, FlyteInteractive é construído a partir de outro projeto de código aberto, Flyte.

Originalmente desenvolvido e de código aberto pela Lyft, Flyte é um orquestrador de fluxo de trabalho para Kubernetes escrito em Go, projetado especificamente para pipelines de dados e aprendizado de máquina, com uma interface que permite aos desenvolvedores criar seus fluxos de trabalho na linguagem mais popular para desenvolvedores de aprendizado de máquina: Python, com verificação de tipo forte para que mais bugs sejam detectados em tempo de compilação (o que pode economizar dinheiro e tempo, dada a infraestrutura cara necessária para aprendizado de máquina).

Flyte se formou na LF AI & Data Foundation no início de 2022 e já está em uso na HBO, Intel, Spotify — e no LinkedIn, que usa IA extensivamente e já migrou todas as suas cargas de trabalho LLM, bem como algumas cargas de trabalho tradicionais de aprendizado de máquina. Flyte cobre mais cenários do que Kubeflow e não exige tanto conhecimento em Kubernetes dos desenvolvedores (mas também tem integrações Kubeflow para pacotes populares como PyTorch e TensorFlow).

Uma grande parte do apelo para as grandes organizações é a escalabilidade que oferece, de acordo com Byron Hsu, colaborador do projeto Flyte que trabalha no dimensionamento da infraestrutura de aprendizado de máquina no LinkedIn. “Temos mais de mil fluxos de trabalho todos os dias e precisamos garantir que cada fluxo de trabalho possa ser iniciado rapidamente.”

Flyte também ajuda com o tipo de experimentação rápida que é tão importante para o aprendizado de máquina, onde os conjuntos de dados mudam e novos algoritmos são lançados com frequência. “O tempo de agendamento é super, super rápido para que os usuários possam fazer experimentos rapidamente”, disse Hsu ao New Stack.

A interface Python também torna o Flyte fácil de aprender para os desenvolvedores de aprendizado de máquina: “Se você deseja adicionar uma tarefa Python personalizada aos seus fluxos de trabalho, é intuitivo e fácil no Flyte. Definitivamente torna os desenvolvedores de aprendizado de máquina muito mais rápidos.”

Flyte também traz alguns recursos DevOps familiares que aceleram o desenvolvimento de aprendizado de máquina, explicou Zhu, que trabalha com modelos extremamente grandes como aquele que alimenta o feed da página inicial do LinkedIn. “Anteriormente, toda vez que construímos nosso pipeline, tínhamos que extrair as dependências localmente e esperar que isso (acontecesse). Mas como o Flyte é baseado em imagem, podemos incluir todas essas dependências na imagem com antecedência, de modo que leva apenas alguns segundos para um usuário carregar seu trabalho e o processo de inserção de todas essas dependências acontece em tempo de execução.” Isso economiza uma quantidade significativa de tempo, inclusive sempre que você atualiza um fluxo de trabalho e executa seu trabalho de aprendizado de máquina novamente.

Para incentivar a reutilização de código e evitar que cada equipe reconstrua os mesmos componentes e processos para cada novo projeto, o LinkedIn criou um Component Hub além do Flyte, que já possui mais de 20 componentes reutilizáveis ​​que economizam muito trabalho repetitivo. “É para ferramentas comuns, como pré-processamento de dados, treinamento ou inferência”, explicou Hsu. “A equipe de treinamento pode criar um componente de treinamento como um treinador do TensorFlow e todos os engenheiros de ML vinculados podem usá-lo sem reimplementá-lo.”

Isso também torna técnicas mais poderosas e complexas, como a quantização de modelo em que Zhu tem trabalhado recentemente, muito mais amplamente disponíveis, transformando-as em uma função ou chamada de API. Existem vários algoritmos para converter a representação de um modelo de alta para baixa precisão, para que você possa compactar modelos e servi-los usando o menor número possível de recursos e, normalmente, os desenvolvedores de aprendizado de máquina precisariam pesquisar os desenvolvimentos mais recentes, escolher um algoritmo e implementá-lo por conta própria projeto.

“Nós o construímos como um componente e como o Flyte tem o conceito de componentes reutilizáveis, para o pipeline de qualquer outro usuário, eles podem optar por chamá-lo como uma interface ou uma API externa. Assim, eles podem quantificar seu modelo depois que o modelo for treinado, não importa se é um modelo para resumo, um modelo para raciocínio ou um modelo para extração de entidades”, disse Zhu.

Os desenvolvedores podem explorar vários algoritmos rapidamente, pois basta conectá-los ao fluxo de trabalho para testar seu efeito no uso de recursos e na precisão do modelo.

“Se você não tem certeza se a quantização funcionará em seu caso de uso específico, podemos ter um hub centralizado com todos os diferentes algoritmos de quantização, para que você possa testar todos eles e observar a matriz de resultados e a latência para entender a negociação -offs e descobrir a abordagem correta. À medida que o campo evolui, mais algoritmos de quantização surgirão, então precisamos ter uma plataforma muito flexível onde possamos testar todos esses algoritmos e adicioná-los ao hub centralizado do qual todos os pipelines downstream possam se beneficiar”, disse Zhu.

Depuração interativa remota

Ser capaz de escrever seu pipeline mais rapidamente e reutilizar componentes acelera o desenvolvimento de aprendizado de máquina o suficiente para que os engenheiros de software do LinkedIn começassem a perceber outras coisas que estavam retardando seu fluxo de trabalho: tudo, desde ter que trabalhar com conjuntos de dados simulados menores que não funcionavam. combinar bem os conjuntos de dados de produção com o ambiente local de desenvolvimento e teste sem o hardware e os recursos do ambiente de produção, o que significa limites artificiais no tamanho dos modelos para o longo ciclo de depuração e ter que esperar que o código seja implantado antes de descobrir se realmente corrigiu o bug.

Graças às diferenças entre os ambientes local e de produção, apenas um em cada cinco bugs foi corrigido na primeira vez: com cada envio de código levando pelo menos 15 minutos para entrar em produção. Rastrear até mesmo um bug menor poderia exigir dezenas de tentativas: em um caso, demorou quase uma semana para encontrar e corrigir um problema.

Esses problemas não são exclusivos do desenvolvimento de aprendizado de máquina, mas são exacerbados não apenas pelo tamanho dos modelos de aprendizado de máquina e pelos conjuntos de dados em que trabalham e pela infraestrutura cara necessária para executar modelos em produção, mas também por um ecossistema que não Nem sempre oferecemos ferramentas que os desenvolvedores em outras áreas consideram certas, como inspeção de código e depuração remota.

Mesmo o menor modelo generativo de IA que seja razoável não pode ser executado em uma CPU, apontou Zhu. “Quando você chega a esse estágio, é realmente natural movermos o processo de codificação e depuração para o pod Kubernetes ou clusters de GPU com os dados reais e os mesmos recursos que você executaria na produção.”

FlyteInteractive pode carregar dados de armazenamento HDFS ou S3 e suporta trabalhos de nó único e configurações mais complexas de vários nós e multi-GPU.

Os desenvolvedores podem simplesmente adicionar o decorador VSCode ao seu código, conectar-se ao servidor VSCode e usar o comando Executar e depurar normalmente para obter uma sessão de depuração interativa que executa sua tarefa Flyte no VSCode. Flyte armazena em cache a saída do fluxo de trabalho para evitar a reexecução de tarefas caras, para que o VSCode possa carregá-la da tarefa anterior.

Você obtém todas as opções usuais, como definir pontos de interrupção (mesmo em um processo de treinamento distribuído) ou executar scripts locais, bem como ferramentas de navegação e inspeção de código que facilitam a compreensão da complexa estrutura de código de um modelo grande com vários módulos e ver como os dados fluem para o modelo.

Você também pode configurar o plug-in para ser executado automaticamente se uma tarefa Flyte falhar, o que interrompe o encerramento da tarefa e lhe dá a chance de inspecionar e depurar a partir do ponto de falha. Depois de descobrir qual é o problema e reescrever seu código, você pode desligar o servidor VSCode e fazer com que Flyte continue executando o fluxo de trabalho. “Você pode retomar o fluxo de trabalho com o código alterado: basta clicar em um botão e a tarefa será executada com o novo código alterado e todo o fluxo de trabalho continuará”, explicou Hsu.

O suporte ao notebook Jupyter no FlyteInteractive também será útil, sugeriu ele: “É um orquestrador rápido com capacidade de notebooks Jupyter e depuração interativa, para que você possa usá-lo para experimentar rapidamente e também para um trabalho agendado ou em lote”.

Embora atualmente seja um plugin, agora é de código aberto, ele espera que, com a contribuição da comunidade, se torne um recurso integrado no Flyte.

Explorando recursos e código

FlyteInteractive já economizou milhares de horas de codificação e depuração no LinkedIn; também pode ajudar no controle de custos com sua opção de monitoramento de recursos. “Se um pod ficar inativo por um determinado período de tempo, vamos apenas limpá-lo e enviar um e-mail para notificar o usuário ‘Ei, seu pod está inativo há algum tempo. Pense em liberar o recurso ou fazer algo para tomar alguma ação sobre ele’.” No futuro, Hsu nos disse que será mais refinado. “Por exemplo, queremos detectar a utilização da GPU. Se eles ocuparam uma GPU, mas na verdade não a usam, podemos querer eliminá-la depois de, digamos, dez minutos, para termos melhor controle de orçamento para nosso sistema de GPU.”

Isso depende do suporte de checkpoint do Flyte porque fazer checkpoints é caro e geralmente não é uma boa opção para os loops de treinamento iterativos usados ​​no aprendizado de máquina. “Temos que fornecer um bom ponto de verificação para que, quando o trabalho do usuário for antecipado, o modelo também seja salvo.”

Mas para os desenvolvedores, o recurso mais atraente não é nem mesmo a depuração rápida, sugeriu Zhu. “Gosto do recurso de inspeção de código porque me permite compreender rapidamente o mecanismo interno de funcionamento dos algoritmos e também me ajuda a criar novas abordagens.”

Isso não é útil apenas para o seu próprio código, ressaltou. “Os engenheiros não só podem aplicar isso aos seus repositórios internos, mas também podem aplicar isso aos repositórios de código aberto. Como campo, o ML é super rápido: toda semana surgem novos algoritmos que engenheiros como nós precisam testar. Podemos simplesmente apontar esta ferramenta para um repositório de código aberto e entender rapidamente se é uma técnica que queremos usar.”

A postagem Depurador interativo de código aberto do LinkedIn para K8s AI Pipelines apareceu pela primeira vez em The New Stack.