Why AI workloads need read and write computation
AI platforms need to support a combination of read and write computed features/embeddings to effectively enable modern workloads.
Read computation: raw data transformed at inference time. For example, make an API call, get the response, transform it then pass it to inference.
Write computation: raw data transformed at event time, then stored and fetched at inference time. For example, as user interactions take place, read from an event bus and update features and embeddings in the KV or vector store.
Note that often features are computed as a combination of read and write computation. For example, in fraud detection use cases, feature pipelines like these are common:
- Fetch the risk score for a request fingerprint from an API like Sift (read computed)
- Fetch the 30 day average score for that fingerprint (write computed)
- Normalize the current risk score by the historical average (read computed)
Below we’ll go through some of the challenges and tradeoffs with both patterns.
Scalability and latency challenges with read computation
Read computation from production databases can be tricky because AI/ML workloads often have drastically different access patterns than the typical production workloads that the DB was configured to handle.
For example, aggregations or embeddings generated over long windows of activities can put significant read pressure on production databases, either degrading SLAs or limiting the data that teams can use in their models.
Latency is the other main concern, especially for read time aggregation over large streams or datasets with potential hot keys.
This is where write computation comes in. By ingesting mutation events and using those to update transformed values in the feature/embedding store in realtime, read patterns for inference no longer put any direct pressure on the production database. Latency is also kept constant for arbitrarily large window sizes (assuming compaction is correctly handled).
This pattern allows AI/ML teams to fully leverage production data for transformation while completely isolating the serving infrastructure, however, managing the realtime infrastructure to transform mutations into features and embeddings poses a number of challenges:
- Large window sizes and hot keys require compaction for state management and latency SLAs.
- Mutations can include updates and deletes that need to be accounted for in aggregations and transformations.
- For windowed aggregation, need to handle events falling out of the tail of the window in addition to new events being appended,
Handling these requirements can quickly lead to complex and difficult to maintain data pipelines without a framework that offers the right level of abstraction.
Iteration Speed
The first step that ML practitioners take when creating a new feature is to generate data for model training and evaluation.
Features coming from external sources like API calls usually need to be generated from fetching-and-logging, because the raw data isn’t available in the warehouse for backfilling transformed feature values. This has the downside of needing to wait sufficient time to generate enough data volume for training, unless a historical dataset can be obtained.
Note that features from external sources also need to be computed at read time, because the concept of write time aggregation simply doesn’t exist for these sources. The same is true for contextual features (i.e. features extracted from the request context that is triggering inference).
Features from event streams, databases, or the data warehouse, however, can be backfilled directly, although users need to be careful to do so in a way that guarantees parity with what’s seen at inference time. This has the advantage of greatly accelerating the iteration speed of new feature development, since long ranges of training data can be immediately generated, tested, and iterated on.
The challenge here is generating point-in-time-correct (PITC) data for training (i.e. data that is temporally consistent with values seen online at inference time). Tools such as Chronon that handle this for you can greatly accelerate iteration.
Zipline
Given the advantages of write computation (namely scalability and iteration speed), Zipline makes it easy to define write computed features on supported data sources, combine them with read computed features, and chain them into further transformations and aggregations.
Transformations defined on data sources like CDC, event streams and warehouse tables are automatically computed at event time and uploaded to scalable low-latency storage. Zipline computed PITC backfills of these features instantly for training data generation.
At inference time, external features are fetched in parallel via fanout, passed into the rest of the pipeline for read computation, and logged. Zipline then joins those logs to backfilled feature data in a log-stitching job, completing the training feature vector.