Uber, una de las empresas líderes a nivel mundial en innovación tecnológica y operaciones a gran escala, ha atravesado una significativa evolución en la forma en que gestiona sus cargas de trabajo de machine learning. La migración de sus procesos desde un sistema tradicional basado en Apache Spark y el servicio Michelangelo Deep Learning Jobs (MADLJ) hacia una plataforma unificada que usa Ray sobre Kubernetes representa un cambio estratégico crucial para enfrentar los retos de escalabilidad, eficiencia y flexibilidad que demandan las aplicaciones de inteligencia artificial modernas. Originalmente, la infraestructura de machine learning de Uber se apoyaba principalmente en MADLJ, un sistema diseñado para coordinar distintas tareas, desde procesos ETL (extracción, transformación y carga) utilizando Apache Spark hasta el entrenamiento de modelos a través de Ray. Sin embargo, este enfoque presentaba limitaciones notables, particularmente en la gestión de recursos computacionales. Los ingenieros de machine learning debían seleccionar manualmente el hardware adecuado, evaluando la disponibilidad de GPUs y la capacidad del clúster, lo que generaba ineficiencias y retrasos considerables en los flujos de trabajo.
Además, la configuración rígida y estática de los recursos y clústeres generaba una distribución de cargas poco equilibrada y un subuso frecuente de las capacidades disponibles. La infraprovisión provocaba fallos en las tareas o demoras significativas, mientras que la sobreprovisión implicaba un desperdicio costoso de recursos. Este estado de rigidez dificultaba la escalabilidad y hacía necesario contar con una solución más adaptable y automatizada para mantener la competitividad y mejorar el rendimiento. Ante estos desafíos, Uber decidió migrar sus cargas de trabajo de machine learning a Kubernetes, declarando un cambio hacia una infraestructura más flexible, elástica y eficiente. Esta transición no solo implicó la adopción de una nueva plataforma tecnológica, sino también un replanteamiento profundo de cómo se gestionan los recursos y cómo se facilita la experiencia de usuario para los científicos y desarrolladores de datos.
El objetivo principal fue crear una plataforma en la que los usuarios pudieran especificar el tipo de trabajo y los recursos necesarios de una manera declarativa, desacoplando el detalle técnico de la infraestructura subyacente. Así, el sistema se encarga automáticamente de asignar y distribuir los recursos de manera óptima, evaluando las condiciones actuales del clúster y las prioridades de los trabajos en ejecución. Una de las innovaciones clave implementadas por Uber fue la introducción de un esquema de recursos jerárquico basado en pools, que organiza los recursos del clúster conforme a límites organizacionales o por equipos. Esta estructura permite una gestión más granular y una mejor visibilidad del uso de recursos dentro de la empresa, asegurando que cada equipo disponga de un presupuesto computacional definido pero flexible. Para maximizar la utilización del hardware, se implementó una política de elasticidad en la cual los pools pueden compartir recursos temporalmente.
Si un pool cuenta con recursos inactivos, estos pueden ser prestados a otro pool que los requiera, aumentando así la eficiencia general sin necesidad de redistribuciones permanentes. Estos recursos compartidos son preemptibles, lo que significa que pueden ser reclamados por el pool original cuando se requiera su uso prioritario. Este mecanismo se rige por principios de max-min fairness para garantizar tanto la equidad como la eficiencia en el acceso a los recursos. Otro aspecto fundamental en la optimización de la infraestructura fue el manejo específico del hardware heterogéneo, dada la presencia de nodos con GPU y otros solo con CPU en los clústeres de Uber. Para sacar el máximo provecho de este entorno, las tareas que no requieren aceleración por GPU, como la carga de datos y la preprocesamiento, son programadas en nodos con CPU únicamente, reservando los nodos GPU para las labores de entrenamiento del modelo, que demandan alta potencia de cómputo.
Para facilitar esta diferenciación, Uber diseñó un plugin de filtrado para pods con requerimientos de GPU, asegurando que solo estos workloads se programen en nodos compatibles. El scheduler de Kubernetes también fue mejorado para soportar distintas estrategias en función del tipo de carga: una estrategia orientada a la distribución equilibrada para pods sin GPU y una estrategia de bin-packing para cargas con GPU, buscando minimizar la fragmentación de recursos y maximizar la densidad de trabajo. Estos avances técnicos han supuesto un salto importante respecto al escenario anterior, brindando a Uber una infraestructura mucho más capaz de responder a la variabilidad y escalabilidad exigidas por sus aplicaciones de inteligencia artificial. La nueva plataforma permite automatizar la asignación de recursos, mejorar la utilización del hardware, y al mismo tiempo garantizar la equidad y estabilidad de operación entre distintos equipos y proyectos. Además, la implementación de Ray sobre Kubernetes ha mostrado otros beneficios como la simplificación del desarrollo y despliegue de modelos, mejor integración con los procesos de CI/CD, y una base más sólida para la experimentación rápida y el escalamiento en producción.