> 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.

# Prediction Guide

> Score a trained model and save predictions to files, object storage, or warehouse tables

Generate predictions with your existing model and save them to a retained destination. Continue with `trainer`, `query`, and `snapshot` from [Training](/kumo-engine/1.0.0/user-code-guide/train-a-model/training), or [retrieve a published model](/kumo-engine/1.0.0/user-code-guide/train-a-model/publish-and-retrieve-models#retrieve-a-published-model). Complete your data connection for [Linux](/kumo-engine/1.0.0/user-code-guide/connect-your-data-linux), [Snowflake](/kumo-engine/1.0.0/user-code-guide/connect-your-data-snowflake), or [Databricks](/kumo-engine/1.0.0/user-code-guide/connect-your-data-databricks) in the same Python session. The prediction calls below are shared across these environments; output setup depends on the destination.

## Select the Population

Generate the rows to score and carry the model's training-table generation settings forward:

```python
prediction_table = query.generate_prediction_table(
    snapshot,
    plan=trainer.training_table_generation_plan,
)
print("Selected rows:", prediction_table.row_count)
print("Prediction rows:", prediction_table.table_path)
```

For temporal queries, set a timezone-naive UTC `prediction_time` for repeatable batches, or omit it to use the latest relevant source timestamp. Reuse the trained generation plan, particularly for forecasts. Static queries do not accept a prediction time. See [population and anchor options](/kumo-engine/1.0.0/user-code-guide/advanced-prediction-configurations#population-and-anchor-options).

## Choose an Export Location

On Linux, save results in the same service as your source, under a dedicated prediction prefix or schema. Snowflake notebooks use a Snowflake output table. Databricks notebooks use a Unity Catalog table, including when reading source data from S3. Reuse your configured source identity where possible; source-read access does not grant output-write access.

Choose your destination: [S3](#amazon-s3), [GCS](#google-cloud-storage), [BigQuery](#bigquery-output), [Snowflake](#snowflake-output), [Databricks](#databricks-output), or [local files](#local-files). Each example defines `output_connector` and `output_name` for writing. Continue to [Generate Predictions](#generate-predictions) after configuring one destination.

For files, S3, and GCS, reuse the source connector when its configured root includes a suitable, separate output prefix: set `output_connector = connector` and `output_name = "kumo-results/batch-001/predictions.parquet"`. Otherwise create an output connector with the same identity as shown below. Choose a new batch path to retain previous results. Do not write into source table directories.

### Amazon S3

The existing AWS identity needs bucket listing and object read/write access to the output prefix, including permission to inspect the export markers. Reuse the installed S3 integration and configured credentials.

```python
from kumo.engine import S3Connector

output_connector = S3Connector(
    identity=connector.identity,
    root_dir="s3://my-bucket/kumo-results/batch-001/",
)
output_name = "predictions.parquet"
```

### Google Cloud Storage

The existing Google identity needs object listing, read, create, and overwrite/delete permissions on the destination. Reuse the installed GCS integration and credentials.

```python
from kumo.engine import GCSConnector

output_connector = GCSConnector(
    identity=connector.identity,
    root_dir="gs://my-bucket/kumo-results/batch-001/",
)
output_name = "predictions.parquet"
```

### BigQuery Output

Use an existing output dataset in the configured location. The identity needs the [BigQuery output permissions](/kumo-engine/1.0.0/user-code-guide/connect-your-data-linux#bigquery-staging-and-prediction-permissions), including table creation and data writes. Reuse the source identity and billing project.

```python
from kumo.engine import BigQueryConnector

output_connector = BigQueryConnector(
    identity=connector.identity,
    project_id=connector.project_id,
    billing_project_id=connector.billing_project_id,
    dataset_id="predictions",
    location=connector.location,
)
output_name = "customer_predictions_batch_001"
```

Use an unqualified table ID without a dataset prefix or file extension. `OVERWRITE` creates the table or replaces its rows and schema.

### Snowflake Output

First [prepare the Snowflake output schema and grants](/kumo-engine/1.0.0/user-code-guide/connect-your-data-snowflake#prepare-prediction-outputs). The same requirements apply when Python runs on Linux. Reuse the source connector's identity and query warehouse:

```python
from kumo.engine import SnowflakeConnector

output_connector = SnowflakeConnector(
    identity=connector.identity,
    warehouse=connector.warehouse,
    database=connector.database,
    role=connector.role,
    schema_name="PREDICTIONS",
)
output_name = "CUSTOMER_PREDICTIONS_BATCH_001"
```

Replace `PREDICTIONS` with your output schema. Use an unqualified table name. `OVERWRITE` retains grants on an existing compatible destination; its column names must match the generated output. If the source connector already targets your output schema, reuse it with `output_connector = connector`. Table retention follows Snowflake policies.

### Databricks Output

Prepare a dedicated output schema and staging volume using [Databricks output setup](/kumo-engine/1.0.0/user-code-guide/connect-your-data-databricks#prepare-prediction-outputs). Use the `output_connector` defined there. With a Unity Catalog source, it reuses the source identity and compute target. With an S3 source in a Databricks notebook, create the separate Databricks connector using notebook authentication and an explicit SQL warehouse, catalog, and schema. S3 access does not grant Unity Catalog access. The Unity Catalog connector settings also apply when Python runs on Linux.

```python
readback_name = "customer_predictions_batch_001"
output_name = f"{output_connector.schema_name}__kumo__{readback_name}"
```

The encoded name is for writes; readback uses the bare table name. Table retention follows Databricks policies.

### Local Files

For a file source, select a writable retained directory separately from the source table directories and execution scratch. Retention follows the lifetime of that filesystem.

```python
from kumo.engine import FileConnector

output_connector = FileConnector(root_dir="/data/kumo-results/batch-001")
output_name = "predictions.parquet"
```

## Generate Predictions

Score and save the selected population in one call. Binary classification requires a finite threshold from `0` to `1`; link prediction requires an explicit candidate count from `1` to `500`. The example sets these only for the applicable task. Other tasks use their query's defaults unless you [customize the output](/kumo-engine/1.0.0/user-code-guide/advanced-prediction-configurations#task-specific-output-options).

```python
from kumo.engine import PredictionOutputConfig
from kumoapi.jobs import PredictionArtifactType, WriteMode
from kumoapi.task import TaskType

prediction_options = {}
task_type = query.get_task_type()
if task_type == TaskType.BINARY_CLASSIFICATION:
    prediction_options["binary_classification_threshold"] = 0.5
elif task_type.is_link_pred:
    prediction_options["num_classes_to_return"] = 10

output_config = PredictionOutputConfig(
    connector=output_connector,
    artifact_type=PredictionArtifactType.PREDICTIONS,
    table_name=output_name,
    write_mode=WriteMode.OVERWRITE,
)
result = trainer.predict(
    prediction_table,
    snapshot=snapshot,
    output_config=output_config,
    **prediction_options,
)
print("Output rows:", result.num_rows)
```

Prediction runs synchronously. The snapshot must support the trained target, entity, relationships, and features, but it can contain newer data. The returned result still references local Parquet files; the saved destination has its own lifetime.

Continue to [Advanced Configurations](/kumo-engine/1.0.0/user-code-guide/advanced-prediction-configurations) for population filters, task output, timestamps, parallel scoring, exports, and embeddings.