Scaling Model Evaluation at Pinterest with Ray (Part 4)
Vishnu Arun | Tech Lead, Feature Importance Platform; Anderson Lam | Sr Software Engineer; Andrew Qiu | Machine Learning Engineer I; Helena Xu | Data Scientist I; Tianyou Wang | Data Scientist II
Introduction
In our multi-part blog series, we have discussed why we chose Ray as a last mile data processing framework, how we integrated Ray into our existing ML Training Infrastructure, and how we enabled efficient Batch Inference at scale. In this blog post, we will discuss how we were able to apply Ray to another critical part of our ML Infrastructure stack: Model Evaluation.
In the Ray Batch Inference Blog Post, we ended on a section on accumulators (counters that would allow us to calculate metrics like ROC-AUC and cross entropy). This post is where those accumulators pay off: a model evaluation platform that combines the strengths of Ray and Apache Spark to compute metrics efficiently at scale.
Model evaluation may appear to be a relatively straightforward step after inference or training. But as ML workflows scale, evaluation quickly becomes an infrastructure challenge. Systems that run many experiments in parallel, such as feature importance, must efficiently compute and aggregate metrics over large volumes of inference data, as shown in Diagram 1.
Later in this post, we’ll show how we amortize the metrics compute step, overlapping it with inference for minimal added runtime cost, cutting end-to-end evaluation time by up to 32% while scaling linearly as we run more experiments.
Terminology
- Feature Importance (FI): A technique for quantifying how much each input feature contributes to a model’s predictions. In this post, we will reference a perturbation-based method where we perturb one or more columns, run inference, and compare the results against an unperturbed dataset.
- Experiment / baseline: One perturbed run vs. the unperturbed run.
- Model head: One output of a multi-task model.
- Batch: an arbitrary, compute-sized group of rows. This is the unit of data that one Ray task processes.
- Bucket: a batch that is group-preserving. It contains one or more complete segments, and no segment may span buckets.
- Segment: The semantic unit over which a group-wise metric is defined, such as a request.
- Serde [serialization/ deserialization]: The process and cost of writing data out and reading it back.
Background
Traditionally, evaluation and inference are two separate stacks stitched together. Diagram 2 below contrasts that unfused pipeline against the fused pipeline we’ll build up over the rest of this post.
Concretely, this evaluation step is typically implemented as a Apache Spark job that computes metrics across one or more inference experiments. As shown in the diagram above, we previously used Ray inference to write predictions to S3, then have a follow-up compute job to reread all this data and compute the metrics.
For a FI run, this approach introduces two sources of overhead:
- Cluster startup: Acquiring and warming up a cluster for the metric job.
- Data movement: Writing predictions to S3 and reading them back to compute metrics.
Because inference is done in batches, and metrics typically depend on all of the inference results, the unfused approach can be used to calculate any metric. To save the overheads, we need to fuse the stages, but this also implies that we need to be able to calculate metrics on batches of rows. Therefore, we need a metric platform that treats decomposability as a first principle.
Metrics fall into three groups:
Metric Contract
We define a common metric contract that captures each metric’s compute layout and provides a consistent interface across execution frameworks. We express this interface concretely as an Abstract Base Class.
from abc import ABC, abstractmethod
class ReducibleMetric(ABC):
@abstractmethod
def dataset_preprocessor(self, ds):
"""Prepare the dataset to satisfy the metric's data layout requirements."""
pass
@abstractmethod
def compute_partials(self, batch):
"""Compute a partial result from one batch"""
pass
@abstractmethod
def reduce_partials(self, partials):
"""Combine partial results into the final result."""
passDecomposable and holistic metrics can calculate partials from arbitrary batches because they do not require related rows to be colocated.
Group-wise metrics introduce an additional constraint. Before handing data to Ray, we organize complete segments into buckets. A bucket is therefore a group-preserving batch: it contains one or more complete segments, and no segment crosses a bucket boundary. This lets us independently tune bucket size for efficient execution without changing the semantic unit of the metric.
Common Metrics found in Recommender Systems
We will walk through three metrics, MAE, ROC-AUC, and Precision@K; one from each group above. We will use the following dataset to compute each metric.
For the examples below:
- model_output contains the predictions, and label_continuous contains the target values.
- label_binarized is derived as 1 if label_continuous ≥ 0.5 and 0 otherwise.
- Batch A contains ReqA, and Batch B contains ReqB. This split lets us demonstrate batch-level computation and aggregation.
MAE (Decomposable)
Mean Absolute Error measures the average absolute difference between predictions and their actual values.
MAE_b = (1/n_b) · Σ_{i ∈ batch b} |pred_i − label_continuous_i|
Unlike MSE/RMSE, MAE does not penalize large errors as heavily.
Per-row absolute errors:
Global: Σ = 1.45, N = 8 → MAE = 0.18125.
Batch decomposition:
Batch A partial (Σ|err| = 0.65, n = 4)
Batch B partial (Σ|err| = 0.80, n = 4)
Merge (0.65 + 0.80) / (4 + 4) = 0.18125
class MAE(ReducibleMetric):
def dataset_preprocessor(self, ds):
# No dataset preparation needed. Rows are independent.
return ds
def compute_partials(self, batch):
# Emit summed absolute error and row counts from each batch.
error = np.abs(batch["model_output"] - batch["label_continuous"])
return error.sum(), len(error)
def reduce_partials(self, partials):
# Sum partials and divide once globally.
error, count = np.sum(partials, axis=0)
return error / countROC-AUC (Holistic)
ROC-AUC is a binary classification metric which measures the ability of a model to distinguish between classes across different decision thresholds.
Method 1 — Global Sort (what sklearn does)
Global Sort
We sort the rows by model_output in descending order, keeping each prediction paired with its corresponding binary label.
Cumulative Sum
After sorting, we calculate cumulative true-positive (TP) and false-positive (FP) counts and prepend 0 to both vectors.
Calculate the TPR and FPR
We calculate the TPR at each threshold by dividing the TP vector by the number of positive labels.
We calculate the FPR at each threshold by dividing the FP vector by the number of negative labels.
Calculate the Trapezoidal area under TPR/FPR
We calculate the AUC of TPR/FPR to get the final ROC_AUC value.
Get Pinterest Engineering’s stories in your inbox
Join Medium for free to get updates from this writer.
Sorting our eight rows by score:
With 4 positives and 4 negatives, the area works out to AUC = 0.9375 (or 15/16 positive/negative pairs are ranked correctly).
We observe that ROC_AUC using this method is calculated over the entire dataset due to a global sort. We could break the dataset into batches, calculate ROC-AUC for each batch, and average the results. However, this does not produce the same result as a global ranking.
Consider the following calculation:
- Batch A (ReqA) = 1.00
- Batch B (ReqB) = 0.75
- Average = 0.875, not 0.9375
Averaging a ranking number across batches is wrong, because AUC depends on how rows rank against all other rows, not just those in their batch.
We could force a real global sort in Apache Spark, but doing that across all metrics, model heads, and experiments is exactly the cross-stack cost we’re trying to avoid.
Method 2 — Confusion Matrix
Consider the standard method to approximate ROC_AUC.
Define thresholds
Pick a fixed set of thresholds in [0, 1].
Calculate Confusion Matrix
Calculate the confusion matrix (TP, FP, TN, FN) for each threshold
Calculate the final metric
Calculate the TPR/FPR and finally derive the AUC.
This method can be calculated batch-wise because the global confusion matrix can be calculated from the batch confusion matrices.
We can create the following confusion matrices per batch and then sum:
From this approximate example, when we derive the area under the curve, we get 0.875. If we use more thresholds, for example [0.00, 1.00] with step size of 0.01, we would get 0.9375. Note: This does not mean that threshold-based aggregation generally produces the exact global AUC. In this toy example, a step size of 0.01 recovers the exact value because the prediction scores align with the threshold grid.
class ROCAUC(ReducibleMetric):
def dataset_preprocessor(self, ds):
# No dataset preparation needed. Rows are independent.
return ds
def compute_partials(self, batch):
# Emit confusion matrix across thresholds.
pred = batch["model_output"]
truth = batch["label_binarized"].astype(bool)
positive = pred[None, :] >= self.thresholds[:, None]
return np.stack([(positive & truth).sum(1), (positive & ~truth).sum(1),
(~positive & ~truth).sum(1), (~positive & truth).sum(1)],1)
def reduce_partials(self, partials):
# Sum confusion matrices and add ROC endpoints.
# Integrate in ascending FPR order.
tp, fp, tn, fn = np.sum(partials, 0).T
tpr, fpr = tp / (tp + fn), fp / (fp + tn)
return np.trapz(np.r_[0, tpr[::-1], 1], np.r_[0, fpr[::-1], 1])Precision@K (Group-wise)
Precision@K is a common metric used for ranking systems. It measures the fraction of the top K ranked items that are relevant.
This metric is calculated within an individual segment of data, such as a session or request. A segment may be identified by a single key or a composite key.
The procedure for calculating Precision@K is as follows.
Group/Rank Dataset
Group the dataset by segment and sort by predicted score.
Window and Count
Select the top K items and count the relevant items within that window.
Calculate Metric
Divide the number of relevant items by K to calculate Precision@K for each segment. Average the per-segment scores to obtain the final result.
On our data with K = 2:
Average = (1.00 + 0.50) / 2 = 0.75.
We observe this metric has a similar decomposition pattern. We can break it down into subproblems where we look at each segment, calculate the metric, then average final scores of all subproblems.
We have two problems when trying to scale this approach:
- The subproblem (batch_size) is tied to the size of each segment. This is inefficient as each executor may receive a small and variably sized amount of data.
- If we combine segments together, we must ensure that no segments span a bucket boundary.
Our goal is to keep every row for a given segment in the same bucket while ensuring that each bucket contains enough segments to saturate the compute resources.
More succinctly, we want all rows within each segment to be colocated and each bucket to be sufficiently large. Since we are operating at the global level, we can define this as an Apache Spark procedure. It keeps all rows for each segment together, then packs complete segments into buckets large enough to use each Ray worker efficiently without splitting a segment across buckets.
class PrecisionAtK(ReducibleMetric):
def dataset_preprocessor(self, ds):
# Colocate each segment and pack into buckets.
return bucket_complete_segments(
ds,
by="request_id",
)
def compute_partials(self, bucket):
# Rank within each segment and calculate top-K score.
scores = []
for segment in group_by(bucket, "request_id"):
ranking = np.argsort(
-segment["model_output"]
)
# Assume each segment contains at least K items.
top_k = ranking[:self.k]
relevant = segment["label_binarized"][top_k].sum()
scores.append(relevant / self.k)
return sum(scores), len(scores)
def reduce_partials(self, partials):
# Sum segment scores and divide by segment count.
score, count = np.sum(
partials,
axis=0,
)
return score / countFusing Metrics with Inference
In the upper figure in diagram 3, we show a typical Sequential Inference-Metrics multi-experiment pipeline. We run model inference over all the experiments, serialize, and persist those results along with the carryover labels. Note this inference dataset scales with the input data and the number of experiments. Next, we deserialize this data and compute metrics. GPU inference is typically more resource intensive than metrics computation. This design causes the CPU workers to sit idle during inference time.
The lower figure in diagram 3 illustrates a fused approach where we pipeline the metric computation once inference completes for each bucket. The shared object store minimizes serde overhead. This enables us to effectively overlap metric computation with inference, reducing total runtime.
Results
We benchmarked two setups against a Pin Ranking Model, sweeping across three dataset sizes (Small ≈ 30K, Medium ≈ 300K, Large ≈ 3M rows) and several experiment counts:
- Unfused: Ray inference, then a separate Ray metric job.
- Fused: Ray inference with metrics in the same pass.
When we gather metrics during inference, we do not see any meaningful runtime increases. The time measured here includes startup time (acquiring the cluster) and runtime (actual computation).
Furthermore, we can see that runtime grows linearly with experiment count while startup stays flat.
As the number of experiments increases from 1 to 300 on the Large dataset, fused runtime grows approximately linearly (~43 min fixed + ~41s per experiment), while startup time remains roughly constant. The separate metric stage (~11s per experiment) is work that fusion absorbs through overlap with inference, keeping the marginal cost per experiment small and predictable.
What’s next
Other upcoming topics include:
- Ray Tune at Pinterest
- Improving the data loader performance in Ray
- Fault-tolerant Ray Batch Inference and Training
- Feature Importance on Ray
Acknowledgments
Related Pins: Shivin Thukral, Zilin Ren
AI Platform: Rubin Fergerson IV, Qingxian Lai
Data Platform Sciences: Darren Reger, Adam Obeng
Agent Platform: Tan Wang
