Skip to navigation

Batch Prediction

View as Markdown

materialize() gathers more entities than one call scores as several calls automatically — you do not split indices into chunks yourself for an ordinary scoring job.

from sdm_client import PredictiveQuery, SDMClient, KumoRelational
query = PredictiveQuery(
"PREDICT COUNT(orders.*, 0, 30, days) FOR EACH customers.customer_id",
schema,
)
with SDMClient(url="http://localhost:8000", max_retries=3) as client:
request = query.materialize(indices=customer_ids) # any size up to 100,000
result = client.predict(request, model=KumoRelational())

materialize() builds one batch per 10,000 entities (200 for a temporal link prediction task), each with its own sampled context. predict() sends one call per batch and merges the answers into one PredictionResult. row_index counts across the whole request, so to_pandas(), which orders by it, gives the entities back in the order you asked for them — not interleaved by batch.

What is still a size limit

LimitValueWhat happens past it
Entities in one materialize() call100,000, in at most 100 calls of up to 10,000 (or 20,000 at 200/call for temporal link prediction)Refused before anything is sampled, naming which ceiling you hit
Request body30 MiB per call, 16 MB on Databricks Model ServingRefused at materialize() (30 MiB) or on the first call (16 MB gateway limit)

materialize() builds every batch, each with its own context, before the first is sent — the whole population is held in memory at once, which is why the ceiling exists; 100,000 entities is roughly a gigabyte of it.

For more than 100,000 entities, chunk indices yourself and call materialize()/predict() once per chunk:

import itertools
def chunks(items, size):
it = iter(items)
while chunk := list(itertools.islice(it, size)):
yield chunk
results = []
with SDMClient(url="http://localhost:8000", max_retries=3) as client:
for chunk in chunks(customer_ids, 100_000):
request = query.materialize(indices=chunk)
results.append(client.predict(request, model=KumoRelational()))

Tune the per-call size with rows_per_call

rows_per_call lowers the per-call size below the endpoint’s own cap — it never raises it. Reach for it when a wide graph’s context pushes a batch over the 30 MiB body, or when your deployment’s own cap is lower than this client’s default:

request = query.materialize(indices=customer_ids, rows_per_call=1000)

It is not the fix for every body-size refusal: a batch carries its context as well as its rows, and where the context alone is already over the limit, no batch size is small enough. There, reduce num_context_samples, num_neighbors, or the columns the schema carries instead.

Ranking and forecasting are not batched the same way

A RANK TOP K query makes one call per candidate, and a forecast one per step — bounded separately by sdm_client.decompose.MOST_PARTS (200), not by the entity limits above. More entities than one call scores is refused outright for either, before anything is gathered, because a ranking can’t be batched: a candidate reported unscored applies to the whole ranking, not one batch, and a batch whose own context held only one class would wrongly drop it from every row. See Prediction Types and Forecasting.

A failed batch keeps what it already paid for

If a call in the middle of a batched request fails, the ones that already answered are not discarded. The error carries them as details['answered'] and the rows they cover as details['rows_answered'] (or details['parts_answered'], by step, for a forecast), so you can retry only what’s left instead of paying for the whole population again.

Sessions and batching

fit() takes a request of a single call — a session pins one context, and a batched request has one context per batch. Use a small initial request to create the session, then score a larger population against the handle.

materialize(indices=...) selects prediction rows; it samples labeled context separately from the schema. Even when the initial and later prediction ID lists are disjoint, later prediction rows can still appear in that sampled context with their answers attached.

Before creating the session, identify every customer it may score and pass those IDs as holdout to fit(). This removes matching rows from the context before it is pinned:

small_request = query.materialize(indices=labeled_customer_ids[:100])
with SDMClient(url="http://localhost:8000") as client:
with client.fit(
small_request,
model=KumoRelational(),
holdout=customer_ids_to_score,
) as fitted:
big_request = query.materialize(indices=customer_ids_to_score)
result = fitted.predict(big_request)

Pass the full future scoring population, not just IDs that overlap the initial request. Omit holdout only if you have independently established that all future scoring rows are excluded from the pinned context.

Keep enough other labeled context after exclusion. If no usable context remains, prepare a separate labeled context or use ordinary client.predict() calls instead of a fitted session. See Holdout for the exclusion rule and Python SDK API for the session interface.

Production jobs

SDMClient(max_retries=3) retries a transient transport failure (429, 502, 503, 504, or a dropped connection) automatically, per call. It does not retry 500, since the endpoint can use that status for a request or model error rather than a transient failure. Record the prediction settings in durable state for a production job: a retry after a timeout can run inference again, so downstream writes should be idempotent and keyed by id (or row_index when there is no entity key).