Cómo Snorkel AI logró ahorrar más del 40 % en costos al escalar las cargas de trabajo de machine learning con Amazon EKS

Las startups de machine learning (ML) suelen utilizar muchos recursos de computación, ya que entrenan modelos de gran tamaño con GPU de gama alta y los implementan a escala para realizar inferencias. Startups de AWS colabora con startups desde su creación hasta su salida a bolsa y ha ayudado a miles de fundadores e innovadores del ámbito de la inteligencia artificial (IA) a desarrollar sus negocios en Amazon Elastic Kubernetes Service (Amazon EKS). Amazon EKS es una opción popular para crear y alojar modelos de ML porque ofrece la flexibilidad de Kubernetes con la seguridad y la resiliencia de ser un servicio administrado de AWS que está optimizado para crear cargas de trabajo en contenedores de alta disponibilidad.

Snorkel AI es una de esas empresas que se beneficia de Amazon EKS. Snorkel AI ayuda a las empresas incluidas en la lista Fortune 500, a las agencias federales y a los innovadores del ámbito de la inteligencia artificial a crear, adaptar y extraer modelos de base (FM) y modelos de lenguaje grandes (LLM) para que funcionen con gran precisión en conjuntos de datos de dominios específicos. Gracias al enfoque de Snorkel centrado en los datos para el desarrollo de la inteligencia artificial, las organizaciones han creado servicios de IA listos para la producción para casos de uso, como el procesamiento de reclamaciones de seguros, la distribución financiera, el análisis de ensayos clínicos y la aceleración de la administración proactiva de pozos para la perforación en alta mar.

Durante los últimos meses, el equipo de Snorkel ha trabajado arduamente para abordar los desafíos únicos de diseñar una infraestructura eficiente para respaldar las cargas de trabajo de desarrollo de ML sin aumentar las facturas de infraestructura, reducir la velocidad de desarrollo ni repercutir en la experiencia del usuario. Su objetivo final era reducir el gasto en computación en clústeres de Snorkel Flow, su plataforma integral de ML, en más de un 40 %.

Descripción general de Snorkel Flow

La plataforma de desarrollo de datos de IA Snorkel Flow de Snorkel permite a los equipos de datos crear rápidamente aplicaciones de IA gracias al uso de un ciclo iterativo de etiquetado programático, el entrenamiento rápido de modelos y el análisis de errores. Cada proyecto comienza cuando los usuarios crean una pequeña cantidad de funciones de etiquetado.

Las funciones de etiquetado emplean heurísticas simples, bases de datos externas, modelos heredados o incluso llamadas a modelos de lenguaje grandes para aplicar etiquetas a franjas de datos sin etiquetar en función de la intuición codificada de expertos. El débil algoritmo de supervisión de la plataforma combina estas funciones basadas en reglas para determinar la etiqueta más adecuada para cada registro. Luego, los usuarios entrenan un modelo simple basado en estos puntos de datos probabilísticos y evalúan el impacto de cada función de etiquetado. En la fase de análisis, los usuarios investigan los segmentos de los datos en los que el modelo tiene un rendimiento inferior. A continuación, crean o modifican las funciones de etiquetado, entrenan otro modelo rápido y continúan el ciclo. Cuando los usuarios están satisfechos con la calidad de sus etiquetas, crean un modelo final a partir de una arquitectura del zoológico de modelos (que abarca desde la regresión logística hasta los FM) y lo exportan para su implementación.

Debido a la naturaleza de este flujo de trabajo, la infraestructura de Snorkel Flow experimenta diversos periodos de elevado uso de computación. Los costos operativos aumentaron naturalmente a medida que aumentaban la base de clientes y las capacidades de los productos de ML de Snorkel Flow. Para lograr un crecimiento eficiente, Snorkel se propuso entender cómo mejorar los beneficios mientras operaba un software de ML de última generación. Snorkel implementó las siguientes prácticas para lograr una reducción de más del 40 % en los costos de computación en clústeres.

Soluciones para la optimización de costos de la nube 

Las startups de software como servicio (SaaS) suelen tener la oportunidad de optimizar sus gastos de la nube. Es esencial comprender los factores únicos que impulsan estos costos.

Para Snorkel, había dos factores importantes:

1. Las cargas de trabajo de desarrollo de ML suelen requerir hardware especializado y caro, como las GPU. Por lo general, estas cargas de trabajo son, en esencia, “ampliables”.
2. Las empresas de la lista Fortune 500 y las grandes agencias federales utilizan Snorkel. Entre estas encontramos las principales instituciones financieras con departamentos de TI sofisticados que tienen requisitos específicos de implementación y privacidad de datos, mediante el uso de una plataforma en contenedores.

El equipo de Snorkel está interesado en crear sistemas que permitan un escalado eficiente sin un aumento lineal de los costos de infraestructura. En consecuencia, Snorkel desarrolló una solución integral de escalado automático adaptada a sus cargas de trabajo de ML en Amazon EKS para abordar los problemas relacionados con los gastos de la nube. Esta solución no solo agilizó las cargas de trabajo que requerían una computación ampliable, sino que también logró sus objetivos de reducción de costos.

Además de la solución de escalado automático, las estrategias clave que contribuyeron a la reducción de más del 40 % en los gastos de la nube incluyen:

• Colaborar con los líderes de ingeniería y el equipo de AWS para adoptar Savings Plans mediante optimizaciones de la configuración de la nube.
• Ajustar el tamaño de los recursos al supervisar el uso de los nodos con Prometheus y consultar a los ingenieros de backend para evaluar las necesidades de los componentes de la plataforma.
• Cambiar a tipos de máquinas virtuales (VM) rentables en Amazon EKS y utilizar instancias de Amazon Elastic Compute Cloud (Amazon EC2) con varias GPU para mejorar la relación precio-rendimiento.
• Establecer modificaciones de procesos internos en las que los ingenieros colaboraran con los equipos de atención al cliente para minimizar la computación inactiva.

En esta publicación, Snorkel comparte el proceso para abordar estos desafíos de escalado a fin de poder facilitar el diseño de una mejor infraestructura para los sistemas de ML. Si está empezando a usar Kubernetes, lea la publicación de Snorkel Introduction to Kubernetes para obtener más información sobre los conceptos básicos y la publicación Machine learning on Kubernetes: wisdom learned at Snorkel para obtener más información sobre su experiencia con Kubernetes hasta el momento.

Cómo funciona Snorkel Flow en AWS

En la práctica, la interacción de Snorkel Flow con AWS sigue la secuencia que se describe a continuación. Como la plataforma Snorkel Flow depende en gran medida de los contenedores, la migración a AWS se llevó a cabo prácticamente sin problemas. 

• Los usuarios acceden a su instancia de Snorkel Flow desde su navegador web, que se asigna a una regla de Amazon Route 53.
• Route 53 reenvía la solicitud a un equilibrador de carga de aplicación.
• A continuación, el equilibrador de carga de aplicación reenvía la solicitud a los pods de Snorkel Flow que se ejecutan en un clúster de EKS compartido. Snorkel cambió el tipo de instancia de EC2 de m5 a m6a para optimizar los costos, lo que se tradujo en un ahorro en computación del 10 % con un impacto insignificante en el rendimiento en función del costo por hora para la misma CPU y RAM.
• Además, pasaron de una única instancia de GPU g4dn.8xlarge a una instancia de múltiples GPU g4dn.12xlarge, lo que permitió ofrecer 4 veces más pods de GPU.
• Cada instancia de Snorkel Flow usa un volumen de Amazon Elastic File System (Amazon EFS) para almacenar los archivos en el disco.
• Una cola de Redis autoalojada en un pod de la instancia de EC2 contiene los trabajos entrantes y espera a que los pods de trabajo los recojan.
• Las métricas de EKS se envían a Amazon CloudWatch y los scripts personalizados supervisan los registros para detectar anomalías en el rendimiento del clúster.

Esta arquitectura ha proporcionado una experiencia estable y ágil para los usuarios de Snorkel Flow.

Redefinición del escalado

Antes de la arquitectura descrita en la figura 1, las primeras iteraciones de la infraestructura de Snorkel utilizaban recursos fijos. Los usuarios de Snorkel señalaban que estas cargas de trabajo acumuladas podían tardar demasiado en completarse, lo que afectaba negativamente a su experiencia.

El escalado manual de los recursos de computación demostró ser poco escalable y propenso a errores, lo que llevó a que los costos de la nube se mantuvieran elevados incluso durante los periodos de bajo uso. Era lo peor de ambos mundos: una baja rentabilidad de la nube y un rendimiento inferior al necesario.

Para abordar estos desafíos, Snorkel implementó el escalado automático en varios niveles de su infraestructura, como se explica en las siguientes secciones.

Diseño de una infraestructura escalable teniendo en cuenta la rentabilidad

La distribución de Snorkel Flow en Kubernetes implica un conjunto de implementaciones que se ejecutan en un clúster de EKS que contiene pods que ejecutan varios componentes de la plataforma.

Como se muestra en la figura 2, para abordar los desafíos únicos de trabajar con cargas de trabajo de computación ampliable, el equipo de Snorkel introdujo un nuevo concepto para los pods de Kubernetes: los categorizó semánticamente como “fijos” o “flexibles”.

• Los pods fijos no se pueden mover de forma segura de un nodo a otro, ya sea porque perderán un estado importante en la memoria (como los trabajos de computación en curso sin puntos de control) o para minimizar el tiempo de inactividad evitable de los componentes fundamentales de la plataforma (por ejemplo, el orquestador del clúster de Ray).
• Los pods que son flexibles se pueden mover de forma segura a un nuevo nodo. Esta distinción es significativa en el contexto del escalado automático, ya que reducir verticalmente los nodos implica alejar los pods de los nodos infrautilizados una vez finalizados.

Este marco fijo y flexible proporciona a Snorkel un medio específico para cada dominio para permitir reducir verticalmente los clústeres, lo que le permite activar el escalador automático de clústeres en Amazon EKS sin que el departamento de finanzas le envíe mensajes cada hora.

El enfoque inicial de Snorkel consistía en implementar podDisruptionBudgets en el clúster de EKS para evitar que el escalador automático de clústeres moviera los pods flexibles durante el día y los pods fijos en cualquier momento. Si bien fue eficaz, este enfoque dejó insatisfecho al equipo de Snorkel, ya que redujo verticalmente muchos menos nodos de lo que era teóricamente óptimo.

Para solucionar este problema, Snorkel combinó una optimización de la programación de pods que aislaba los pods fijos en un pequeño grupo fijo de nodos. Programó los pods flexibles y de trabajo (que se consideran pods fijos, pero son efímeros debido al escalado automático de los nodos de trabajo) en el resto del grupo flexible de nodos.

Estos cambios permitieron a Snorkel reducir verticalmente de manera eficiente los nodos flexibles durante la noche, cuando era seguro mover pods flexibles y reducir verticalmente la gran mayoría de los pods de trabajo.

Al permitir reducir verticalmente de manera eficiente la gran mayoría de los nodos del clúster (es decir, los nodos flexibles), Snorkel pudo cumplir su objetivo de reducir los costos de la nube derivados del alojamiento de Snorkel Flow en más de un 40 %.

Más detalles sobre la solución de escalado automático de Snorkel

Snorkel divide la implementación de la solución descrita en la sección anterior en tres esfuerzos secuenciales:

1. En primer lugar, Snorkel implementó el “escalado automático de procesos de trabajo”, un servicio personalizado de escalado automático basado en Redis que permite a sus pods de trabajo escalar y reducir verticalmente en función de los trabajos en las colas de los procesos de trabajo.
2. En segundo lugar, implementó el “escalado automático de clústeres” mediante la reconfiguración de sus implementaciones de Kubernetes para permitir que el escalador automático de clústeres de Kubernetes redujera los nodos, además de escalarlos verticalmente.
3. En tercer lugar, Snorkel implementó “optimizaciones de reducción vertical de nodos” mediante la agrupación de los pods fijos en un pequeño grupo de nodos fijos para evitar que los pods fijos interfirieran con la reducción vertical de los nodos restantes.

Escalado automático de procesos de trabajo

La plataforma Snorkel Flow abstrae la computación en un paradigma en el que los trabajos esperan en las colas de Redis y los procesos de trabajo se ejecutan como procesos en pods de trabajo.

Snorkel implementó una solución de escalado automático de procesos de trabajo (figura 3) para los pods de trabajo mediante la ejecución de una función recurrente en la API de backend de Snorkel Flow. Cada pocos segundos, esta función comprueba si el clúster de Kubernetes y Redis cumplen los requisitos de escalado y reducción vertical.

Si hay trabajos esperando en una o más colas relevantes basadas en Redis, la función solicitará a la API de Kubernetes que aprovisione pods de trabajo adicionales para procesar estos trabajos. Si la cola de Redis está vacía y no hay trabajos en ejecución en el registro de trabajos, solicitará a la API de Kubernetes que destruya los pods de trabajo para liberar los recursos de CPU y RAM reservados.

Como se muestra en la figura 4, con la implementación de esta implementación de escalado automático de procesos de trabajo, los pods de trabajo de Snorkel Flow pasaron a ser efímeros y aparecieron en el clúster solo cuando era necesario procesar los trabajos.

Escalado automático de clústeres

El recurso PodDisruptionBudget protege ciertos pods contra interrupciones (por ejemplo, reinicios voluntarios) al permitir especificar la cantidad máxima de réplicas de pods que pueden no estar disponibles en un momento dado. Como se muestra en la figura 5, establecer este valor de forma explícita en 0 para una implementación garantiza que el escalador automático de clústeres no reduzca verticalmente los nodos que ejecutan los pods de la implementación.

La implementación de este recurso en instancias alojadas de Snorkel Flow de forma segura permitió al escalador automático del clúster reducir la escala de los nodos infrautilizados. Sin embargo, los ahorros de costos que consiguió Snorkel fueron mínimos: ni con esta implementación lograron reducir la escala de la mayoría de sus nodos porque todos los pods de Snorkel Flow estaban protegidos por un podDisruptionBudget asociado.

Tras un análisis más detenido, el equipo de Snorkel se dio cuenta de que esta protección no tiene por qué existir todo el tiempo. Las cargas de trabajo son intensas y la mayoría de las interacciones de los usuarios con Snorkel Flow se producen durante el horario laboral del cliente, por lo que es seguro relajar esta protección fuera del horario laboral. De forma similar al escalado automático de los procesos de trabajo, Snorkel implementó una función recurrente que desactivaba podDisruptionBudgets de un día para otro en los pods flexibles de una instancia. Para ello, establecía el número máximo de réplicas de pods no disponibles en 1, en lugar de 0 (figura 6). La solución anterior de escalado automático para procesos de trabajo, combinada con las características ClusterAutoscaler y PodDisruptionBudget, permitía reducir la escala de muchos más nodos de trabajo infrautilizados que antes. Los clientes que implementen Snorkel Flow en su nube pueden configurarlo según sea necesario.

Optimizaciones de la reducción automática de nodos

Incluso con estas mejoras, Snorkel vio que la mayoría de los nodos infrautilizados no se estaban reduciendo en absoluto.

Tras una investigación más detallada, Snorkel se dio cuenta de que el problema se debía a que los pods fijos y flexibles ocupaban el mismo nodo. Esto resultaba problemático, porque un pod fijo, asignado de forma seudoaleatoria a un nodo que contenía pods flexibles, bloqueaba ese nodo y evitaba que se redujera su tamaño, incluso si no se utilizaba lo suficiente. Esta falta de control sobre la programación de los pods fijos provocó periodos en los que no era posible reducir la escala de la gran mayoría de los nodos del clúster, a pesar de que representaban mucha más potencia de computación de la necesaria en ese momento.

Snorkel utilizó el recurso podAffinities de Kubernetes para solucionar este problema, lo que le permitió restringir los nodos en los que podía ejecutarse un pod en función de las etiquetas de otros pods que ya se estaban ejecutando en un nodo determinado.Snorkel agregó etiquetas a los pods para diferenciar entre los pods fijos y los flexibles. Agregó también un bloque podAntiAffinity a la configuración de sus implementaciones para garantizar que los pods fijos no estuviesen programados en los nodos que ejecutan pods flexibles, y viceversa.

Esta implementación de podAffinities permitió a Snorkel AI dividir los nodos en dos grupos funcionales: el grupo fijo de nodos que contienen pods fijos, que nunca se pueden mover de forma segura entre los nodos (por ejemplo, Redis debido a la caché), y el grupo flexible de nodos que contienen pods “flexibles” que son efímeros (como los pods de trabajo) o que se pueden mover de forma segura fuera del horario laboral (por ejemplo, durante la noche).

Aunque es posible con la intervención manual durante el mantenimiento de la plataforma, Snorkel no puede reducir verticalmente de forma automática los nodos fijos. Sin embargo, esta solución permite reducir verticalmente los nodos flexibles, ya que ahora se han aislado los pods que no se pueden mover en nodos fijos.

Conclusión

Los equipos de Snorkel y Startups de AWS esperan que compartir este proceso de pensamiento y estas soluciones ayude a otras startups a crear una mejor infraestructura para las cargas de trabajo de ML, que están adquiriendo cada vez más importancia a medida que el ML, los modelos de lenguaje grandes y otros FM se introducen en la producción para organizaciones de todo el mundo.

¿Va a asistir a AWS re:Invent 2023? Esta implementación se destacará como parte de la sesión de Snorkel Navigating the future of AI: Deploying generative models on Amazon EKS (sesión CON312). ¡No se la pierda!

Muchas gracias a David Hao, Edmond Liu y Alec Xiang por ayudar a hacer realidad esta visión técnica para Snorkel. Un agradecimiento especial a los ya mencionados, así como a Matt Casey, Henry Ehrenberg, Anthony Bishopric y a todo el infrastructure engineering team de Snorkel, por sus considerados comentarios sobre este artículo.