Autor: Pedro Viana Gomes Pena

Revisado por Kevin Rossell

Los flujos de trabajo de machine learning (ML) son el corazón de la innovación basada en datos, pero asegurar que sean escalables y adaptables puede ser todo un reto. Por eso creamos Common Python Workflows (CPW), una solución versátil que simplifica cómo los equipos crean, gestionan y ejecutan pipelines de machine learning.  

CPW aborda problemas comunes como la complejidad de los flujos de trabajo y la falta de flexibilidad, permitiendo que los equipos se concentren en mejorar sus modelos en lugar de lidiar con procesos ineficientes. En este blog, exploraremos las razones detrás del desarrollo de CPW, sus características únicas y los aprendizajes clave de nuestra experiencia hasta ahora.  

Desafios inciales

Antes de CPW, enfrentábamos varias limitaciones con nuestra herramienta existente, Sheep:  

  • Complejidad y confusión: A lo largo de los años, Sheep se sobrecargó con numerosas funciones, lo que llevó a una experiencia de usuario confusa. Escribir archivos de configuración adecuados para controlar el comportamiento del entrenamiento también era complicado, dificultando a los usuarios aprovechar todo su potencial.  
  • Falta de extensibilidad: Sheep fue diseñada como una biblioteca Python única con todo el código y dependencias agrupadas. Esto dificultaba extender o sustituir ciertas funcionalidades, como la carga de datos. Los usuarios tenían poca flexibilidad para modificar la pipeline: solo podían reemplazar 3 etapas predefinidas.  
  • Soporte para frameworks: Dar soporte a varios frameworks (como Scikit-learn, Pytorch y Tensorflow) no era sencillo, y distribuir el entrenamiento era laborioso. Los usuarios debían entrenar el modelo con fklearn.
  • Entrenamiento out-of-core: Los usuarios no podían realizar eficientemente el entrenamiento out-of-core. 
  • Puntuación y evaluación: No había abstracciones para la puntuación, evaluación o comparación de modelos.  
  • Re-entrenamiento automático: Las pipelines de re-entrenamiento automático no se proporcionaban de forma predeterminada.  
  • Gestión de dependencias: Incluso los cambios menores requerían la actualización de múltiples dependencias en los modelos, lo que frecuentemente desembocaba en el infierno de las dependencias.   

Estos desafíos nos llevaron a desarrollar CPW: una solución flexible y preparada para el futuro que simplifica los flujos de trabajo de ML y resuelve estos puntos problemáticos. 

¿Qué es CPW?  

CPW es una biblioteca diseñada para simplificar la creación y gestión de flujos de trabajo de ML. Introduce varias características clave:  

  • Workflows de entrenamiento personalizados: Adapta los flujos de trabajo a necesidades específicas.  
  • Heredar y adaptar flujos de entrenamiento: Reutiliza workflows con mínimas modificaciones.  
  • Compartir flujos entre equipos: Facilita la colaboración permitiendo compartir fácilmente los flujos de trabajo.  
  • Manejo de serialización/deserialización: Simplifica el guardado y carga de modelos y datos.  
  • Implementación de re-entrenamiento automático: Implementa fácilmente mecanismos de re-entrenamiento automático.

Para ejecutar pipelines de modelos, utilizamos Kubeflow Pipelines: una plataforma open-source de machine learning que facilita la implementación y gestión de flujos de trabajo de ML de extremo a extremo. Además de orquestar pipelines de ML en Kubernetes, ofrece herramientas para gestionar y monitorear ejecuciones de pipeline.  

CPW se integra perfectamente con Kubeflow a través de Argo Workflows. Una vez que se crea un flujo (Flow), este puede compilarse en un Argo Workflow y someterse a ejecución en Kubeflow, que ofrece características como:  

  • Ejecutar tareas en paralelo, cada una en un contenedor diferente.  
  • Ajustar hiperparámetros con Katib (Componente de Kubeflow para la búsqueda de hiperparámetros).  
  • Asignar diferentes recursos de máquina para cada etapa.  

CPW se basa en Dagster, aprovechando principalmente los conceptos de activos definidos por software y IOManagers.  

Por qué elegimos Dagster  

Al diseñar CPW, evitamos deliberadamente reinventar soluciones existentes para Grafos Acíclicos Dirigidos (DAG) y orquestación. Había varios frameworks open-source que lo hacían bien. Queríamos enfocarnos en la experiencia de creación de pipelines.  

Después de evaluar frameworks como Metaflow y Dagster, elegimos Dagster por su:  

  • Diseño modular: El diseño modular de Dagster permite una separación clara de responsabilidades. Cada etapa en un flujo de trabajo puede definirse de manera independiente, y las dependencias se infieren a partir de los artefactos producidos por cada etapa.  
  • Activos definidos por software: El concepto de activos definidos por software en Dagster se alinea con nuestro objetivo de crear workflows reutilizables y flexibles.  
  • IOManagers: Los IOManagers de Dagster proporcionan una forma robusta de manejar operaciones de entrada y salida, asegurando que los datos se gestionen de manera eficiente y consistente en las distintas etapas del workflow.  
  • Extensibilidad: El diseño de Dagster permite una mayor extensibilidad, permitiendo a los usuarios personalizar solo los componentes que necesitan.  

Diseño de la clase workflow  

Para facilitar la creación de workflows reutilizables, introdujimos la clase FlowSpec, que permite que los equipos hereden y adapten workflows fácilmente. Esta clase permite cambios mínimos, lo que facilita la colaboración y la adaptación de flujos entre equipos.

Este enfoque aporta ventajas adicionales:

  • Integración con la infraestructura: Crea un punto central para que nuestro equipo lo mantenga y evolucione. Cuando algo cambia en nuestra infraestructura, podemos actualizar la clase base y estar seguros de que todos los flujos de trabajo se adaptarán al cambio.
  • Facilidad de desarrollo: El desarrollo de nuevas funcionalidades se ha vuelto más fácil porque podemos confiar en los conceptos modulares que teníamos para Assets, IOManagers y Flows (repositorios en Dagster).
  • A modo de ejemplo, desde que introdujimos CPW, hemos habilitado las siguientes funcionalidades para los usuarios:
  • Diferentes backends para la ejecución: cambiando una palabra clave en la definición del asset, los usuarios pueden ejecutar este paso en clusters Spark en Databricks o Sagemaker.
  • Profiling: pasando una palabra clave en la línea de comandos, se habilita cProfiling para todos los pasos de Flow.

Esto permite realizar cambios mínimamente intrusivos que requieren poca o ninguna modificación del código del usuario para poder ser utilizados.

Este diseño también ha facilitado la colaboración, ya que los equipos desarrollan Flujos de formación básicos que se utilizan para varios modelos, lo que permite aplicar las mejoras una sola vez a todos los modelos.

Conceptos centrales

CPW se basa en varios conceptos centrales de Dagster, principalmente activos definidos por software e IOManagers. Estos conceptos permiten definir y gestionar componentes de flujo de trabajo de forma modular y reutilizable.

He aquí un ejemplo de definición de un flujo en CPW:

En este ejemplo, hemos definido un flujo de trabajo sencillo con tres activos:

  • params: Un diccionario que contiene definiciones estáticas necesarias a lo largo del flujo de trabajo, como hiperparámetros.
  • my_dataset: Un conjunto de datos que se divide en función de los parámetros definidos en params.
  • train_model: Un modelo entrenado usando el conjunto de datos y los parámetros.

Basándose en los nombres y argumentos de los activos (funciones), CPW establece automáticamente las dependencias entre las funciones. El DAG (grafo acíclico dirigido) resultante tendrá el siguiente aspecto:

Descubre las oportunidades

Artefactos y persistencia

CPW garantiza que los artefactos de cada activo se conservan en un directorio S3 compartido (o local). Por ejemplo:

  • params: Almacenado como un archivo JSON.
  • my_dataset: Guardado como un Parquet.
  • train_model: Serializado y guardado como un objeto Pickle.

CPW gestiona la serialización y deserialización de estos artefactos, garantizando que se almacenan y recuperan de forma eficiente. Esta persistencia automática permite a los usuarios concentrarse en definir su flujo de trabajo sin preocuparse de la gestión subyacente de los datos.

Al aprovechar los activos definidos por software y los IOManagers de Dagster, CPW proporciona un marco sólido y flexible para crear y gestionar flujos de trabajo de ML. Este enfoque modular facilita la definición, el uso compartido y la adaptación de los flujos de trabajo, lo que permite a los equipos trabajar de forma más eficiente y eficaz, centrándose únicamente en lo que importa:

  • La función que debe realizar tu Flow, y;
  • Los artefactos que debe producir su flujo.

Casos prácticos

CPW ha permitido a varios equipos simplificar considerablemente sus flujos de trabajo:

Compartir el mismo flujo base para diversos modelos de PLN

Uno de los éxitos más significativos procede de un equipo que trabaja con modelos de procesamiento del lenguaje natural (PLN). Al compartir un flujo base común para varios modelos de PNL, el equipo pudo reducir en un 80% el esfuerzo de actualización del flujo de trabajo de entrenamiento de los modelos.

Esto se consiguió reutilizando la misma estructura de flujo de trabajo y modificando únicamente los parámetros específicos de cada modelo. El diseño modular de CPW permitió al equipo gestionar y adaptar eficazmente sus flujos de trabajo, lo que se tradujo en un importante ahorro de tiempo y un aumento de la productividad.

Distribuir las mejoras entre los equipos

Otro caso de éxito es el de un científico de datos que desarrolló código de evaluación y gráficos para su modelo. Gracias a CPW, pudo desplegar inmediatamente estas mejoras en todos los demás modelos de su equipo. Esto le ahorró semanas de esfuerzo y trabajo de ingeniería, ya que podían integrarse sin problemas en los flujos de trabajo existentes. La posibilidad de compartir y reutilizar componentes del flujo de trabajo entre distintos modelos y equipos pone de manifiesto el poder de CPW para fomentar la colaboración y la eficacia.

Estos casos prácticos demuestran las ventajas tangibles de utilizar CPW para gestionar los flujos de trabajo de ML. Al aprovechar la naturaleza modular y reutilizable de CPW, los equipos pueden mejorar significativamente la eficiencia de sus flujos de trabajo y reducir el tiempo y el esfuerzo necesarios para desarrollar y mantener modelos de ML.

Lecciones aprendidas

Durante el desarrollo y el uso de CPW, identificamos varias lecciones clave:

  1. Equilibrio entre flexibilidad y sencillez: aunque CPW permitió casos de uso complejos, también introdujo un nuevo lenguaje específico del dominio (DSL), con una curva de aprendizaje más pronunciada. Proporcionar abstracciones más sencillas para flujos de trabajo básicos es crucial para una adopción más amplia. Con puntos de entrada directos, podemos hacer que CPW sea más accesible para los nuevos usuarios, sin dejar de ofrecer capacidades avanzadas para flujos de trabajo complejos.
  2. Proporcionar ejemplos y plantillas: los ejemplos del mundo real y las plantillas completas son esenciales para demostrar las capacidades del CPW y facilitar su adopción.
  3. Identificación de responsables en cada equipo: Identificar y apoyar a los responsables dentro de los equipos ha favorecido una adopción más profunda y ha reducido la necesidad de orientación constante por parte de nuestro equipo central.

Estas lecciones han servido de base a nuestro plan de futuro, garantizando que CPW siga evolucionando y satisfaciendo las necesidades de nuestros usuarios, al tiempo que mantiene su flexibilidad y potencia.

Descubre las oportunidades