Batch Inference

Use Lightning prediction with rf.Writer to write results directly to Parquet. This keeps large jobs from accumulating every prediction in memory.

import lightning.pytorch as lit
import pyarrow.dataset as ds

import relflow as rf

# model is trained; request files are supplied by the application.
data = rf.ArrowDataModule(
    model=model,
    predict=ds.dataset("warehouse/requests", format="parquet"),
    retain=("request_id",),
)
trainer = lit.Trainer(callbacks=[rf.Writer("predictions")])
trainer.predict(model=model, datamodule=data, return_predictions=False)

Each rank writes predictions/rank-{global_rank}.parquet. The files contain the canonical inputs and predictions columns described in Prediction Output. Retain a business key to join or sort distributed results downstream; relflow does not add row identity.

With multiple ranks or data workers, the loader trims the final group to divide evenly among consumers, so trailing rows can be omitted even with drop_last=False. Use one device and num_workers=0 when every input row must be processed.

A postprocessor can give files a consumer-specific schema:

writer = rf.Writer("warehouse-predictions", postprocessor=warehouse)

Here warehouse is an application-defined decorated processor. The first written batch fixes that rank’s Parquet schema; later batches must match it. Processors may change rows and columns for offline output.

Open the resulting shards as one Arrow dataset for further batch processing:

result = ds.dataset("predictions", format="parquet")

Use model.predict(request_table) for bounded in-memory requests and Deployment for online requests. Direct model.predict processes its supplied input as one unit; use the data module to batch larger datasets.