> 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()` when you want to provide the labeled examples that Kumo Relational uses as context. Unlike `predict()`, which derives context examples from a Predictive Query Language (PQL) query, `predict_task()` accepts the context and prediction entities as pandas DataFrames.

You still bind a relational graph to the client. The graph provides the tables, relationships, and metadata that Kumo Relational uses to construct relational context.

## Before you begin

You need:

* A configured and validated Kumo Relational graph.
* A `RelationalClient` connected to a Kumo Relational NIM.
* A context DataFrame containing labeled examples.
* 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` | Provides labeled examples for the task. |
| `predict` | `ENTITY`           | `ANCHOR_TIMESTAMP` | Identifies the entities to score.       |

`ENTITY` values reference entities from `entity_table`. `TARGET` contains the known label or value for each context example.

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

### Time-column behavior

Kumo Relational selects the task time 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

Set `task_type` 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](/rfm/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, including:

* `run_mode`
* `num_neighbors`
* `num_hops`
* `inference_config`
* `use_prediction_time`
* `return_embeddings`
* `random_seed`
* `explain`
* `batch_size` and `num_retries`
* `top_k`
* `exclude_cols_dict`
* `num_forecasts` and `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](/rfm/configuration) for inference and neighborhood settings, [Batch Prediction](/rfm/batch-prediction) for batching behavior, and [Prediction Explainability](/rfm/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](/rfm/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.