Autor: Pedro Viana Gomes Pena

Revisado por Kevin Rossell

Os fluxos de trabalho de machine learning (ML) são a espinha dorsal da inovação baseada em dados, mas garantir que eles sejam escaláveis e adaptáveis pode ser uma tarefa desafiadora. É por isso que construímos o Common Python Workflows (CPW), uma solução versátil que simplifica como as equipes criam, gerenciam e executam pipelines de machine learning.

O CPW enfrenta problemas comuns, como a complexidade dos fluxos de trabalho e a falta de flexibilidade, capacitando as equipes a focarem em refinar seus modelos em vez de lutarem contra processos ineficientes. Neste post do blog, vamos explorar o raciocínio por trás do desenvolvimento do CPW, seus recursos diferenciados e os principais aprendizados da nossa jornada até agora.

Desafios Iniciais

Antes do CPW, enfrentávamos várias limitações com nossa ferramenta existente, Sheep:

  • Complexidade e confusão: Ao longo dos anos, Sheep tornou-se sobrecarregada com inúmeras funções, levando a uma experiência de usuário confusa. Escrever arquivos de configuração adequados para governar o comportamento do treinamento era igualmente desafiador, dificultando para os usuários entenderem seu pleno potencial.
  • Falta de extensibilidade: A Sheep foi projetada como uma única biblioteca Python com todo o código e dependências agrupadas. Isso tornava difícil estender ou substituir certas funcionalidades, como carregamento de dados. Os usuários tinham pouca flexibilidade para alterar a pipeline: podiam apenas substituir 3 etapas pré-definidas.
  • Suporte a frameworks: Dar suporte a vários frameworks (como Scikit-learn, Pytorch e Tensorflow) não era fácil, e distribuir o treinamento era trabalhoso. Os usuários tinham que treinar o modelo com fklearn.
  • Treinamento out-of-core: Os usuários não podiam realizar treinamento out-of-core de forma eficiente.
  • Predição e avaliação: Não havia abstração para predição, avaliação ou comparação de modelos.
  • Re-treinamento automático: Pipelines de re-treinamento automático não eram fornecidas por padrão.
  • Gestão de dependências: Mesmo mudanças menores exigiam a atualização de múltiplas dependências nos modelos, o que frequentemente levava ao inferno das dependências.

Esses desafios nos levaram a desenvolver o CPW—uma solução flexível e à prova de futuro para simplificar fluxos de trabalho de ML e resolver esses pontos problemáticos.

O que é o CPW?

O CPW é uma biblioteca projetada para simplificar a criação e gestão de fluxos de trabalho de ML. Ela introduz vários recursos chave:

  • Workflows de treinamento personalizados: Adapte os workflows para atender necessidades específicas.
  • Herdar e adaptar fluxos de treinamento: Reutilize workflows com mínimas alterações.
  • Compartilhamento de fluxos entre equipes: Promova a colaboração permitindo fácil compartilhamento de workflows.
  • Manipulação de serialização/deserialização: Simplifique a salvaguarda e o carregamento de modelos e dados.
  • Implementação de re-treinamento automático: Implemente mecanismos de re-treinamento automático facilmente.

Para executar pipelines de modelos, usamos o Kubeflow Pipelines—uma plataforma open-source de machine learning que facilita a implantação e gestão de piplines de ML ponta a ponta. Além de orquestrar pipelines de ML no Kubernetes, ele fornece ferramentas para gestão e monitoramento de execuções de pipeline.

O CPW integra-se perfeitamente com o Kubeflow através do Argo Workflows. Uma vez que um fluxo (Flow) é criado, ele pode ser compilado para um Argo Workflow e submetido para execução no Kubeflow, que oferece recursos como:

  • Executar tarefas em paralelo, cada uma em um contêiner diferente.
  • Ajustar hiperparâmetros com o Katib (Componente do Kubeflow para busca de hiperparâmetros).
  • Alocar diferentes recursos de máquina para cada etapa.

O CPW baseia-se no Dagster, aproveitando principalmente os conceitos de software defined assets e IOManagers.

Por que Escolhemos o Dagster

Ao desenhar o CPW, evitamos deliberadamente re-implementar soluções existentes para Grafos Acíclicos Direcionados (DAG) e orquestração. Havia vários frameworks open-source que faziam isso bem. Queríamos focar na experiência de autoria de pipelines.

Após avaliar frameworks open-source como Metaflow e Dagster, escolhemos o Dagster pelas seguintes características:

  • Design modular: O design modular do Dagster permite uma separação clara de responsabilidades. Cada etapa em um fluxo de trabalho pode ser definida de forma independente, e as dependências são inferidas com base nos artefatos produzidos por cada etapa. Isso contrasta com o Metaflow, que exigia que os usuários agrupassem o que uma etapa faria com a descrição de quais etapas seguiriam.
  • Ativos definidos por software: O conceito de ativos definidos por software do Dagster alinha-se bem com nosso objetivo de criar workflows reutilizáveis e flexíveis. Isso nos permite gerenciar os componentes e artefatos de nossos workflows de maneira modular e reutilizável.
  • IOManagers: Os IOManagers do Dagster oferecem uma maneira robusta de lidar com operações de entrada e saída, assegurando que os dados sejam gerenciados de maneira eficiente e consistente em diferentes etapas do workflow, liberando assim nosso usuário final desse fardo.
  • Extensibilidade: O design do Dagster permite uma maior e focada extensibilidade, permitindo aos usuários customizarem apenas o componente do qual precisam para atender suas necessidades específicas, como um IOManager para imagens, que pode ser além disso reutilizado e compartilhado com outras equipes.

Design da Classe de Workflow

Para facilitar a criação de workflows reutilizáveis, introduzimos uma classe FlowSpec. Esta classe permite que as equipes criem workflows que podem ser herdados e adaptados por outras equipes. Por exemplo, uma equipe pode herdar um workflow (um DAG) e substituir uma etapa (um nó) ao sobrescrever o atributo apropriado do workflow herdado.

Essa abordagem traz benefícios adicionais:

  • Integração com infraestrutura: Cria um ponto central para nossa equipe manter e evoluir. Uma vez que algo mude em nossa infraestrutura, podemos atualizar a classe base e ter certeza de que todos os workflows se adaptarão à mudança.
  • Facilidade de desenvolvimento: Desenvolver novas funcionalidades tornou-se mais fácil porque podemos contar com os conceitos modulares que tínhamos para Ativos, IOManagers, e Flows (repositórios no Dagster).

Como exemplo, desde que introduzimos o CPW, habilitamos as seguintes funcionalidades para os usuários:

  • Diferentes backends para execução: ao mudar uma palavra-chave na definição do ativo, os usuários podem executar essa etapa em clusters Spark no Databricks ou no Sagemaker.
  • Profiling: ao passar uma palavra-chave na linha de comando, cProfiling é habilitado para todas as etapas do Flow.

Isso permite mudanças minimamente intrusivas que requerem pouco ou nenhuma modificação no código do usuário para serem usadas.

Este design também facilitou a colaboração, com equipes desenvolvendo Flows básicos de treinamento a serem usados para múltiplos modelos, permitindo que melhorias sejam aplicadas uma vez em todos os modelos.

Conceitos Centrais

O CPW baseia-se em vários conceitos centrais do Dagster, principalmente software defined assets e IOManagers. Esses conceitos permitem definir e gerenciar os componentes dos workflows de forma modular e reutilizável.

Aqui está um exemplo de definição de Flow no CPW:

Neste exemplo, definimos um workflow simples com três ativos:

  • params: Um dicionário contendo definições estáticas necessárias ao longo do workflow, como hiperparâmetros.
  • my_dataset: Um dataset que é dividido com base nos parâmetros definidos em params.
  • train_model: Um modelo treinado usando o dataset de treino e parâmetros.

Com base nos nomes e argumentos dos ativos (funções), o CPW automaticamente estabelece as dependências entre as funções. O DAG resultante terá a seguinte aparência:

Artefatos e Persistência

O CPW assegura que os artefatos de cada ativo sejam persistidos em um diretório S3 compartilhado (ou local). Por exemplo:

  • params: Armazenado como um arquivo JSON.
  • my_dataset: Salvo como Parquet.
  • train_model: Serializado e salvo como um objeto Pickle.

O CPW lida com a serialização e deserialização desses artefatos, garantindo que eles sejam armazenados e recuperados de maneira eficiente. Essa persistência automática permite que os usuários se concentrem na definição do seu workflow sem se preocupar com a gestão subjacente dos dados.

Ao aproveitar os ativos definidos por software e os IOManagers do Dagster, o CPW fornece uma estrutura robusta e flexível para construir e gerenciar workflows de ML. Essa abordagem modular torna fácil definir, compartilhar e adaptar workflows, capacitando as equipes a trabalharem de forma mais eficiente e eficaz, focando apenas no que importa:

  • A função que seu Flow deve executar, e;
  • Os artefatos que seu Flow deve produzir.

Estudos de Caso

O CPW permitiu que múltiplas equipes simplificassem significativamente seus workflows:

Compartilhando um mesmo fluxo base para uma variedade de modelos de NLP

Uma das histórias de sucesso mais significativas vem de uma equipe trabalhando com modelos de processamento de linguagem natural (NLP). Compartilhando um Flow base comum para vários modelos de NLP, a equipe conseguiu reduzir o esforço para atualizar o workflow de treinamento dos modelos em 80%.

Isso foi alcançado reutilizando a mesma estrutura de workflow e apenas modificando parâmetros específicos para cada modelo. O design modular do CPW permitiu que a equipe gerenciasse e adaptasse seus workflows de maneira eficiente, resultando em economias substanciais de tempo e aumento da produtividade.

Distribuindo melhorias entre as equipes

Outra história de sucesso envolve um cientista de dados que desenvolveu ativos de avaliação e gráficos para seu modelo. Usando o CPW, ele conseguiu implantar imediatamente essas melhorias em todos os outros modelos dentro de sua equipe. Isso economizou semanas de esforço de engenharia e retrabalho, à medida que os novos ativos puderam ser integrados perfeitamente nos workflows existentes. A capacidade de compartilhar e reutilizar componentes de workflow entre diferentes modelos e equipes destaca o poder do CPW em promover a colaboração e eficiência.

Esses estudos de caso demonstram os benefícios tangíveis de usar o CPW para gerenciar workflows de ML. Ao alavancar a natureza modular e reutilizável do CPW, as equipes podem melhorar significativamente a eficiência de seus workflows e reduzir o tempo e esforço necessários para desenvolver e manter modelos de ML.

Lições Aprendidas

Durante o desenvolvimento e uso do CPW, identificamos várias lições-chave:

  1. Equilibrando flexibilidade e simplicidade: Enquanto o CPW permitiu casos de uso complexos, também introduziu uma nova linguagem de domínio específico (DSL), com uma curva de aprendizado mais acentuada. Fornecer abstrações mais simples para workflows básicos é crucial para uma adoção mais ampla. Com pontos de entrada diretos, podemos tornar o CPW mais acessível a novos usuários, oferecendo ainda capacidades avançadas para workflows complexos.
  2. Fornecendo exemplos e templates: Exemplos do mundo real e templates abrangentes são essenciais para demonstrar as capacidades do CPW e facilitar sua adoção.
  3. Identificação de “campeões” em cada equipe: Identificar e apoiar “campeões” dentro das equipes favoreceu uma adoção mais profunda e reduziu a necessidade de orientações constantes do nosso time principal.

Essas lições informaram nosso plano futuro, garantindo que o CPW continue a evoluir e atender às necessidades dos nossos usuários enquanto mantém sua flexibilidade e poder.

Conheça nossas oportunidades