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

# Predict with Custom Context

> Supply labeled context examples directly to a Kumo Relational prediction task

Use `predict_task()` to control which labeled examples Kumo Relational uses as context for a prediction.
Each example pairs an entity with a known outcome, giving the model examples of the task you want it to perform.
Unlike `predict()`, which derives context examples from a Predictive Query Language (PQL) query, `predict_task()` accepts both the context examples and the entities to score as pandas DataFrames.

You still bind a relational graph to the client.
The graph describes the tables, relationships, and metadata that Kumo Relational uses to gather information about each entity.

## Before you begin

You need:

* A configured and validated Kumo Relational graph
* A `RelationalClient` connected to a Kumo Relational NVIDIA NIM, the service that runs predictions
* A context DataFrame containing entities and their known outcomes
* A prediction DataFrame containing the entities to score

The table named by `entity_table` must exist in the graph.

## 1. Prepare the input tables

By default, `predict_task()` expects the following columns:

| DataFrame | Required columns   | Optional column    | Purpose                                               |
| --------- | ------------------ | ------------------ | ----------------------------------------------------- |
| `context` | `ENTITY`, `TARGET` | `ANCHOR_TIMESTAMP` | Supplies examples with known outcomes.                |
| `predict` | `ENTITY`           | `ANCHOR_TIMESTAMP` | Lists the entities for which to generate predictions. |

`ENTITY` values identify rows in `entity_table`.
In the context DataFrame, `TARGET` contains the known label or value for each example.

To use different column names, set `entity_column`, `target_column`, or `time_column` in the `predict_task()` call.

### Time-column behavior

For a time-dependent task, the time column anchors each example or prediction entity at a specific point in time.
Kumo Relational selects this column in the following order:

1. The column passed through `time_column`.
2. `ANCHOR_TIMESTAMP`, when that column appears in either input DataFrame.
3. The time column configured for the entity table.

When you explicitly set `time_column`, the context DataFrame must contain that column.

## 2. Select the task type

The `task_type` argument tells Kumo Relational how to interpret the `TARGET` values and format the predictions.
Set it to one of the task types supported by Kumo Relational:

* `binary_classification`
* `multiclass_classification`
* `regression`
* `forecasting`
* `temporal_link_prediction`

You can inspect the supported values through `client.capabilities("kumo-relational").tasks`.

For classification, regression, and forecasting tasks, set `entity_table` to the name of the entity table in the graph.
For temporal link prediction, pass the source and target entity-table names as a pair.

Forecasting requests require `step_size` and also accept `num_forecasts`.
`step_size` is an integer number of nanoseconds.
For example, calculate a 30-day step with pandas:

```python
import pandas as pd

step_size = int(pd.Timedelta(days=30).value)
```

See [Forecasting](/kumo-relational/user-guide/make-predictions/forecasting) for background on Kumo Relational forecasting tasks.

For temporal link prediction, `entity_table` must be a `(source_table, target_table)` pair and each context target must be a list of target IDs.

## 3. Run the prediction

Bind the graph to the client and call `predict_task()`:

```python
from kumo_relational_client import RelationalClient

with RelationalClient(url="http://localhost:8000") as client:
    result = client.relational(graph).predict_task(
        context=train_df,      # ENTITY, TARGET, [ANCHOR_TIMESTAMP]
        predict=predict_df,    # ENTITY, [ANCHOR_TIMESTAMP]
        task_type="multiclass_classification",
        entity_table="users",
    )
```

This example uses rows in `train_df` as labeled context and scores the entities in `predict_df`.

## Configure the prediction

`predict_task()` supports the same prediction controls used by the PQL workflow.
The controls fall into the following groups:

| Purpose                                 | Settings                                                                      |
| --------------------------------------- | ----------------------------------------------------------------------------- |
| Control context construction            | `run_mode`, `num_neighbors`, `num_hops`                                       |
| Configure inference inputs and behavior | `inference_config`, `use_prediction_time`, `random_seed`, `exclude_cols_dict` |
| Choose additional outputs               | `return_embeddings`, `explain`, `top_k`                                       |
| Split and retry requests                | `batch_size`, `num_retries`                                                   |
| Configure forecasting output            | `num_forecasts`, `step_size`                                                  |

`return_embeddings=True` is supported for binary classification and regression.
It appends an `EMBEDDINGS` column to each returned row without changing the task's row count.

See [Configuration](/kumo-relational/user-guide/configure-and-scale/configuration) for inference and neighborhood settings, [Batch Prediction](/kumo-relational/user-guide/configure-and-scale/batch-prediction) for batching behavior, and [Prediction Explainability](/kumo-relational/user-guide/explain-predictions/prediction-explainability) to request an explanation.

## Work with the result

By default, `predict_task()` returns a pandas DataFrame.
The result follows the same task-specific schema as a PQL-based prediction.
When `explain=True`, it returns an `Explanation` object whose `prediction` attribute contains the prediction DataFrame.

See [Prediction Results](/kumo-relational/user-guide/work-with-prediction-results/prediction-results) for the fields returned by each task.

## Troubleshoot input validation

The SDK rejects the request before prediction when:

* `task_type` is not supported.
* `entity_table` does not exist in the graph.
* The context DataFrame is missing the entity or target column.
* The prediction DataFrame is missing the entity column.
* An explicitly selected time column is missing from the context DataFrame.

Use the column names reported in the error to correct the DataFrame or the corresponding `entity_column`, `target_column`, or `time_column` argument.