Advanced Configurations

For a temporal query, prediction_time is a timezone-naive Python datetime interpreted as UTC. If you omit it, Kumo uses the latest relevant timestamp in the target source data. Pass an explicit time for repeatable batch runs:

from datetime import datetime
prediction_table = query.generate_prediction_table(
    snapshot,
    prediction_time=datetime(2026, 1, 31),
    plan=trainer.training_table_generation_plan,
)

Timezone-aware values and generated anchors before the target source-data range are rejected. Anchors at the start of that range and future anchors are allowed. The prediction table retains its selected rows and anchor times, so you do not pass prediction_time again to trainer.predict(). The snapshot supplied to prediction selects graph features independently of those retained rows. For forecasts, reuse the fitted or loaded trainer’s training_table_generation_plan so the horizon, lag, year-over-year features, and step size match training. Do not pass a prediction time for static queries. For individual entities or entity-specific anchor times, use a supplied prediction table. For static list-target recommendations, the typed supplied-row example uses a graph without time columns. Adding supplied timestamps does not make a static LIST_DISTINCT query valid on a graph whose time columns require temporal aggregation. For forecasts, prepare the required sequence columns with the trained generation plan before supplying prediction rows.

You can change a query’s WHERE filter when scoring with the same fitted model. Keep the target expression, entity table and key, task, relationship paths, and forecast length compatible with training. Reuse the trained generation plan and a compatible snapshot.

For example, after fitting PREDICT LIST_DISTINCT(orders.PID) RANK TOP 3 FOR EACH users.ID on a graph with a boolean users.LABEL column, select just the matching users:

from kumo.engine import PredictiveQuery
filtered_query = PredictiveQuery(
    graph,
    "PREDICT LIST_DISTINCT(orders.PID) RANK TOP 3 "
    "FOR EACH users.ID WHERE users.LABEL = True",
)
filtered_population = filtered_query.generate_prediction_table(
    snapshot,
    plan=trainer.training_table_generation_plan,
)
filtered_result = trainer.predict(
    filtered_population,
    snapshot=snapshot,
    num_classes_to_return=3,
)
print("Selected users:", filtered_population.row_count)
print("Ranked candidates:", filtered_result.num_rows)

Use your graph’s actual table, column, and key names. Changing the prediction entity or target requires a compatible newly trained model; an incompatible query raises a target/entity contract error.

For binary classification, supply a finite threshold from 0 to 1. The TARGET_PRED value is true when True_PROB is strictly greater than the threshold.

result = trainer.predict(
    prediction_table,
    snapshot=snapshot,
    binary_classification_threshold=0.5,
)

For multiclass, multilabel, ranking, or link prediction, use num_classes_to_return to choose how many classes, labels, or candidates to retain per population row. Link prediction requires an explicit integer from 1 to 500. For the other applicable tasks, a supplied count must be a positive integer. Use this example only with one of those tasks; binary classification and regression do not accept a class count.

prediction_options["num_classes_to_return"] = 5
result = trainer.predict(
    prediction_table,
    snapshot=snapshot,
    **prediction_options,
)

For multiclass and multilabel tasks, an omitted count uses the query’s TOP value when present. Without that value, multiclass classification defaults to at most 20 classes, and multilabel classification and ranking return the available labels. The number returned can be smaller than the count requested when fewer candidates are available. See the worked recommendation and multilabel examples.

Use num_workers to choose how many scorers process the selected population. The default is 1, which runs one scorer in the current process. Reuse the fitted trainer, prediction_table, snapshot, and task-specific prediction_options from the Prediction Guide. For link prediction, first add a positive num_classes_to_return to prediction_options.

result = trainer.predict(
    prediction_table,
    snapshot=snapshot,
    num_workers=1,
    **prediction_options,
)
print("Output rows:", result.num_rows)

Set num_workers to a positive integer. Zero, negative values, booleans, and non-integer values raise ValueError. When the prediction population count is available, Kumo caps the effective scorer count at the number of population rows. An empty population produces an error rather than launching scorers.

With more than one effective scorer, Kumo spawns separate processes that each load their own model and encoders. This setting controls scorer processes, not DataLoader workers. In Python scripts, put executable setup and prediction calls under if __name__ == "__main__": so spawned processes can import the script without starting the workflow again. For CUDA prediction, each concurrent scorer needs a separate visible GPU with enough memory for its model, encoders, and prediction batches. For example, use num_workers=2 only when two GPUs are visible and available for the run. If Kumo reports too few visible CUDA devices, reduce num_workers or correct GPU visibility before retrying. More workers do not guarantee faster prediction.

Custom embedding classes used by multiple scorers must be importable by the spawned processes. Move classes defined in an interactive notebook or session into an importable .py module, or keep num_workers=1. Prediction still returns a PredictionResult after the scorers finish.

The returned PredictionResult is backed by local Parquet files:

Use result.frame only when the complete output fits in memory. For large outputs, process the Parquet files or write directly to a durable destination.

Output columns depend on the task:

Output order is not guaranteed. Join results to other data by entity and, where applicable, time and class. Forecast output includes ENTITY, TIMESTAMP, and TARGET_PRED by default. TIMESTAMP identifies the start of each forecast window. To include the prediction anchor as ANCHOR_TIMESTAMP, import MetadataField from kumoapi.jobs and set extra_fields=[MetadataField.ANCHOR_TIMESTAMP] in the PredictionOutputConfig passed to trainer.predict(). This optional column does not replace the forecast window timestamp. The Linux forecasting example selects this metadata and checks both timestamps in its saved output.

Choose fields before scoring. Reuse the output connector and a new batch path or table name from the basic workflow. This example requests scoring time:

from kumoapi.jobs import MetadataField
result = trainer.predict(
    prediction_table,
    snapshot=snapshot,
    output_config=PredictionOutputConfig(
        connector=output_connector,
        artifact_type=PredictionArtifactType.PREDICTIONS,
        table_name=output_name,
        write_mode=WriteMode.OVERWRITE,
        extra_fields=[MetadataField.JOB_TIMESTAMP],
    ),
    **prediction_options,
)

Add MetadataField.ANCHOR_TIMESTAMP when the population carries temporal anchors. Omitting extra_fields, or supplying an empty list during prediction, omits both optional columns. Exporting an existing result preserves its already generated columns and cannot add these fields.

Configure copy_connector for your chosen destination using the same setup and grants as the basic workflow, then choose a new copy_name. Use a new path or table to preserve earlier batches:

result.export(
    PredictionOutputConfig(
        connector=copy_connector,
        artifact_type=PredictionArtifactType.PREDICTIONS,
        table_name=copy_name,
        write_mode=WriteMode.OVERWRITE,
    )
)

result.export() returns None. It does not change result.files or make the result object usable after its local files are removed. Keep those files until every export finishes.

The copy preserves scores, columns, and the original timestamp values. A nonempty extra_fields list is rejected, even if those fields already exist. To choose different metadata, score again with the desired configuration.

For compatible warehouse tables, see Append to a Warehouse Table. File, S3, and GCS destinations require OVERWRITE and follow the file destination rules.

Select PredictionArtifactType.EMBEDDINGS to export embeddings only if the original run generated them. Export checks requested artifacts before writing; it cannot derive missing embeddings from saved predictions. A failed write may leave partial output. Inspect the destination and use a new batch location for a retry when you need to preserve completed output.

Reuse the output connector and table name from Choose an Export Location. Append creates the table if it does not exist:

result.export(
    PredictionOutputConfig(
        connector=output_connector,
        artifact_type=PredictionArtifactType.PREDICTIONS,
        table_name=output_name,
        write_mode=WriteMode.APPEND,
    )
)

For an existing table with no concurrent writers, compare its counts before and after export: the increase should equal result.num_rows. For a new table, compare the final count. With concurrent writers, verify this batch by keys and metadata or use a separate batch table.

In a Databricks notebook, a Spark SQL readback uses the fully qualified SQL name, such as `catalog`.`schema`.`customer_predictions_batch_001`, rather than the encoded write name. The Spark identity needs USE CATALOG, USE SCHEMA, and SELECT independently of the output connector’s write grants.

Snowflake overwrite replaces rows through a staging table and retains existing destination grants; the destination column names must match. Kumo automatically retries an expired-credential file export only for OVERWRITE, not for APPEND.

Set embedding_name to a separate table or path in the output connector, using the destination naming rules. For Databricks, use f"{output_connector.schema_name}__kumo__customer_embeddings_batch_001" for writes and the bare customer_embeddings_batch_001 for readback. For file and object storage, use Parquet to retain array types. Binary embedding calls still require a threshold; link embedding calls do not require a candidate count.

embedding_options = {}
if query.get_task_type() == TaskType.BINARY_CLASSIFICATION:
    embedding_options["binary_classification_threshold"] = 0.5
embedding_result = trainer.predict(
    prediction_table,
    snapshot=snapshot,
    output_config=PredictionOutputConfig(
        connector=output_connector,
        artifact_type=PredictionArtifactType.EMBEDDINGS,
        table_name=embedding_name,
        write_mode=WriteMode.OVERWRITE,
    ),
    **embedding_options,
)
print("Embedding rows:", embedding_result.num_rows)

Embedding width depends on the trained model. Node-task embeddings include ENTITY and EMBEDDING, plus any selected timestamp metadata. Link-task embeddings use ID, TYPE (the entity table), and EMBEDDING, and can include both source and candidate entities. Links within one entity table also include NAME to distinguish the two roles. Use these identifiers and any temporal coordinates to associate vectors with their task context; row order is not guaranteed.

For FileConnector, S3Connector, and GCSConnector destinations, choose the format in table_name, or in each destination value when using output_paths. Pass the configuration to trainer.predict() when scoring, or to result.export() to copy an existing result to one of these destinations. The returned result still references local Parquet files, even when the destination copy uses CSV.

A trailing slash requests a directory; the suffix before it selects the format. Without a recognized suffix, the format defaults to Parquet and a single-file destination receives a .parquet suffix. A name ending in .csv or .parquet can therefore identify a directory when the result has multiple parts. Retain every part and the directory’s markers together. Do not read only the first part or rename the directory as though it were a single file.

Use CSV for scalar columns when the consuming tool expects text. One-dimensional vectors become space-separated strings, such as 1.0 2.0, rather than typed arrays. Use Parquet for embeddings and nested arrays to retain their array types. Nested values may raise a CSV conversion error or be written as text, depending on their representation. A successful CSV export does not guarantee that a nested array can be reconstructed faithfully. If conversion fails, choose a new Parquet destination and export the existing result while its local files remain available.

For a small Databricks result, you can also retain a Parquet copy in the Unity Catalog volume configured as ARTIFACT_ROOT in Databricks setup. This collects the full result in memory; use the warehouse output above for larger results. Choose a new directory:

from uuid import uuid4
run_root = ARTIFACT_ROOT / uuid4().hex
run_root.mkdir(parents=True, exist_ok=False)
result.frame.to_parquet(run_root / "predictions.parquet", index=False)
print("Saved prediction file:", run_root / "predictions.parquet")

The file’s lifetime and access follow the volume’s policies.

These file destinations support WriteMode.OVERWRITE; WriteMode.APPEND is rejected. For each prediction batch, choose a new destination path or an empty directory. OVERWRITE does not authorize replacing a nonempty directory belonging to another export. Kumo records directory ownership in _KUMO_EXPORT and writes _SUCCESS only after that export finishes. Do not remove ownership markers to force a write into an occupied directory.

An internal retry of the same export can reuse its owned directory, but this does not mean a new prediction job can reuse that path: prediction checks can reject an occupied destination before scoring. If you receive an output-directory conflict, choose a new batch path. If an export fails after writing some parts, treat that destination as incomplete and retry to a new path as described in Export an Existing Result.

For a local run, use output_paths to save both artifacts from one prediction call. The same extra_fields list applies to every selected artifact; it cannot select different metadata for predictions and embeddings within one call. This output_paths configuration requires a filesystem or object-storage connector; warehouse connectors do not accept it. Give the artifacts distinct, nonoverlapping destinations that do not resolve to the same location. Identical paths raise ValueError; choose a separate file or directory for each output. Use output_paths instead of artifact_type and table_name in this configuration. The same configuration can be passed to combined_result.export() afterward, with a new connector or distinct paths and no extra_fields, to copy both already generated artifacts without rescoring. Writes run sequentially without a transaction across destinations; inspect each destination if one fails.

The example below includes scoring time and the generated prediction anchor in both outputs. Reuse trainer and snapshot for a fitted binary classification model, with a generated prediction_table for its temporal query. Reuse prediction_options from the Prediction Guide with binary_classification_threshold=0.5. Replace the root path with a writable retained location and use a new batch directory to preserve earlier results:

from pathlib import Path
from kumo.engine import FileConnector, PredictionOutputConfig
from kumoapi.jobs import MetadataField, WriteMode
output_root = Path("/path/to/retained-results/batch-with-embeddings")
combined_result = trainer.predict(
    prediction_table,
    snapshot=snapshot,
    output_config=PredictionOutputConfig(
        connector=FileConnector(root_dir=str(output_root)),
        output_paths={
            "predictions": "predictions/",
            "embeddings": "embeddings/",
        },
        write_mode=WriteMode.OVERWRITE,
        extra_fields=[
            MetadataField.JOB_TIMESTAMP,
            MetadataField.ANCHOR_TIMESTAMP,
        ],
    ),
    **prediction_options,
)
print("Prediction rows:", combined_result.num_rows)
embedding_files = sorted((output_root / "embeddings").glob("*.parquet"))
print("Embedding files:", embedding_files)

The returned PredictionResult selects predictions when both artifacts are requested: its num_rows, files, and frame describe the prediction output. Read embeddings from their separately configured destination, using the listed embedding_files. For an embedding-only call, those result attributes describe embeddings instead.

Both outputs contain the selected JOB_TIMESTAMP and ANCHOR_TIMESTAMP columns alongside their task-specific columns. They share the scoring start time because they come from the same call. Match entities and temporal coordinates when associating predictions with embeddings; do not rely on row order.