Feature stores en ML en tiempo real, parte 3:  Ejemplos de Nubank

Escrito por: Felipe Almeida, con contribuciones de Luiz Felix, Vívian Miranda, Lucas Nolasco y Caroline Custodio.

Esta es la parte 3 de una serie de 3 partes sobre Feature Stores para ML en tiempo real. Consulta la parte 1 y la parte 2.

Breve resumen e introducción

En la parte 1 de esta serie describimos qué son las feature stores y por qué tienen sentido cuando una organización orientada a datos alcanza cierto nivel de madurez, especialmente al usar sistemas en tiempo real habilitados por ML. En la parte 2, enumeramos las lecciones que aprendimos al aplicar conceptos de feature store a problemas reales en Nubank.

En las siguientes secciones repasaremos algunos patrones arquitectónicos concretos que hemos implementado. Algunos de ellos son lo que podríamos llamar una feature store completa (con definiciones de features comunes para los caminos de entrenamiento y de serving, por ejemplo), mientras que otros se enfocan en resolver solo un problema particular.

Features precalculadas a partir de datos batch, materializadas en una capa de serving clave-valor

En esta variante producimos un dataset totalmente desnormalizado (con todas las features de cada entidad combinadas en una sola fila) como parte de nuestros jobs nocturnos de ETL/ELT y lo ponemos a disposición en una “capa de serving” (por ejemplo, DynamoDB), indexado por el ID, para usarlo en tiempo de inferencia. Esto se puede ver en la Figura 1:

Figura 1: En esta estrategia se agrega un paso extra a los jobs nocturnos de ETL/ELT: producir un dataset, indexado por el ID de la entidad (por ejemplo, ID de cliente, ID de cuenta), donde cada columna es el valor de una feature para esa entidad. Luego, ese snapshot se envía a un almacén clave-valor (KV) de baja latencia como DynamoDB o Redis. En tiempo de serving, el modelo de ML en tiempo real  recuperará una sola fila de ese almacén, ya que está totalmente desnormalizado. Esto es rápido y fácil.

Este suele ser el primer “patrón” de feature store al que convergen muchas empresas, a menudo sin saber que se trata de una feature store. Se basa en la arquitectura lambda (sin la speed layer).

Es un caso de uso simple, pero bastante poderoso:

• Es barato y fácil: No se necesitan habilidades especializadas de ingeniería de machine learning para usarlo en producción; quien construyó el dataset de entrenamiento (por ejemplo, los científicos de datos) puede construir el dataset de la capa de serving, mientras que los ingenieros que despliegan el modelo solo necesitan una única solicitud para obtener varias features a la vez. 

• No habrá training-serving skew “inducido por lógica” si el entrenamiento y el serving reutilizan la misma lógica de cálculo de features del data lake, pero aún puede haber skew debido a retrasos o fallas en las ejecuciones de ETL.

• La latencia de recuperación es muy baja porque es una simple búsqueda por clave que devuelve una sola fila.

La principal desventaja es que los valores de las features servidos en tiempo de inferencia estarán desactualizados, ya que son un snapshot de la última vez que se ejecutaron los jobs de ETL/ELT (no una representación del mundo tal como es ahora). Por lo tanto, hay que tener en cuenta cuándo cada snapshot habría estado realmente disponible al construir el dataset de entrenamiento, para evitar fuga de datos (los snapshots no están disponibles instantáneamente; pueden tardar horas y a veces fallan).

Como ocurre con todas las estrategias de implementación de features “shift-left”, esto puede volverse costoso porque se precalculan features para todas las entidades (por ejemplo, clientes, cuentas), aunque solo una minoría de ellas necesite ser evaluada en un día determinado.

Features tabulares on-demand con Flink y Pinot

Las features on-demand se calculan en tiempo de serving, en lugar de precalcularse con antelación. Esto es útil cuando los modelos necesitan la información más actualizada y precalcular features para entidades que quizá nunca las usen sería un desperdicio.

Pueden aumentar la latencia si no tomamos precauciones, pero tienen ventajas claras:

• Son lo más cercano que podemos llegar a features realmente online cuando necesitamos calcular agregaciones en tiempo de inferencia
• Solo se calculan cuando realmente se usan (a diferencia de las features precalculadas, que se materializan para toda la base de clientes por si acaso llegan a necesitarse)

Una de las estrategias que usamos para features on-demand es la de escritura dual (dual-write), en la que los microservicios transaccionales escriben eventos tanto en el camino offline (archivos Parquet en S3) como en el camino online (una app de Apache Flink y luego Pinot), como podemos ver en la Figura 2 a continuación.

Figura 2: Cada evento de comportamiento se enruta al camino offline y al camino online. Esto ayuda a reducir el training-serving skew en este caso, pero la lógica de agregación sigue siendo diferente en el camino offline (normalmente jobs de Spark) y en el camino online (queries SQL en Pinot).

En el camino de escritura, la estrategia de escritura dual anterior busca enviar los mismos eventos a los almacenes offline y online. Esto es importante para reducir el training-serving skew (pero no para eliminarlo, porque aún puede haber retrasos y eventos faltantes). Además, la lógica de agregación batch (usada en el entrenamiento) y las queries usadas en tiempo de serving son artefactos diferentes. Por esta razón, el monitoreo del skew sigue siendo necesario.

El camino de lectura consume los datos producidos por el camino de escritura (Figura 2). Tiene dos ramas: una para entrenamiento y otra para serving.

La lógica de tiempo de entrenamiento suele escribirse como jobs de procesamiento batch, ya que implica joins y agregaciones pesadas sobre horizontes de tiempo largos. La lógica de tiempo de serving es una query SQL ejecutada en tiempo real sobre una base de datos de procesamiento analítico en línea (OLAP) en tiempo real como Apache Pinot, como se puede ver en la Figura 3 a continuación. Esto forma parte de nuestro stack “Avalanche”, sobre el cual escribimos en este post del blog.

Figura 3: En tiempo de entrenamiento, construimos el dataset de entrenamiento a partir de un spine (un dataset con IDs de entidades y timestamps de referencia) y el event store, mediante un job de Spark. En tiempo de serving, ejecutamos la agregación en tiempo real, con una query SQL normal, sobre un backend de base de datos Pinot.

La query SQL enviada a Pinot es una query de agregación. Este es un ejemplo de la query que usaríamos para obtener la cantidad de veces que un cliente determinado abrió su app en las últimas 2 horas:

SELECT

  customer_identifier,

  count(DISTINCT event_timestamp) AS app_opened__last_2h__count

FROM clickstream_table

WHERE type = ‘app_opened’

  AND customer_identifier = ‘{{customer_identifier}}’

  AND event_timestamp >= AGO(‘PT2H’)

GROUP BY customer_identifier

Algunos puntos para que esto funcione a nuestra escala:

• Sin joins: Normalmente no hacemos ningún join en tiempo de serving (solo agregaciones group-by sobre la misma tabla). Si necesitamos más datos en la tabla, por lo general los enriquecemos en la capa de Flink (para no necesitar joins costosos en tiempo de query)
• TTL agresivo: Debido a nuestra escala, aplicamos un time to live (TTL) a los eventos en Pinot, para que el tamaño de la base de datos no crezca demasiado.
  • Por lo tanto, solo obtenemos datos de muy corto plazo (por ejemplo, los últimos días) de Pinot. Si necesitamos features de más largo plazo (por ejemplo, agregaciones del último año), las precalculamos a partir de datos batch, ya que es más fácil, rápido y barato obtenerlas de un snapshot con retraso calculado diariamente.
• Evitar cambios de schema: Los cambios de schema son un problema. Una solución es tener campos string con contenido JSON, ya que la estructura del JSON puede evolucionar sin cambiar el schema de la tabla (los cambios aún deben ser retrocompatibles para los consumidores existentes).

Una ventaja secundaria de este enfoque es que cualquier otro equipo puede usar los eventos que guarda un equipo: toda la lógica de la feature está contenida en la query, así que varios equipos pueden compartir la base de datos Pinot y ejecutar queries diferentes sobre ella.

Esta no es una feature store completa; solo usamos esta configuración para soportar el camino de serving de features online. El código de entrenamiento está en un camino completamente distinto, y por eso debemos seguir monitoreando el training-serving skew.

Uso de una arquitectura de materialización y serving en dos etapas basada en Flink

En esta estrategia, los eventos de comportamiento fluyen por dos apps de Flink, una tras otra (Flink app #1 y Flink app #2, como se puede ver en la Figura 4 a continuación). Esto es útil cuando los modelos necesitan features actualizadas, pero calcularlas en tiempo de serving añadiría demasiada latencia:

Figura 4: En este ejemplo, cada nuevo evento “purchase-requested” se transmite a través de Flink app #1, y un conteo continuo de las últimas 24 horas para cada cliente se mantiene en Flink app #2 y luego se envía mediante upsert a DynamoDB. Cada vez que cualquier modelo en tiempo real necesita esta información, solo tiene que llamar a DynamoDB pasando el ID del cliente y el nombre del feature set para obtener el conteo materializado más reciente de ese cliente.

Flink app #1 escucha eventos de comportamiento (logs de CDC de bases de datos transaccionales y eventos crudos de Kafka), combina datos de múltiples shards en un solo stream, filtra los eventos no deseados y normaliza nombres y schemas. En esta primera app no se mantiene ningún estado.

Flink app #2 escucha los eventos de la primera app y mantiene cierto estado, guardando todos los eventos de cada cliente, con un TTL de 24 horas. Cuando un nuevo evento entra en la app #2, todas las features relacionadas se recalculan y se envía mediante upsert una nueva fila a una tabla de DynamoDB (indexada por un ID de entidad, como el ID de cliente).

Esta estrategia depende del precálculo de features, por lo que su latencia en tiempo de serving es muy baja (ya que es una simple búsqueda). Al igual que en el primer ejemplo de este artículo, el trade-off es que la mayoría de las features se calcularán para clientes que nunca las necesitarán.

Como en el ejemplo anterior, esta configuración solo cubre el camino de serving, así que no es una feature store completa. El código de entrenamiento corre en Spark, por lo que puede haber skew inducido por lógica, que debe monitorearse.

Próximos pasos

Además de mejorar constantemente nuestra configuración y aprender nuevas lecciones, hay al menos otras dos líneas que estamos explorando actualmente:

• Pipelines de Feldera: Recientemente hemos estado experimentando con una herramienta llamada Feldera. Los pipelines de Feldera soportan mantenimiento incremental de vistas, lo que significa que las features basadas en agregaciones (sum, count, avg, etc.) se actualizan en tiempo real a medida que llegan nuevos eventos, sin recalcular toda la agregación desde cero cada vez).

Además, Feldera podría permitirnos unificar la lógica de features para entrenamiento y serving, reduciendo los problemas de training-serving skew y acercando a Nubank a una feature store en tiempo real completa y para toda la empresa.

• Features secuenciales: Aunque la mayoría de los ejemplos de esta serie involucran features tabulares, también hemos estado usando features secuenciales en varios casos de uso, desde 2023.

Las features secuenciales generan mucho valor, particularmente en redes neuronales, pero también introducen complejidad adicional y nuevos modos de falla; descubrir cómo integrarlas en nuestras plataformas de features en tiempo real es un desafío en curso.

Conclusión

En este post repasamos algunos patrones arquitectónicos que usamos en Nubank para poner features a disposición de modelos de ML en tiempo real.

• Las features batch son simples y baratas y no deben pasarse por alto, incluso cuando eso signifique que el modelo en tiempo real usará información ligeramente desactualizada.

• Decidir si precalcular las features o calcularlas on-demand es uno de los principales trade-offs al pensar en feature stores. Distintos casos de uso requieren distintos enfoques. Una buena feature store debería soportar ambos.

• Unificar los caminos de definición de features de entrenamiento y de serving es difícil y, aunque no resuelve por completo el training-serving skew, sí lo mitiga sustancialmente.

Con esto cerramos nuestra serie de 3 partes sobre feature stores para ML en tiempo real en Nubank. No dejes de leer las partes 1 y 2  y mantente atento a otros posts sobre este tema, ¡ya que es algo en lo que estamos trabajando activamente!