[Vedaant-J]

Series: kumo-engine MLOps series: every snapshot, training table, model and prediction run carries a self-contained _provenance.json; with kumo.mlops(...) stores each artifact under one artifact root, tracks fits and prediction runs in MLflow or Snowflake and registers models; volumes and stages can hold files; prediction runs are monitored for drift and realized performance.
Stack: Snapshot portability: #54 → #51
Merge order: one linear stack; merge bottom to top (#53 first)
This PR: Makes Trainer.predict(snapshot) the one scoring entry point, adds PredictiveQuery.from_snapshot(snapshot, query) and TrainingResult.predict, so a job that has a trained model bundle and a graph snapshot, but no Graph code, can score in one call.

Summary

This PR makes Trainer.predict(snapshot) the one scoring entry point, adds PredictiveQuery.from_snapshot(snapshot, query) and TrainingResult.predict, so a job that has a trained model bundle and a graph snapshot, but no Graph code, can score in one call.

Problem. Teams running Kumo in production need to answer questions like these:

• A scheduled scoring job has the published model bundle and the snapshot it should score against. How does it score without re-declaring every table, key and edge in Python?
• An MLflow model packages a bundle plus a snapshot. How does its predict() build a prediction table when nothing in the package can rebuild the Graph?
• Where does the job get the query text the model was trained with?
• How does the job make sure the snapshot it was handed is of the graph the model was trained on?

Today none of these work: PredictiveQuery can only be bound to a Graph, and a loaded Trainer keeps the training query privately.

Cause. A scoring job has these on hand, and no Graph:

<registry>/churn-model/                 # published model bundle
  training-spec.json                    # model plan, query text, snapshot_id
  metrics.json  prediction-schema.json  publication-manifest.json
  manifests/{graph,input,task}.json
  model/{checkpoint/, config.json, encoders.json, resolved-model-plan.json, target-encoder.json}

<snapshots>/f8fbba06…/07d6b2ca…/        # graph snapshot (local, or loaded from an export)
  _SUCCESS  _MANIFEST.json  diskgraph_config.json
  col_stats.json  edge_stats.json  ingest_stats.json  index_metadata.json
  kumo_col_stats/  nodes/{users,orders}/  train_tables/

GraphSnapshot already carries everything needed to compile a query: graph_config is the same portable graph config the source Graph produces, and Trainer.predict already scores against it. But PredictiveQuery read self.graph.native_graph to parse, and checked a snapshot with snapshot.id.startswith(f'{self.graph.id}@'). Passing a GraphSnapshot in place of the Graph, on main today:

>>> PredictiveQuery(snapshot, QUERY).get_task_type()
AttributeError: 'GraphSnapshot' object has no attribute 'native_graph'
>>> PredictiveQuery(snapshot, QUERY).generate_prediction_table(snapshot)
ValueError: snapshot 'f8fbba06…@07d6b2ca2d3ecf02' was not materialized from self.graph (id 'f8fbba06…@07d6b2ca2d3ecf02') -- call self.graph.snapshot() and pass its result.

Change. No new files or on-disk changes. PredictiveQuery.graph widens from Graph to Graph | GraphSnapshot, from_snapshot binds a snapshot, and Trainer.predict takes the snapshot first and uses the query of the last fit or load (read from the bundle's training-spec.json on load):

Before, the job must rebuild the graph:

from kumo.engine import FileConnector, Graph, PredictiveQuery, Trainer

QUERY = "PREDICT COUNT(orders.*, 0, 30, days) > 0 FOR EACH users.USER_ID"   # the training query
registry = FileConnector(root_dir="<registry>")
graph = Graph(tables={...}, edges=[...])          # every table, key, stype and edge again
snapshot = graph.snapshot()                       # re-reads sources to find the snapshot
trainer = Trainer.load_from_connector(registry, path="churn-model")
table = PredictiveQuery(graph, QUERY).generate_prediction_table(
    snapshot, plan=trainer.training_table_generation_plan)
scores = trainer.predict(table, snapshot=snapshot, binary_classification_threshold=0.5)

After, the snapshot is enough. It is loaded from the export the training job wrote (location is what graph.snapshot().export(SnapshotOutputConfig(store, "snapshots")) returned there):

from kumo.engine import FileConnector, GraphSnapshot, Trainer

store = FileConnector(root_dir="<snapshot store>")
registry = FileConnector(root_dir="<registry>")
location = "snapshots/<graph id>/<data version prefix>"
snapshot = GraphSnapshot.load_from_connector(store, path=location)
trainer = Trainer.load_from_connector(registry, path="churn-model")
scores = trainer.predict(snapshot, binary_classification_threshold=0.5)

predict(snapshot, *, table=None, entities=None, prediction_time=None, output_config=None, num_classes_to_return=None, num_workers=1, binary_classification_threshold=None) binds the saved query to the snapshot, builds the prediction table with the saved table-generation plan (trainer.training_table_generation_plan), and scores it. table= scores a prediction table generated beforehand instead. Without entities it scores the query's population at prediction_time (latest source time when omitted); entities is a source table of ENTITY1 (and TIMESTAMP) rows to score instead, and passing both raises ValueError. When the snapshot's graph key (the graph id half of snapshot.id) differs from the graph key of the snapshot the model was fitted or loaded with, it warns (UserWarning) and still scores: the model may be scoring a graph it was not trained on. It is the one scoring path the MLflow model (#65) and the job scripts (#71) call. main's form, predict(prediction_table, snapshot=...), keeps working, positionally or with prediction_table= as a keyword. A TrainingResult from fit or a loaded trainer's trainer.result also scores (result.predict(snapshot, ...)), with the model that produced it. The pieces stay public for jobs that need them:

query = PredictiveQuery.from_snapshot(snapshot, QUERY)
table = query.generate_prediction_table(snapshot, plan=trainer.training_table_generation_plan)
scores = trainer.predict(table, snapshot=snapshot, binary_classification_threshold=0.5)

On a synthetic users/orders graph (50 users, 3,000 orders, one-epoch CPU fit, bundle published to a FileConnector), a fresh process running the "after" code (with the query text passed explicitly; this run predates the merged predict(snapshot), which the tests below cover) returned 50 rows identical to the Graph-bound scores from the training process:

   ENTITY  TARGET_PRED  False_PROB  True_PROB
0       0         True    0.041282   0.958718
1      37         True    0.041199   0.958801

How it works:

• The query parses against the snapshot's graph_config in kumo-core's normalized form (a private cached GraphSnapshot._native_graph) instead of graph.native_graph. A Graph-bound query still uses the Graph's own.
• The snapshot check only establishes that the snapshot belongs to the query's graph. It is not a check that the snapshot is compatible with a trained model: that would compare what the model needs (features, types, relations, encoders, query contract, temporal settings) with the snapshot, and is not done here. The check compares the snapshot's graph key (the graph id part of the snapshot id, Export and load graph snapshots through any file destination #54) with the query's graph id, which is the same rule as before for a Graph-bound query. A snapshot-bound query also accepts its own snapshot, or any snapshot with the same graph_key, so a snapshot of the same graph over newer data is accepted; a snapshot of a different graph is refused with ValueError for both bindings.
• get_task_type, validate, suggest_model_plan, generate_training_table, generate_prediction_table and PredictionTable/TrainingTable.from_source_table all work on a snapshot-bound query.
• docs/mlops.md ("Scoring without the graph") says where a graph-free job gets its snapshot and model, scores with trainer.predict(snapshot), and states what the graph check does and does not mean.

Back-compat

This PR only adds an entry point. It adds no new checks or errors on any existing path.

Situation	Behaviour
Existing call PredictiveQuery(graph, query)	Unchanged: parses against graph.native_graph, same snapshot check. Only the refusal message is reworded ("was not materialized from this query's graph (id ...) -- pass a snapshot of that graph.").
New call PredictiveQuery.from_snapshot(snapshot, query)	Parses against snapshot.graph_config. Accepts any snapshot with the same graph key (the graph id half of snapshot.id).
Snapshot of a different graph	Refused with ValueError, for both bindings.
Snapshot whose id has no @	A snapshot-bound query accepts that same snapshot.
Trainer.predict(snapshot) (new form)	Before any fit or load: RuntimeError. Otherwise the same result as building the table by hand and calling predict. A snapshot of a different graph than the training snapshot gets a UserWarning and is still scored.
Existing call trainer.predict(prediction_table, snapshot=...) (positional or keyword)	Unchanged. Argument checks (num_workers, num_classes_to_return) still run first.
New TrainingResult.predict	Scores with the trainer that produced or loaded the result; a result built by hand raises RuntimeError.
Code that reads pquery.graph	Gets the GraphSnapshot for a snapshot-bound query; Graph otherwise.

Known limitations

• The graph check only establishes that a snapshot belongs to the model's graph (a warning in predict, a ValueError in PredictiveQuery). It does not check that the snapshot's features, types or encoders match what the model was trained on.
• predict(snapshot) needs the trained query and plan, so it works only on a trainer that has been fitted or loaded from a bundle; a hand-built TrainingResult cannot predict.
• mlops= on predict (storing and tracking prediction runs) comes with Record a manifest and output score statistics for every prediction run #67 and Log prediction runs to MLflow and Snowflake through the tracking adapters #73.

Testing Done

• test_query_from_snapshot.py (4 tests, run at the current head: 4 passed):
  • end to end: the fitted snapshot is exported through a FileConnector and loaded into a fresh scratch directory with no Graph; the model is reloaded from its bundle, and both result.predict(snapshot) and PredictiveQuery.from_snapshot + generate_prediction_table + predict reproduce the Graph-bound scores;
  • predict(snapshot) returns the same scores as the Graph-bound path with and without supplied entities, with a pre-generated table=, and in main's predict(prediction_table=..., snapshot=...) form; it refuses entities together with prediction_time;
  • a snapshot-bound query produces prediction tables (generated, and from supplied rows) that score identically;
  • a snapshot of another graph is refused, and a snapshot whose id has no @ accepts itself;
  • predict on a snapshot whose graph key differs from the training snapshot's warns and returns the same scores.
• Also at this head: test_prediction_workers.py (55), test_prediction_output_columns.py (25), test_prediction_completion.py (14), test_prediction_output_preflight.py (4), test_trainer_registry.py (27), test_warm_start.py (6), test_embedding_models.py (4), test_semantic_carry_forward.py (20), connector/test_file_connector.py (77). test_accelerator.py has 2 failures that also fail on Export and load graph snapshots through any file destination #54's head (environment).
• flake8, isort clean.

Checklist

• No breaking changes: PredictiveQuery(graph, query) and predict(prediction_table, snapshot=...) behave as before.
• Docs updated (docs/mlops.md).
• Unit tests added.

🤖 Generated with Claude Code

[Vedaant-J]

Change map

File	Function / scope	What changed	Why
packages/kumo-engine/docs/mlops.md	docs	"Scoring without the graph" section	user docs for graph-free scoring
packages/kumo-engine/kumo/engine/graph.py	GraphSnapshot._native_graph	private cached graph_config in kumo-core's normalized form	a snapshot-bound query parses against it
packages/kumo-engine/kumo/engine/pquery/predictive_query.py	PredictiveQuery.graph, from_snapshot, _native_graph, _parsed, _require_matching_snapshot	field accepts a GraphSnapshot; named constructor; parsing uses either binding's normalized graph; the snapshot check compares graph keys	bind a query to a snapshot without its Graph
packages/kumo-engine/kumo/engine/trainer/trainer.py	Trainer.predict, _prediction_table, _predict, TrainingResult.predict, TrainingResult._trainer	predict(snapshot, *, table=, entities=, prediction_time=, ...) builds the prediction table from the saved query and plan (warns on another graph); main's predict(prediction_table, snapshot=) still dispatches; results score through their trainer	one scoring entry point; graph-free jobs score in one call
packages/kumo-engine/test/test_query_from_snapshot.py	tests	export -> load into a fresh scratch -> result.predict equals Graph-bound scores; snapshot, entities, table and legacy forms; other-graph warning; snapshot check	behaviour coverage

Observed requirement: scheduled/MLflow scoring jobs hold a snapshot and a bundle but no Graph. Existing behaviour: PredictiveQuery needs a Graph and the trained query is private. Smallest change: accept a snapshot as the binding and let predict take the snapshot first. Verification: equality of predictions against the Graph-bound path. Stopping condition: every PredictiveQuery method works with either binding.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

• Last reviewed commit: 3e2701e
• Latest commit, not reviewed: 803cd82

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

Walkthrough

PredictiveQuery now accepts a Graph or GraphSnapshot. The new from_snapshot() method supports snapshot-bound construction. Query parsing and snapshot validation use the resolved graph configuration. GraphSnapshot.native_graph exposes the normalized native graph. Trainer.predict_snapshot generates prediction tables from saved training queries and delegates scoring to predict. Documentation and tests cover graphless snapshot workflows.

Sequence Diagram(s)

sequenceDiagram
  participant Caller
  participant Trainer
  participant PredictiveQuery
  participant GraphSnapshot
  Caller->>Trainer: predict_snapshot(snapshot, options)
  Trainer->>Trainer: read training_query
  Trainer->>PredictiveQuery: generate_prediction_table(snapshot)
  PredictiveQuery->>GraphSnapshot: validate snapshot compatibility
  PredictiveQuery-->>Trainer: return prediction table
  Trainer-->>Caller: return PredictionResult

Loading

Priority: Low

Merge Risk: Low · up to 3e270

Binary-classification users following the examples may encounter a scoring error. Clarify the required threshold and test scoring with restored artifacts; the remaining merge risk is bounded.

Pre-merge checks: 2 · 1

Failed checks: 1 inconclusive

Check name	Status	Explanation	Resolution
Change Map Design Review		The required PR conversation comment beginning with <!-- kumo-pr-change-map --> is not available. Repository inspection found only the template and review instructions, not the authored Change map. …	Provide the complete conversation comment starting with <!-- kumo-pr-change-map --> (and the full PR description if it contains the map link or required context), then reassess the five changed files against that map.

Passed checks: 2 passed

Check name	Status	Explanation
Pr Size		The pull request changes 5 files. This does not exceed the 40-file limit.
Public Resource Access		No added or modified external resource reference is present in the pull-request diff. The README additions use API names and illustrative local variables, with no URL or external repository/package lo…

Full details: Public Resource Access

Explanation

No added or modified external resource reference is present in the pull-request diff. The README additions use API names and illustrative local variables, with no URL or external repository/package location. The kumoai/backend/local/pquery.py reference already existed in the base and its reference text is unchanged.

Full details: Change Map Design Review

Explanation

The required PR conversation comment beginning with <!-- kumo-pr-change-map --> is not available. Repository inspection found only the template and review instructions, not the authored Change map. The PR description is also explicitly truncated. Without that comment, this check cannot assess the proposed design, rationale, verification, stopping condition, or omitted behavior against the implementation.

Resolution

Provide the complete conversation comment starting with <!-- kumo-pr-change-map --> (and the full PR description if it contains the map link or required context), then reassess the five changed files against that map.

✨ Finishing Touches: 💡 1

🛠️ Fix failing CI checks: 💡

• Commit to this branch
• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Note

Quiet mode is enabled, so only the most important comments were posted inline. Other review comments are grouped below.

🟡 Other comments: 1

packages/kumo-engine/README.md: 1

409-410: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win

Clarify snapshot reuse in the reload guidance.

The graph-free scoring path is documented earlier, but the reload section still instructs every scoring job to rebuild the live graph and snapshot. Distinguish graph-based scoring from graph-free scoring.

Suggested documentation fix

-For an exported model, use
-`Trainer.load_from_connector(connector, path="model/version")`, or
-`location=ModelRegistryLocation(...)` for Snowflake Model Registry. Rebuild
-the live graph and snapshot before generating prediction inputs. Bundles
-contain trained artifacts, not source data or graph storage.
+For an exported model, use
+`Trainer.load_from_connector(connector, path="model/version")`, or
+`location=ModelRegistryLocation(...)` for Snowflake Model Registry. If the
+scoring job has `Graph` code, rebuild the live graph and snapshot before
+generating prediction inputs. A graph-free scoring job can instead load the
+exported snapshot with `GraphSnapshot.load_from_connector(...)` and bind the
+query with `PredictiveQuery.from_snapshot(snapshot, trainer.training_query)`.
+Bundles contain trained artifacts, not source data or graph storage.

---

ℹ️ Review info

⚙️ Run configuration

• Configuration used: Repository: NVIDIA/kumo-engine/.coderabbit.yaml
• Review profile: QUIET
• Plan: Enterprise
• Run ID: 4372abd0-2cd3-4a74-a511-01be4ef9315b

📥 Commits

Reviewing files that changed from the base of the PR and between b4f46b8 and 2f97f33.

📒 Files selected for processing: 3

• packages/kumo-engine/README.md
• packages/kumo-engine/kumo/engine/graph.py
• packages/kumo-engine/kumo/engine/pquery/predictive_query.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 10 remain after this review.

[aaron-kumo]

Discussion — Confirm the compatibility boundary and scoring ergonomics.

PredictiveQuery.from_snapshot() and Trainer.training_query look consistent with the existing explicit, stage-oriented interface. I would keep these entry points.

Could we clarify that the graph-id check establishes membership between the query binding and supplied snapshot, rather than compatibility with the trained model? Future model compatibility should compare the trained requirements—features, types, relations, encoders, query contract, and temporal settings—with the scoring snapshot. Requiring equal training/scoring graph ids would unnecessarily restrict compatible graphs from another environment.

Also, callers must remember to pass trainer.training_table_generation_plan when generating prediction inputs. Would a future trainer helper that supplies the saved query/settings and performs compatibility checks be useful? I wouldn’t require that convenience API in this PR, but documenting the intended direction would help avoid multiple competing scoring paths.

[aaron-kumo]

Re-reviewed at 49bae78. Approved. PredictiveQuery.from_snapshot() and Trainer.training_query remain explicit entry points, and Trainer.predict_snapshot() supplies the saved query and table-generation plan for both generated and supplied entities. The docs correctly distinguish query/snapshot graph membership from model compatibility, with full model-compatibility checks explicitly deferred; a different training graph produces a warning rather than an unconditional refusal.

No additional issue identified in this PR's changes. Validation was source review and inspection of the added regression coverage; the full integration suites were not rerun because this checkout lacks the required Python environment.

[Vedaant-J]

Folded into #50 (artifacts & provenance), which now contains this PR's final code: the consolidated stack shows only the end state. Review continues there.