Rai Predictive Training

Predictive Training

> Early access. The RAI predictive reasoner (GNN) is in early access — APIs, engine requirements, and behavior may change. Confirm the latest surface with the RelationalAI team before production use.

Summary

What: Training, evaluation, and model management workflow for GNN pipelines.

When to use:

• Configuring the GNN estimator and hyperparameters
• Training models with fit()
• Generating predictions on test data
• Evaluating and debugging results
• Registering or loading saved models

When NOT to use:

• Defining concepts, loading data, building graphs -- see rai-predictive-modeling

Overview: 4 steps: configure GNN -> train -> predict/evaluate -> optional: register/load.

By user intent — sections to focus on:

• Train + read validation metric → Quick Reference + GNN Constructor + gnn.fit()
• + predict + downstream rule / optimization → also Predictions + Using Predictions Downstream
• + register + reload across sessions → also Model Management

Quick Reference

Node Classification (minimal)

gnn = GNN(
    exp_database="DB", exp_schema="EXPERIMENTS",
    graph=gnn_graph, property_transformer=pt,
    train=Train, validation=Val,
    task_type="binary_classification", eval_metric="roc_auc",
    has_time_column=True, device="cuda", seed=42,
)
gnn.fit()
User.predictions = gnn.predictions(domain=Test)

Default Metrics

| Task Type | Suggested Metric | |-----------|-----------------| | binaryclassification | roc_auc | | multiclassclassification | accuracy | | multilabelclassification | multilabel_auprc_macro | | regression | rmse | | linkprediction | link_prediction_precision@5 | | repeatedlinkprediction | link_prediction_precision@5 |

Prediction Attributes

| Task Type | Attributes | |-----------|-----------| | binary_classification, multiclass_classification, multilabel_classification | .probs, .predicted_labels | | regression | .predicted_value | | link_prediction, repeated_link_prediction | .rank, .scores, .predicted_ |

---

GNN Constructor

Required Parameters

| Parameter | Description | |-----------|-------------| | exp_database, exp_schema | Snowflake location for experiment artifacts | | graph | Graph object with edges defined | | train, validation | Relationship objects | | task_type | Task type string | | eval_metric | Evaluation metric string |

Optional Parameters

| Parameter | Default | Description | |-----------|---------|-------------| | property_transformer | None | PropertyTransformer instance (omit for auto-inference) | | has_time_column | False | Set True if the task table contains a time column. See § Triple-coupling rule below. | | dataset_alias | None | Custom alias for the dataset |

Operational flags (stream_logs, parallel_reasoners_init, use_current_time, test_batch_size) are documented in [references/hyperparameters.md](references/hyperparameters.md) § GNN constructor operational flags.

Triple-coupling rule for has_time_column

When has_time_column=True, three things must move together:

1. Train/Val/Test Relationship(...) signatures use the at {Type:} clause.
2. PropertyTransformer(...) is constructed with time_col=[.] for at least one source concept — the same property must also appear in datetime=.
3. GNN(has_time_column=True, ...).

Setting only has_time_column=True without (2) raises ValueError: has_time_column=True is set but time_col is not defined in the PropertyTransformer from validate_time_col (relationalai.semantics.reasoners.predictive.preparation).

has_time_column is not a train/test split. It makes the relationship time-aware for the model; it does not partition train/val/test by time. For a forecast, also build the task tables with a temporal split (see rai-predictive-modeling § Define and Populate Concepts) — these are separate requirements.

Node Classification Example

gnn = GNN(
    exp_database="DB", exp_schema="EXPERIMENTS",
    graph=gnn_graph, property_transformer=pt,
    train=Train, validation=Val,
    task_type="binary_classification",
    eval_metric="roc_auc",
    has_time_column=True,
    device="cuda", seed=42,
    # Sweep n_epochs/lr from defaults (10 / 0.001) if learning is flat
)
gnn.fit()

Link Prediction Example (temporal)

gnn = GNN(
    exp_database="DB", exp_schema="EXPERIMENTS",
    graph=gnn_graph, property_transformer=pt,
    train=Train, validation=Val,
    task_type="repeated_link_prediction",
    eval_metric="link_prediction_precision@5",
    has_time_column=True,
    device="cuda", seed=42,
    head_layers=2,           # deeper head helps rank quality (default 1)
    label_smoothing=True,    # skill recommends True for link prediction (library default False)
)
gnn.fit()

Note: gnn.fit() trains at most once per GNN instance. If training has already completed (or is in progress), subsequent calls to fit() are silent no-ops. To retrain -- e.g. with different hyperparameters -- construct a new GNN instance.

Multi-GNN pipelines on the same model. Train multiple GNNs over the same entity set (e.g. regression + classification + link-prediction on the same graph) by reusing one Graph and one PropertyTransformer across all GNN instances; vary task_type, eval_metric, train/validation, and the source/target concepts. Bind each task's predictions to a distinct attribute name -- the convention Source.predictions collides if one source concept hosts more than one task.

shared = dict(graph=gnn_graph, property_transformer=pt)
gnn_a  = GNN(**shared, train=TrainA, validation=ValA, task_type="regression", eval_metric="rmse")
gnn_b  = GNN(**shared, train=TrainB, validation=ValB, task_type="binary_classification", eval_metric="roc_auc")
gnn_c  = GNN(**shared, train=TrainC, validation=ValC, task_type="repeated_link_prediction", eval_metric="link_prediction_precision@5")
for g in (gnn_a, gnn_b, gnn_c): g.fit()

# Distinct attributes when a source concept hosts multiple predictions:
Item.value_predictions = gnn_a.predictions(domain=TestA)
User.label_predictions = gnn_b.predictions(domain=TestB)
User.link_predictions  = gnn_c.predictions(domain=TestC)

Hyperparameters can also be passed as a dictionary:

train_config = {"device": "cuda", "n_epochs": 10, "lr": 0.001, "train_batch_size": 512, "seed": 42}
gnn = GNN(exp_database="DB", exp_schema="EXPERIMENTS", graph=gnn_graph, property_transformer=pt, **train_config)
# train=, validation=, task_type= are still required (as in the calls above) — the dict carries only tuning params
gnn.fit()

Unknown keys raise actionable errors. validate_train_params (relationalai.semantics.reasoners.predictive.preparation) rejects unknown **train_params keys with a ValueError and uses difflib to suggest near-matches; if you accidentally pass a GNN(...) constructor parameter (e.g. task_type=) inside train_params, the message tells you it belongs on the constructor. Read the suggestion before second-guessing the typo.

---

Common Hyperparameters

The full hyperparameter reference (every TrainerConfig field, library defaults, and typical-override examples) lives in [references/hyperparameters.md](references/hyperparameters.md). The most common knobs and their library defaults:

| Parameter | Default | Description | |-----------|---------|-------------| | device | "cuda" | "cuda" (GPU) or "cpu" | | n_epochs | 10 | Number of training epochs | | lr | 0.001 | Learning rate | | train_batch_size | 128 | Training batch size |

For link prediction, also consider: head_layers=2, num_negative=20, label_smoothing=True.

device="cuda" is a paired requirement. The client-side flag alone is not enough — the predictive reasoner engine must also be GPU-sized in raiconfig.yaml. Configure both or neither; mismatched settings silently fall back or fail. Heuristic: CPU HIGHMEM tiers trade training speed for more RAM; GPU is faster per epoch when the dataset fits in the GPU VM's CPU memory, HIGHMEM otherwise.

A GNN workflow touches multiple reasoner engines — size each per its role. At minimum, the predictive engine runs fit() and prediction jobs, and the logic engine runs model definitions, queries, and downstream rule evaluation; predict-then-optimize pipelines also need the prescriptive engine. Each is configured independently in raiconfig.yaml under reasoners: with its own name and size — appropriate sizing differs per role (GPU for predictive training, HIGHMEM CPU for logic query and rule workloads, and per-problem sizing for prescriptive). Mis-sizing one engine doesn't error loudly; the workflow still runs and silently under-performs or hits memory limits on that engine's step.

Auto-suspend during iteration. Set a low auto_suspend_mins on every engine you're using — idle pool cost can dominate total spend on small workloads. Warm pools make sense only for scheduled/production cadence. Specific tier names and per-cloud memory-vs-compute tradeoffs change over time — ask the RelationalAI team for current sizing. Full raiconfig.yaml structure (including the reasoners: block for all engine types) lives in the RAI configuration/setup skill.

Pre-flight check the GPU compute pool before long fit() runs. The predictive reasoner provisions onto a Snowpark Container Services GPU compute pool. Confirming the pool is active up front is cheap hygiene; if it's auto-suspended, jobs may queue without surfacing a clear progress signal. Substitute the pool name your account uses — SYSTEM_COMPUTE_POOL_GPU is the Snowflake-provided default; some accounts have a custom pool referenced from raiconfig.yaml:

SHOW COMPUTE POOLS;                                    -- list all + check state
ALTER COMPUTE POOL  RESUME;             -- if state=SUSPENDED

For all hyperparameters and tuning guidance, see [references/hyperparameters.md](references/hyperparameters.md).

---

Training

fit() Stages

gnn.fit() runs three stages internally:

1. Data preparation and feature extraction
2. Model training over n_epochs
3. Evaluation on the validation set

Timing expectations

gnn.fit() and gnn.predictions() both submit Snowpark Container Services jobs that can run for many minutes — the long quiet between submission and completion is expected, not stuck.

| Mode | Behavior | |------|----------| | stream_logs=True (default) | fit() blocks until training completes — log streaming runs synchronously inside fit() (relationalai.semantics.reasoners.predictive.estimator._stream_logs_formatted). The console silence after "Training job submitted" is the streamer waiting on log buffers, not a stalled client. | | stream_logs=False | fit() returns shortly after submission with "Job submitted and running in background." predictions() then waits — _wait_obtain_model_run_id blocks for training completion before submitting the prediction job. | | In both modes | predictions() always blocks until the prediction job completes (_wait_for_completion). |

Treat the run as "long-running" until it crosses ~5× the dataset-prep time printed at Step 1 before suspecting it's stuck. At that point run the diagnostic ladder below before suspending or killing anything.

"Training appears stuck"

Once the run crosses the ~5×-prep-time threshold above, run the three-step diagnostic ladder in [references/known-limitations.md](references/known-limitations.md) § "Training appears stuck" — diagnostic ladder before suspending or killing anything: (1) GET_REASONER('predictive', …) for pod status, (2) client.jobs.list("Predictive", …) for job state, (3) SHOW EXPERIMENTS for artifact creation. Each step localizes the failure before the next so you don't suspend the wrong reasoner.

Known Limitations & Runtime Troubleshooting

GNN training has runtime gotchas that surface as opaque or no-error symptoms in the client. Use this table to recognize each one; load references/known-limitations.md for full causes (with SDK source citations), the before/after fallback code for has_time_column=True at scale, and the GET_TRANSACTION_ARTIFACTS recipe.

| Symptom | Recover via | |---|---| | has_time_column=True fails with no time column defined in data tables | Temporal setup is incomplete. Verify (1) the relationship signature includes at {Type:}, (2) the time column is bound in define(Train(Source, train., ...)), (3) the column is a true DATE/TIMESTAMP* per rai-predictive-modeling § Auto-Discovery step 8. If all three are correct and the error persists, fall back to has_time_column=False + non-temporal Relationships. | | Train job stays QUEUED indefinitely while reasoner reports READY | rai-health § Predictive train jobs stuck QUEUED (SUSPEND_REASONER + RESUME_REASONER_ASYNC) | | gnn.fit() returns a model_run_id from an earlier job after a partial failure or notebook re-run | gnn.fit() is idempotent if self.train_job exists and isn't FAILED — re-instantiate GNN(...) on every retry, not bump Model("...") | | Client polls forever with no progress | JobMonitor._wait_for_completion has no timeout — kill the client manually + recover via the QUEUED runbook | | Failed to pull data into index: transaction was aborted (runtime error) | Opaque wrapper — pull RELATIONALAI.API.GET_TRANSACTION_ARTIFACTS('') -> problems.json for the real error. For the schema-drift / compiled-relation-cache cause: rename Model(...) |

For predictive train issues, stay on the supported RELATIONALAI.API.* surface — SUSPEND_REASONER / RESUME_REASONER_ASYNC for recovery, DELETE_REASONER + CREATE_REASONER_ASYNC('predictive', '', 'GPU_NV_S', OBJECT_CONSTRUCT()) for a fresh rebuild. See rai-health § Predictive train jobs stuck QUEUED.

---

Predictions

After training, generate predictions on the test set. Two valid binding patterns:

# Pattern 1 — bind to a concept attribute (queryable via select()):
Source.predictions = gnn.predictions(domain=Test)

# Pattern 2 — assign to a plain Python variable (re-callable):
predictions = gnn.predictions(domain=Test)

Each concept-attribute name can be assigned once per session — re-binding Source.predictions raises [Duplicate relationship]. To call predictions() multiple times in one session, use Pattern 2 or a fresh attribute name (e.g. Source.predictions_v2).

Silent corruption — Pattern 2 cartesian-joins against bare source columns. A plain-variable prediction relation is not bound to the source concept, so selecting the source's own columns alongside it cross-joins every source row with every prediction — N² rows with plausible-looking values and no error. Row count is the only tell: verify it equals the domain size.

# WRONG — Source.id is unbound relative to preds: N×N rows, no error.
preds = gnn.predictions(domain=Test)
select(Source.id, preds.probs).where(preds).inspect()

# CORRECT (option A) — bound attribute joins per-row:
Source.predictions = gnn.predictions(domain=Test)
select(Source.id, Source.predictions.probs).where(Source.predictions).inspect()

# CORRECT (option B) — reach the source through the relation's own field,
# indexed by the source concept's field name (see [evaluation-debugging.md](references/evaluation-debugging.md)):
select(preds["source"].id, preds.probs).where(preds).inspect()

With Pattern 2, never select the bare source concept's columns next to the relation — go through the relation's own fields, or use Pattern 1.

Per-task result access

Each task type exposes its results as attributes on the prediction relation (summary table in Quick Reference above) — classification: .probs / .predicted_labels; regression: .predicted_value; link prediction: .rank / .scores / .predicted_ (alwa

…

Source & license

This open-source skill is cataloged on AgentStack and links to its original source — we do not rehost the code.

• Author: RelationalAI
• Source: RelationalAI/rai-agent-skills
• License: Apache-2.0
• Homepage: https://relational.ai

Install and usage instructions live in the source repository linked above.