> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/sdgm/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/sdgm/_mcp/server.
# Advanced Configurations
> Customize prediction populations, task output, timestamps, parallel scoring, exports, and embeddings
Continue from the [Prediction Guide](/kumo-engine/1.0.0/user-code-guide/making-predictions) with `trainer`, `query`, `snapshot`, `prediction_table`, and your output configuration. Customize population filters, task output, timestamps, parallel scoring, exports, and embeddings. Select optional metadata before scoring and retain local result files until all exports finish.
## Population and Anchor Options
Set temporal anchors and use supplied prediction rows
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:
```python
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](/kumo-engine/1.0.0/user-code-guide/train-a-model/predictive-queries#use-supplied-prediction-rows).
For static list-target recommendations, the [typed supplied-row example](/kumo-engine/1.0.0/user-code-guide/train-a-model/predictive-queries#use-supplied-training-rows) 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](/kumo-engine/1.0.0/user-code-guide/train-a-model/predictive-queries#supply-forecast-sequences) with the trained generation plan before supplying prediction rows.
## Change the Scoring Filter
Change the Scoring Filter
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:
```python
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.
## Task-Specific Output Options
Choose thresholds and returned candidates
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.
```python
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.
```python
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](/kumo-engine/1.0.0/user-code-guide/code-examples/linux-deployment/product-recommendations) and [multilabel](/kumo-engine/1.0.0/user-code-guide/code-examples/linux-deployment/multilabel-classification) examples.
## Parallel Prediction
Parallel Prediction
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](/kumo-engine/1.0.0/user-code-guide/making-predictions#generate-predictions).
For link prediction, first add a positive `num_classes_to_return` to `prediction_options`.
```python
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.
## Understand Prediction Results
Understand Prediction Results
The returned `PredictionResult` is backed by local Parquet files:
| Attribute | Meaning |
| ---------- | ----------------------------------------------------------------- |
| `num_rows` | Number of output rows from result metadata. |
| `files` | Paths to the selected local Parquet files. |
| `frame` | A pandas DataFrame collected from all result files when accessed. |
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:
| Task | Main columns and interpretation |
| ------------------------- | ---------------------------------------------------------------------------------------------------- |
| Binary classification | One row per population row, with `ENTITY`, `TARGET_PRED`, `False_PROB`, and `True_PROB` |
| Multiclass classification | One row per returned class for each population row, with `ENTITY`, `CLASS`, `SCORE`, and `PREDICTED` |
| Multilabel classification | One row per returned label for each population row, with `ENTITY`, `CLASS`, `SCORE`, and `PREDICTED` |
| Multilabel ranking | One row per returned label for each population row, with `ENTITY`, `CLASS`, and `SCORE` |
| Regression | One row per population row, with `ENTITY` and `TARGET_PRED` |
| Link prediction | One row per returned candidate, with the source entity identifier, `CLASS`, and `SCORE` |
| Forecasting | One row per entity and forecast timestamp, with `ENTITY`, `TIMESTAMP`, and `TARGET_PRED` by default |
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](/kumo-engine/1.0.0/user-code-guide/code-examples/linux-deployment/forecasting#inspect-the-saved-forecast) selects this metadata and checks both timestamps in its saved output.
## Select Timestamp Metadata
`JOB_TIMESTAMP` records when scoring began in UTC without timezone information. `ANCHOR_TIMESTAMP` records each row's prediction anchor. Forecast `TIMESTAMP` identifies the start of each forecast window and is separate from both optional fields.
Choose scoring and anchor timestamps
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:
```python
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.
## Preserve Supplied Row Anchors
For row-specific times, [supply prediction rows](/kumo-engine/1.0.0/user-code-guide/train-a-model/predictive-queries#use-supplied-prediction-rows) and select `MetadataField.ANCHOR_TIMESTAMP` before scoring. Match outputs by entity and anchor together when the same entity occurs at several times. `JOB_TIMESTAMP` instead records the shared scoring start. Forecast supplied rows need the [required sequence columns](/kumo-engine/1.0.0/user-code-guide/train-a-model/predictive-queries#supply-forecast-sequences).
## When an Anchor Is Unavailable
Selecting `ANCHOR_TIMESTAMP` does not create an anchor for a population that has none.
For example, a generated static direct-column classification population can score without optional timestamps, but requesting its unavailable anchor raises `ValueError`: `Requested PREDICTIONS metadata field(s) were not produced: ANCHOR_TIMESTAMP`.
This error can occur when Kumo finalizes the output after scoring, rather than before prediction starts.
Remove `MetadataField.ANCHOR_TIMESTAMP` from `extra_fields` and rerun prediction when your task has no anchor.
You can still select `MetadataField.JOB_TIMESTAMP` to record scoring time.
If your temporal task needs row-specific anchors, supply its timestamps before scoring as shown above.
## Export an Existing Result
Copy a completed result without rescoring while its local Parquet files still exist.
Copy results and handle export retries
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:
```python
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](#append-to-a-warehouse-table). File, S3, and GCS destinations require `OVERWRITE` and follow the [file destination rules](#choose-file-formats-and-batch-paths).
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.
## Append to a Warehouse Table
BigQuery, Snowflake, and Databricks support `WriteMode.APPEND` for compatible tables. File and object-storage destinations do not. Keep output columns and types compatible, including metadata selected before scoring.
Append an existing result and check the saved rows
Reuse the output connector and table name from [Choose an Export Location](/kumo-engine/1.0.0/user-code-guide/making-predictions#choose-an-export-location). Append creates the table if it does not exist:
```python
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`.
> **Warning**
>
> Append is not deduplicated. A failed export may have written rows, so inspect the destination before retrying. Reusing a local result avoids rescoring but does not make writes exactly once. Configure recurring batch acceptance and retries in [Schedule Prediction Jobs](/kumo-engine/1.0.0/user-code-guide/deploy-a-model/schedule-prediction-jobs#validate-outputs-and-handle-retries).
## Generate and Save Embeddings
An embedding is a numeric vector representing an entity for downstream analysis. ID-GNN does not support embedding output. Generate embeddings separately using the same population and a new destination name.
Generate, save, and interpret entity embeddings
Set `embedding_name` to a separate table or path in the output connector, using the [destination naming rules](/kumo-engine/1.0.0/user-code-guide/making-predictions#choose-an-export-location). 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.
```python
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.
## Include Metadata in Embedding Outputs
Add `extra_fields` to the embedding call's `PredictionOutputConfig` to select scoring time or available row anchors, as in [Select Timestamp Metadata](#select-timestamp-metadata). Omitting the list omits both optional timestamps.
## Choose File Formats and Batch Paths
Choose File Formats and Batch Paths
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](#export-an-existing-result) to one of these destinations.
The returned result still references local Parquet files, even when the destination copy uses CSV.
| Destination name | Format and layout |
| --------------------- | ------------------------------------------------------------------------------------------------------------------------------------------ |
| `predictions.parquet` | Parquet; a single file when the result has one part, otherwise a directory of parts |
| `predictions.csv` | CSV; a single file when the result has one part, otherwise a directory containing `part-00000.csv`, `part-00001.csv`, and subsequent parts |
| `predictions.csv/` | CSV directory, including when the result has only one part |
| `predictions/` | Parquet directory, including when the result has only one part |
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](#export-an-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](/kumo-engine/1.0.0/get-started-with-kumo-engine/installation/databricks-notebook). This collects the full result in memory; use the warehouse output above for larger results. Choose a new directory:
```python
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](#export-an-existing-result).
## Save Predictions and Embeddings Together
Save Predictions and Embeddings Together
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](/kumo-engine/1.0.0/user-code-guide/making-predictions#generate-predictions) 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:
```python
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.
## Repeat a Prediction Run
Use [Production Workflow](/kumo-engine/1.0.0/user-code-guide/deploy-a-model/overview) to load a retained model in a fresh process and score updated data. Then [schedule prediction jobs](/kumo-engine/1.0.0/user-code-guide/deploy-a-model/schedule-prediction-jobs) with batch validation and retry policies. Retrain separately when you [refresh a model](/kumo-engine/1.0.0/user-code-guide/deploy-a-model/refresh-a-model).