Feature stores em ML em tempo real, parte 3:  Exemplos do Nubank

Escrito por: Felipe Almeida, com contribuições de Luiz Felix, Vívian Miranda, Lucas Nolasco e Caroline Custodio.

Esta é a parte 3 de uma série de 3 partes sobre Feature Stores para ML em tempo real. Veja a parte 1 e a parte 2.

Breve recapitulação e introdução

Na parte 1 desta série, descrevemos o que são feature stores e por que elas fazem sentido quando uma organização orientada a dados atinge um certo nível de maturidade, especialmente ao usar sistemas em tempo real habilitados por ML. Na parte 2, listamos as lições que aprendemos ao aplicar conceitos de feature store a problemas reais no Nubank.

Nas próximas seções, vamos apresentar alguns padrões arquiteturais concretos que implementamos. Alguns deles são o que poderíamos chamar de uma feature store completa (com definições de features comuns para os caminhos de treinamento e de serving, por exemplo), enquanto outros focam em resolver apenas um problema específico.

Features pré-computadas a partir de dados em batch, materializadas em uma camada de serving chave-valor

Nesta variante, produzimos um dataset totalmente desnormalizado (com todas as features de cada entidade combinadas em uma única linha) como parte dos nossos jobs noturnos de ETL/ELT e o disponibilizamos em uma “camada de serving” (por exemplo, DynamoDB), indexado pelo ID, para ser usado em tempo de inferência. Isso pode ser visto na Figura 1:

Figura 1: Nesta estratégia, uma etapa extra é adicionada aos jobs noturnos de ETL/ELT: produzir um dataset, indexado pelo ID da entidade (por exemplo, ID do cliente, ID da conta), em que cada coluna é o valor de uma feature para aquela entidade. Em seguida, esse snapshot é enviado para um armazenamento chave-valor (KV) de baixa latência, como DynamoDB ou Redis. Em tempo de serving, o modelo de ML em tempo real vai  recuperar uma única linha desse armazenamento, já que ele é totalmente desnormalizado. Isso é rápido e fácil.

Esse costuma ser o primeiro “padrão” de feature store para o qual muitas empresas convergem, muitas vezes sem saber que se trata de uma feature store. Ele é baseado na arquitetura lambda (sem a speed layer).

É um caso de uso simples, mas bastante poderoso:

• É barato e fácil: Não são necessárias habilidades especializadas de engenharia de machine learning para usar isso em produção; quem construiu o dataset de treinamento (por exemplo, cientistas de dados) pode construir o dataset da camada de serving, enquanto os engenheiros que fazem o deploy do modelo só precisam de uma única requisição para buscar várias features de uma vez. 

• Não haverá training-serving skew “induzido por lógica” se o treinamento e o serving reutilizarem a mesma lógica de cálculo de features do data lake, mas ainda pode haver skew devido a atrasos ou falhas nas execuções de ETL.

• A latência de recuperação é muito baixa porque é uma simples busca por chave, retornando uma única linha.

A principal desvantagem é que os valores das features servidos em tempo de inferência estarão desatualizados, pois são um snapshot da última vez em que os jobs de ETL/ELT rodaram (e não uma representação do mundo como ele é agora). Portanto, é preciso considerar quando cada snapshot de fato teria ficado disponível ao construir o dataset de treinamento, para evitar vazamento de dados (snapshots não ficam disponíveis instantaneamente; podem levar horas e às vezes falham).

Como acontece com todas as estratégias de implementação de features “shift-left”, isso pode ficar caro, porque você pré-computa features para todas as entidades (por exemplo, clientes, contas), mesmo que apenas uma minoria delas precise ser escorada em um determinado dia.

Features tabulares on-demand via Flink e Pinot

Features on-demand são calculadas em tempo de serving, em vez de pré-computadas com antecedência. Isso é útil quando os modelos precisam das informações mais atualizadas e pré-computar features para entidades que talvez nunca as usem seria um desperdício.

Elas podem aumentar a latência se não tomarmos precauções, mas há vantagens claras:

• São o mais próximo que conseguimos chegar de features verdadeiramente online quando precisamos calcular agregações em tempo de inferência
• Só são calculadas quando realmente são usadas (ao contrário das features pré-computadas, que são materializadas para toda a base de clientes, caso venham a ser necessárias)

Uma das estratégias que usamos para features on-demand é a de escrita dupla (dual-write), na qual microsserviços transacionais escrevem eventos tanto no caminho offline (arquivos Parquet no S3) quanto no caminho online (um app Apache Flink e depois o Pinot), como podemos ver na Figura 2 abaixo.

Figura 2: Todo evento comportamental é roteado para o caminho offline e para o caminho online. Isso ajuda a reduzir o training-serving skew neste caso, mas a lógica de agregação ainda é diferente no caminho offline (geralmente jobs Spark) e no caminho online (queries SQL no Pinot).

No caminho de escrita, a estratégia de escrita dupla acima tem o objetivo de enviar os mesmos eventos para os armazenamentos offline e online. Isso é importante para reduzir o training-serving skew (mas não para eliminá-lo, porque ainda pode haver atrasos e eventos faltantes). Além disso, a lógica de agregação em batch (usada no treinamento) e as queries usadas em tempo de serving são artefatos diferentes. Por esse motivo, o monitoramento de skew continua sendo necessário.

O caminho de leitura consome os dados produzidos pelo caminho de escrita (Figura 2). Ele tem dois ramos: um para treinamento e outro para serving.

A lógica de tempo de treinamento geralmente é escrita como jobs de processamento em batch, pois envolve joins e agregações pesadas sobre horizontes de tempo longos. A lógica de tempo de serving é uma query SQL executada em tempo real em um banco de dados de processamento analítico online (OLAP) em tempo real, como o Apache Pinot, como pode ser visto na Figura 3 abaixo. Isso faz parte da nossa stack “Avalanche”, sobre a qual escrevemos neste post do blog.

Figura 3: Em tempo de treinamento, construímos o dataset de treinamento a partir de um spine (um dataset com IDs de entidades e timestamps de referência) e do event store, por meio de um job Spark. Em tempo de serving, executamos a agregação em tempo real, com uma query SQL comum, em um backend de banco de dados Pinot.

A query SQL enviada ao Pinot é uma query de agregação. Aqui está um exemplo da query que usaríamos para recuperar o número de vezes que um determinado cliente abriu o app nas ú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

Alguns pontos para tornar isso viável na nossa escala:

• Sem joins: Normalmente não fazemos nenhum join em tempo de serving (apenas agregações group-by na mesma tabela). Se precisamos de mais dados na tabela, geralmente os enriquecemos na camada do Flink (para não precisarmos de joins custosos em tempo de query)
• TTL agressivo: Devido à nossa escala, aplicamos um time to live (TTL) aos eventos no Pinot, para que o tamanho do banco de dados não cresça demais.
  • Portanto, buscamos apenas dados de curtíssimo prazo (por exemplo, os últimos dias) no Pinot. Se precisamos de features de prazo mais longo (por exemplo, agregações do último ano), nós as pré-computamos a partir de dados em batch, pois é mais fácil, rápido e barato recuperá-las de um snapshot defasado calculado diariamente.
• Evitar mudanças de schema: Mudanças de schema são um problema. Uma solução é ter campos string com conteúdo JSON, já que a estrutura do JSON pode evoluir sem alterar o schema da tabela (as mudanças ainda precisam permanecer retrocompatíveis para os consumidores existentes).

Uma vantagem secundária dessa abordagem é que qualquer outro time pode usar os eventos que um time salva: toda a lógica da feature está contida na query, então vários times podem compartilhar o banco de dados Pinot e rodar queries diferentes nele.

Isso não é uma feature store completa; usamos essa configuração apenas para suportar o caminho de serving de features online. O código de treinamento está em um caminho totalmente diferente, e por isso precisamos continuar monitorando o training-serving skew.

Usando uma arquitetura de materialização e serving em dois estágios baseada em Flink

Nesta estratégia, os eventos comportamentais fluem por dois apps Flink, um após o outro (Flink app #1 e Flink app #2, como pode ser visto na Figura 4 abaixo). Isso é útil quando os modelos precisam de features atualizadas, mas calculá-las em tempo de serving adicionaria latência demais:

Figura 4: Neste exemplo, todo novo evento “purchase-requested” é transmitido pelo Flink app #1, e uma contagem contínua das últimas 24 horas para cada cliente é mantida no Flink app #2 e então enviada via upsert ao DynamoDB. Sempre que qualquer modelo em tempo real precisa dessa informação, basta chamar o DynamoDB passando o ID do cliente e o nome do feature set para recuperar a contagem materializada mais recente daquele cliente.

Flink app #1 escuta eventos comportamentais (logs de CDC de bancos de dados transacionais e eventos brutos do Kafka), combina dados de vários shards em um único stream, filtra eventos indesejados e normaliza nomes e schemas. Nenhum estado é mantido neste primeiro app.

Flink app #2 escuta os eventos do primeiro app e mantém algum estado, guardando todos os eventos de cada cliente, com um TTL de 24 horas. Quando um novo evento entra no app #2, todas as features relacionadas são recalculadas e uma nova linha é enviada via upsert para uma tabela do DynamoDB (indexada por um ID de entidade, como o ID do cliente).

Essa estratégia depende da pré-computação de features, então sua latência em tempo de serving é muito baixa (já que é uma simples busca). Assim como no primeiro exemplo deste artigo, o trade-off é que a maioria das features será calculada para clientes que nunca precisarão delas.

Como no exemplo anterior, essa configuração cobre apenas o caminho de serving, então não é uma feature store completa. O código de treinamento roda em Spark, então pode haver skew induzido por lógica, que precisa ser monitorado.

Próximos passos

Além de melhorar constantemente nossa configuração e aprender novas lições, há pelo menos outras duas frentes que estamos explorando atualmente:

• Pipelines Feldera: Recentemente, temos experimentado uma ferramenta chamada Feldera. Os pipelines do Feldera suportam manutenção incremental de views, o que significa que features baseadas em agregação (sum, count, avg etc.) são atualizadas em tempo real à medida que novos eventos chegam, sem recalcular toda a agregação do zero a cada vez).

Além disso, o Feldera pode nos permitir unificar a lógica de features para treinamento e serving, reduzindo problemas de training-serving skew e aproximando o Nubank de uma feature store em tempo real completa e para toda a empresa.

• Features sequenciais: Embora a maioria dos exemplos desta série envolva features tabulares, também temos usado features sequenciais em diversos casos de uso, desde 2023.

Features sequenciais destravam muito valor, especialmente em redes neurais, mas também introduzem complexidade extra e novos modos de falha; descobrir como integrá-las às nossas plataformas de features em tempo real é um desafio em andamento.

Conclusão

Neste post, apresentamos alguns padrões arquiteturais que usamos no Nubank para disponibilizar features para modelos de ML em tempo real.

• Features em batch são simples e baratas e não devem ser negligenciadas, mesmo quando isso significa que o modelo em tempo real usará informações ligeiramente desatualizadas.

• Decidir entre pré-computar features ou calculá-las on-demand é um dos principais trade-offs ao pensar em feature stores. Casos de uso diferentes pedem abordagens diferentes. Uma boa feature store deve suportar ambas.

• Unificar os caminhos de definição de features de treinamento e de serving é difícil e, embora não resolva completamente o training-serving skew, o mitiga substancialmente.

Isso encerra nossa série de 3 partes sobre feature stores para ML em tempo real no Nubank. Não deixe de ler as partes 1 e 2  e fique de olho em outros posts sobre o tema, já que é algo em que estamos trabalhando ativamente!