Taming the ML Firehose: Scaling Feature Consistency
Staff Software Engineer
Senior ML Engineer
Senior Software Engineer
Introduction
Figure1: ML feature life cycle.
Our goal is to make the features used at inference the same features that are used to train the next model iteration. The single source of truth for features leads to stronger model performance and faster detection and remediation of issues. This blog describes how we scale feature consistency at Uber with a feature logging framework.
The Problem
Figure 2 shows how, in Uber’s ML environment, inconsistencies between how online and offline training pipelines compute and consume features can introduce regressions in ML model performance.

Figure 2: How online inference and offline training pipelines compute and consume features.
Online Path
Figure 2: How online inference and offline training pipelines compute and consume features.
In the online inference environment, the following features are passed into the model for serving predictions:
- User session info like user_id and language
- Store info like restaurant metadata
- Precomputed behavioral features like past store clicks, orders, and so on
During training, models learn from a fixed vocabulary, for example, languages encoded as en-US or fr-FR. At serving time, however, those same features may arrive in a different format, such as en instead of en-US. This mismatch means the serving layer is passing values the model was never trained on. As a result, strong signals silently degrade: feature distributions shift, and prediction stability suffers without obvious failures.
Offline Path
In contrast, offline training aggregates historical data from multiple logging sources:
- App logs
- ML feature Apache Hive™ tables
- Click logs
These are transformed into offline feature tables, which become the input for model training. The final dataset includes:
- User session info
- Store info
- ML features
- User actions (labels) such as did_click, did_order
The offline feature computation pipeline began generating values with different formatting conventions, for example, using _ instead of -. This results in training data containing values like jp_JA and online data containing values like jp-JA. This subtle formatting mismatch causes the model to learn from categories that never appear during serving, leading to yet another source of drift.
Further issues include:
- Fragile ETL lineages. Training data was assembled from many interdependent ETL jobs, which might have missing partitions or upstream changes that can silently degrade training datasets.
- Freshness gaps. Some critical features had multi-day latency from production changes to training visibility, limiting how quickly models could adapt.
- Developer productivity loss. Developers would have to spend several weeks uncovering the issues with ETL and freshness.
Our Solution: Feature Logging
Figure 3: How online inference logs features and offline training pipelines consume those features.
Scale Challenges
When implementing this framework, we faced some challenges with scale.
Bandwidth and Volume Constraints
Transformer-powered recommendation systems can generate prediction traffic at an extremely high throughput of 8 million QPS. If feature logs were collected for every single prediction across a full day or month, this would translate into massive data volumes for our use case at Uber on the order of hundreds of billions of records per day and 5 trillion rows per month.
Even before considering compression, the required storage footprint would reach multi-petabyte scale with around 1.7 PB per day, and the bandwidth needed to ship this data from the online prediction service to downstream pipelines would be costing us in the order of prohibitively high, multi-million-dollar infrastructure costs over time. When estimating end-to-end infrastructure cost (network, Kafka™, storage), the numbers quickly grow beyond what’s practical, even for logging just the basic feature sets.
Overly Long Feature Names
Feature names within the system tend to be verbose and highly descriptive, often including namespaces and multiple levels of aggregation. For example, store_unique_identifier_operational_meal_period_context_key. Transmitting these long strings for every prediction dramatically increases the payload size.
For large-scale models like the Uber Eats Restaurant Recommendation Model, the estimated outbound data rate for sending all features can grow to roughly 9.3 GiB per second, far beyond the capacity of existing Kafka clusters.
Unnecessary Features
Clients supply many more features than a model actually uses, resulting in bloated request payloads to the inference service. These duplications contribute significantly to bandwidth and storage requirements.
Not All Predictions Are Useful for Training
In recommendation systems, scoring happens on a batch of thousands of candidates and only a fraction of those, around 40%, of candidates are left after filtering and ranking. Only 5% of those ranked candidates ultimately become impressions on a person’s device. This means:
- Most scored/ranked candidates never actually get in front of people
- Logging all of them is wasteful for training data pipelines
- The majority of stored data offers little incremental model value
Selective logging is therefore much more efficient and cost-effective than logging the full prediction universe.
Online Architecture Deep Dive
Let’s look at how we designed online inference logs to address our challenges with scale.
Figure 4: How online inference logs features are generated.
Feature Allow List
Currently, prediction requests often include many features that the model doesn’t use. By implementing a feature allow list, we log only the necessary features fetched from the feature store. This approach:
- Reduces payload size by 4–5×
- Significantly lowers Kafka capacity requirements
- Streamlines the feature logging pipeline
Feature Name Aliasing via Enums
Instead of transmitting raw strings like store_unique_identifier_operational_meal_period_context_key across Kafka network pipes billions of times a day, the platform automatically maps each feature name to a compact, deterministic integer ID (an enum value) during serialization at the inference endpoint.
Impression Filtering
Not all predictions contribute to training. Only a fraction of candidates generated by inference are seen by users. To reduce unnecessary storage and computation, we use Flink to join prediction data with client events, where joins are performed on Kafka streams and time-windowed joins handle event delays (like user impressions arriving minutes after predictions). Based on analysis, we hold only the minimum required window in memory to a few minutes to account for the 90th percentile of session-to-impression time. This approach reduces memory requirements for the number of predictions to hold for joining against client events.
Scaling a production grade Flink job taught us that distributed systems are rarely limited by compute alone. The hardest challenges were hidden beneath the surface: state management, correctness, observability, and gradual optimization.
Scaling Requires Profiling
Increasing parallelism alone doesn’t solve performance bottlenecks. By profiling individual Flink operators, we identified several bottlenecks and tuned each stage of the Flink job independently. This reinforced that every operator has unique scaling characteristics and must be optimized based on data. We ended up allocating 512 parallelism to pre-join operators and 768–the largest Flink parallelism scale deployed at Uber–to join operator so that it has the maximum throughput available to do impression filtering.
State Management
We were using RocksDB for job state management. The default checkpointing strategy we used led to checkpoint size growing on the order of over 12 TB/hour that our writes for the state to our cloud provider started failing. We couldn’t use the default checkpointing strategy. As traffic grew, RocksDB state became a major contributor to latency, so we pivoted to a custom state management strategy. We reduced the state footprint by storing only essential metadata, aggressively evicting records immediately after stream joins occurred, and deduplicating incoming data. Combined with careful state retention tuning, this significantly minimized metadata overhead leading to high throughput.
Observability Enables Better Optimization
Before making large-scale tuning changes, we invested in detailed metrics around job output rates, time-window distributions, and pipeline behavior. These insights allowed every optimization to be measured and validated, replacing trial-and-error tuning with data-driven engineering.
High Throughput Comes from Many Small Improvements
Rather than relying on a single optimization, we improved throughput through a series of targeted enhancements, including object reuse, typed payloads, configuration caching, and reduced serialization overhead. Individually these changes were modest, but together they significantly reduced CPU utilization and garbage collection pressure and helped us reduce peak consumer lag processing by around 70%.
Build Validation and Automation into the Platform
Reliable production systems require more than fast code. Automated rollbacks, alerts-as-code, deterministic validation queries, and testing beds replicating prod-like setup enabled us to evolve the pipeline safely while maintaining correctness and minimizing operational overhead.
Transitioning to a Transformer Architecture and Usage of Sequence Features
Instead of logging single-row items, the transformer-based model leads to logging features as multi-dimensional arrays of stores (like store_uuid: [taco_store, McD, burger_store...]). By default, this creates a massive problem for downstream Kafka payload limits and also handling per-store-level watermarking for the impression filtering Flink job. The system instead inspects the model’s structural dimensions via reflection caching to detect array features, and flatten array payloads element-wise into explicit, non-array key-value records before emitting them to Kafka. This ensures non-transformer model generated logs and new sequence-based logs generated by the transformer model can be co-mingled perfectly to train next-generation models without exhausting resources. Also, to scale such massive record volume across Kafka, we shard the data into multiple Kafka clusters in a round robin fashion.
Offline Consumption
Once feature logging is enabled and we start accumulating features through the Inference endpoints, we build a table that can be used to train the next iteration of the model. Once data validation is complete and enough data has been accumulated, we train the production model using this new source of data. We also retrain the production model on the old ETL pipeline generated table for the same date range for a fair comparison.
Once this newly trained model’s performance is acceptable (typically this would mean the performance is at par or better than the prod model for metrics such as AUC and MAP), we move on to train the candidate model that would be a part of the next experiment.
Conclusion
After rolling the system into search and various ranking models, we observed measurable and operational improvements:
- 0% mismatch on key features that previously experienced mismatch rates of over 10%
- Freshness improvements, with SLAs for many priority features improved from multi-day latency down to hours, enabling faster model iteration
Through feature allow lists, name optimization, and selective impression logging we’ve reduced feature logging overhead significantly. These optimizations not only save infrastructure cost but also improve pipeline reliability, making large-scale model training more efficient and sustainable.
Acknowledgments
The rollout of this functionality couldn’t have happened without the many team members who contributed to it. A huge thank you to Abhi Kune from the Delivery team and engineers from the Michelangelo, Storage, and Streaming teams.
Cover Photo Attribution: The “'belllerophon taming pegasus' no.2” image is covered by a CC BY 2.0 license and is credited to llahbocaj. No changes have been made to the image.
Android® is a registered trademark of Google LLC.
Apache®, Flink™, Hive™, and Kafka™ are either registered trademarks or trademarks of the Apache Software Foundation in the United States and/or other countries. No endorsement by The Apache Software Foundation is implied by the use of these marks.
iOS is a trademark or registered trademark of Cisco in the U.S. and other countries and is used by Apple under license.
Paarth Chothani
Staff Software Engineer
Paarth Chothani is a Staff Software Engineer on the Uber AI Gen AI/CoreML team in the San Francisco Bay area. He specializes in building distributed systems/Gen AI solutions at scale.
Chirag Agrawal
Senior ML Engineer
Chirag Agrawal is a Senior ML Engineer in Applied AI team, based out of Bengaluru. He specializes in building ranking and recommendation systems at scale.
Amrith M
Senior Software Engineer
Amrith is a Senior Software Engineer on the Michelangelo team in Amsterdam. He specializes in building distributed systems at scale.
Products
Company