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

Prediction Attributes

GNN Constructor

Required Parameters

Optional Parameters

Operational flags (stream_logs, parallel_reasoners_init, use_current_time, test_batch_size) are documented in 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:<slot>} clause.
2. PropertyTransformer(...) is constructed with time_col=[<Concept>.<datetime_property>] 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. The most common knobs and their library defaults:

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 <your_gpu_pool> RESUME;             -- if state=SUSPENDED

For all hyperparameters and tuning guidance, see 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.

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

For predictive train issues, stay on the supported RELATIONALAI.API.* surface — SUSPEND_REASONER / RESUME_REASONER_ASYNC for recovery, DELETE_REASONER + CREATE_REASONER_ASYNC('predictive', '<name>', '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_<target> (always lowercase: Target Item -> .predicted_item). Swap .inspect() for .to_df() to get a pandas DataFrame. Full per-task code shapes including the link-prediction target join: references/prediction-attributes.md; dictionary-style field indexing (when the source concept name conflicts with an existing attribute) and gnn.prediction_concept direct access: references/evaluation-debugging.md.

Using Predictions Downstream

Once Source.predictions = gnn.predictions(...) runs, predictions are bound to the source concept and accessible via Source.predictions.<attribute> throughout the same Model. Other reasoners (rules, prescriptive, graph) consume them by deriving new properties from those attributes.

Same-model pattern (default)

Keep training, prediction, and downstream reasoning in one Model. This is the idiomatic RAI flow for predict-then-optimize and predict-then-rules chains:

# 1. Train and bind predictions
Item.predictions = gnn.predictions(domain=Test)

# 2. Derive a regular property from the prediction
Item.predicted_value = model.Property(f"{Item} has {Float:predicted_value}")
model.define(Item.predicted_value(Item.predictions.predicted_value))

# 3a. Predictive -> Rules: boolean flag
Item.is_high = model.Relationship(f"{Item} is high")
model.where(Item.predicted_value > threshold).define(Item.is_high())

# 3b. Predictive -> Prescriptive: Item.predicted_value can appear in
# Problem(model, Float) constraint / objective expressions.

Cross-session pattern (explicit persistence)

If training and downstream reasoning run in separate processes, persist predictions to Snowflake and reload them as a fresh Concept:

# Training session: save predictions DataFrame
df = select(Source.source_id, Source.predictions.predicted_value) \
     .where(Source.predictions).to_df()
# Then write_pandas(conn, df, "MY_PREDICTIONS", auto_create_table=True, overwrite=True)
# Grant SELECT on MY_PREDICTIONS to APPLICATION RELATIONALAI.

# Downstream session: load as a Concept in a new Model
Prediction = Concept("Prediction", identify_by={"source_id": Integer})
model.define(Prediction.new(Table("DB.SCHEMA.MY_PREDICTIONS").to_schema()))
# Derive properties from Prediction, apply rules, run a solver, etc.

database= and schema= on GNN(...) are optional and omitted throughout this skill. For durable persistence, use the explicit write_pandas path above.

Aggregation and bridge concepts

When the downstream reasoner's scope differs from the GNN source -- e.g. per-source predictions feeding a per-target optimizer -- aggregate predictions via aggregates.<agg>(...).per(Target).where(join) and attach the result to a bridge concept representing the downstream scope:

# GNN source predicts a value per Source (e.g. per-event regression);
# downstream scope is OptTarget, one row per coarser entity.
OptTarget = Concept("OptTarget", identify_by={"opt_target_id": Integer})
OptTarget.total_predicted_value = model.Property(f"{OptTarget} has {Float:total_predicted_value}")

agg = aggregates.sum(Source.predictions.predicted_value).per(OptTarget).where(
    Interaction.target_id == OptTarget.opt_target_id,
    Interaction.source_id == Source.source_id,
)
model.define(OptTarget.total_predicted_value(agg))

The bridge concept (OptTarget) separates what the GNN predicted at Source scope from what the downstream reasoner consumes at Target scope. Skipping the bridge and trying to use Source.predictions.predicted_value directly in a Target-scoped constraint forces ad-hoc joins inside each rule or objective expression. For classification or link-prediction predictions, swap sum/predicted_value for avg/probs or count/scores per the rule below.

Choose the aggregation function by target shape. Use sum for additive or count-like predictions (per-event regression values rolled up to an entity total), avg for proportional or probability-like predictions (mean predicted score across related source entities), count for link-prediction hits. Mixing them produces values that look numerically fine but don't mean what downstream expects.

Non-additive blending of multiple signals (e.g. combining several GNN outputs, or a GNN probability with a rule-derived flag) is also a derived-property step, not a built-in aggregates.<fn>. Express it as ordinary arithmetic in the property definition: a multiplicative composite (predicted_a * (1 - w * avg_b) * (1 + w * avg_c)) for risk-uplift-style logic, or a weighted interpolation (alpha * rule_signal + (1 - alpha) * gnn_probs) for hybrid scoring. Keep the bridge concept distinct from the GNN source so the blend is a regular Property the downstream reasoner can consume.

Denormalize if the target was pre-scaled at training. If the training target was normalized (e.g. to [0, 1], or z-scored), raw predictions carry that scale too. Record the denormalization factor alongside the derived property and apply it before feeding into constraints or objectives that expect real-world units -- otherwise the downstream reasoner sees tiny numbers where it expected the real-world quantity.

For a full predict-then-optimize example chaining multiple GNNs into optimizers with bridge + aggregation, see the retail_planning template in the templates repo.

Evaluation & Debugging

After gnn.fit(), inspect what data the engine received:

# Visual schema with data types — inline (Jupyter) or saved to file
gnn.display_dataset_diagram(show_dtypes=True)
gnn.save_dataset_diagram("dataset_schema.svg", show_dtypes=True)

# Full metadata dict (debugging feature types) and data-config printout
config = gnn.dataset.metadata_dict
gnn.dataset.print_data_config()

If results are poor, see references/evaluation-debugging.md § Tuning Poor Results for the ordered checklist (dataset inspection → text-feature reduction → hyperparameter tuning) plus regression-specific sanity checks, multi-metric framing, and leakage diagnostics.

Model Management

Register a Model

After gnn.fit() completes:

gnn.register_model(
    model_database="DB",
    model_schema="MODEL_REGISTRY",
    model_name="my_predictor",
    version_name="v1",
    comment="Initial training run",  # optional
)

Load by Registry Key

gnn = GNN(
    exp_database="DB", exp_schema="EXPERIMENTS",
    graph=gnn_graph, property_transformer=pt,
    source_concept=User,
    task_type="binary_classification",
    has_time_column=True,
    model_database="DB", model_schema="MODEL_REGISTRY",
    model_name="my_predictor", version_name="v1",
)
gnn.load()
User.predictions = gnn.predictions(domain=Test)

Load by Run ID

Same as above, replacing the registry key params with model_run_id="<run_id>".

What to Include vs. Omit When Loading

Train-Register-Load Workflow

Session 1: Train and Register

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", n_epochs=5, seed=42,
)
gnn.fit()
gnn.register_model(
    model_database="DB", model_schema="MODEL_REGISTRY",
    model_name="my_predictor", version_name="v1",
)

Session 2: Load and Predict

# Rebuild graph and property_transformer (same structure as training)
gnn_graph = Graph(model, directed=True, weighted=False)
# ... define edges ...
pt = PropertyTransformer(...)

gnn = GNN(
    exp_database="DB", exp_schema="EXPERIMENTS",
    graph=gnn_graph, property_transformer=pt,
    source_concept=User,
    task_type="binary_classification",
    has_time_column=True,
    model_database="DB", model_schema="MODEL_REGISTRY",
    model_name="my_predictor", version_name="v1",
)
gnn.load()
User.predictions = gnn.predictions(domain=Test)

Common Pitfalls

Examples

Reference Files