Prediction callbacks

Prediction callbacks observe placed batches and model outputs on the Python host. Use them for streaming or artifact writing rather than changing the numerical prediction step.

Prediction writer

PredictionWriter writes the canonical NPZ after joining host outputs and can also write a MATLAB file. A Trainer accepts at most one prediction writer. Projects may include it only for prediction runs or reuse the same Trainer and in-memory state after fitting.

The writer receives application metadata through PredictionContext. SIGTERM ends the active task after cleanup; project entrypoints decide whether a later prediction task should run.

PredictionWriter [source]

PredictionWriter(
    output_dir: str | Path | None,
    *,
    save_file_name: str = "predictions",
    save_mat: bool = False,
    mat_field_names: Mapping[str, str] | None = None,
    rank_zero_only: bool = True,
)

Bases: Callback

Write reconstructed prediction artifacts after a prediction lifecycle completes.

The writer keeps persistence outside the model and prediction entrypoint. It receives the concatenated host predictions and immutable reconstruction pool through :class:PredictionContext. The artifact serializer applies configured output scales so NPZ and optional MATLAB values share physical units.

Attributes:

Name	Type	Description
`output_path`		Resolved canonical NPZ destination.
`artifact_path`	`Path \| None`	Saved canonical path after successful prediction, otherwise `None`.

Initialize prediction artifact policy.

Parameters:

Name	Type	Description	Default
`output_dir`	`str \| Path \| None`	Directory receiving the canonical NPZ and optional MATLAB sidecar.	required
`save_file_name`	`str`	Suffix-free filename stem without parent-directory components.	`'predictions'`
`save_mat`	`bool`	Whether to write a MATLAB sidecar with the same stem.	`False`
`mat_field_names`	`Mapping[str, str] \| None`	Optional generic-to-MATLAB variable-name mapping.	`None`
`rank_zero_only`	`bool`	Whether only global rank zero may write artifacts.	`True`

Raises:

Type	Description
`ValueError`	If `output_dir` is absent or `save_file_name` is not a suffix-free local filename stem.

setup [source]

setup() -> None

Reset the artifact result before a new prediction lifecycle.

on_predict_end [source]

on_predict_end(context: PredictionContext) -> None

Reconstruct and save the completed prediction on the permitted process.

Parameters:

Name	Type	Description	Default
`context`	`PredictionContext`	Final prediction context containing concatenated host outputs and the reconstruction pool.	required

Raises:

Type	Description
`RuntimeError`	If prediction collection was disabled or the prediction source exposed no host pool.