From 68152b02ef716eee473a629a9a7fc1e8a5784bed Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 20 Oct 2026 20:05:48 +0000 Subject: [PATCH 01/10] Add PredictiveQuery.from_snapshot for graph-free scoring A job that loads a model bundle and a stored snapshot has no Python Graph, but PredictiveQuery required one, so it could not compile a prediction table. Prediction-table compilation only needs the snapshot's graph config. PredictiveQuery.from_snapshot(snapshot, query) binds the query to a GraphSnapshot: parsing uses the snapshot's graph_config and the snapshot match check uses the graph id embedded in snapshot.id, so a snapshot of a different graph is still refused. PredictiveQuery(graph, query) is unchanged. Verified: a snapshot-bound query generates prediction tables (generated and supplied rows) that Trainer.predict scores identically to the Graph-bound path; existing predictive-query and custom-source tests pass. Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/pquery/predictive_query.py | 83 ++++++++++++------- .../test/test_query_from_snapshot.py | 54 ++++++++++++ 2 files changed, 107 insertions(+), 30 deletions(-) create mode 100644 packages/kumo-engine/test/test_query_from_snapshot.py diff --git a/packages/kumo-engine/kumo/engine/pquery/predictive_query.py b/packages/kumo-engine/kumo/engine/pquery/predictive_query.py index 5cbf4771c..a967774f1 100644 --- a/packages/kumo-engine/kumo/engine/pquery/predictive_query.py +++ b/packages/kumo-engine/kumo/engine/pquery/predictive_query.py @@ -6,15 +6,19 @@ `PredictiveQuery(graph, query)` binds a `Graph` at construction, matching the cloud SDK's `kumo.PredictiveQuery(graph=graph, query=QUERY)` -- not a -`GraphSnapshot`. Parsing (`kumo_core.pquery.parse_and_validate_query` against -`self.graph.native_graph` -- `Graph`'s own `cached_property`, needing no -materialized index) happens lazily, on first access to the `_parsed` -`cached_property`, and only once: `validate()`/`get_task_type()` both read -that same property rather than re-parsing on every call. Confirmed against -the existing local backend's own `_LocalPQueryBackend.validate` -(`kumoai/backend/local/pquery.py`), which parses against -`to_portable(definition)` -- the schema-level conversion, not an ingested -snapshot -- that this doesn't need a `GraphSnapshot` either. +`GraphSnapshot`. `PredictiveQuery.from_snapshot(snapshot, query)` binds a +snapshot instead, for callers that load a snapshot without its `Graph`; the +query then parses against `snapshot.graph_config`, the same portable config +the source `Graph` produces. + +Parsing (`kumo_core.pquery.parse_and_validate_query` against +`self._native_graph`, needing no materialized index) happens lazily, on first +access to the `_parsed` `cached_property`, and only once: +`validate()`/`get_task_type()` both read that same property rather than +re-parsing on every call. Confirmed against the existing local backend's own +`_LocalPQueryBackend.validate` (`kumoai/backend/local/pquery.py`), which parses +against `to_portable(definition)` -- the schema-level conversion, not an +ingested snapshot -- that this doesn't need a `GraphSnapshot` either. Construction itself stays I/O-free, consistent with `Table`/`Graph`: nothing parses until `_parsed` is actually touched (by `validate()`, `get_task_type()`, @@ -25,8 +29,8 @@ `suggest_model_plan` uses resolved graph/query metadata. Methods that generate training or prediction tables also read the materialized index. These methods -accept an explicit `GraphSnapshot` produced by `self.graph.snapshot()` and -never trigger graph ingest themselves. +accept an explicit `GraphSnapshot` of this query's graph and never trigger +graph ingest themselves. Supported, matching `kumo.engine.pquery.predictive_query.PredictiveQuery` in the cloud SDK: `validate`, `get_task_type`, `suggest_model_plan`, @@ -57,7 +61,9 @@ ) from kumo_core.model_plan_conversion import public_model_plan_from_internal from kumo_core.pquery import ( + NativeGraph, ValidatedNativeQuery, + normalize_graph, parse_and_validate_query, query_is_temporal, resolve_query_split, @@ -97,15 +103,33 @@ @dataclass(frozen=True) class PredictiveQuery: - r"""An immutable PQL query bound to a `Graph`.""" + r"""An immutable PQL query bound to a `Graph` or a `GraphSnapshot`.""" - graph: Graph + graph: Graph | GraphSnapshot query: str + @classmethod + def from_snapshot(cls, snapshot: GraphSnapshot, + query: str) -> 'PredictiveQuery': + r"""Bind `query` to `snapshot`'s graph without its `Graph` object.""" + return cls(graph=snapshot, query=query) + + @cached_property + def _native_graph(self) -> NativeGraph: + if isinstance(self.graph, GraphSnapshot): + return normalize_graph(self.graph.graph_config) + return self.graph.native_graph + + @property + def _graph_id(self) -> str: + if isinstance(self.graph, GraphSnapshot): + return self.graph.id.partition('@')[0] + return self.graph.id + @cached_property def _parsed(self) -> ValidatedNativeQuery: - r"""Parses and validates `self.query` against `self.graph. - native_graph`, via `kumo_core.pquery.parse_and_validate_query` -- + r"""Parses and validates `self.query` against `self._native_graph`, + via `kumo_core.pquery.parse_and_validate_query` -- the same parser `Trainer._predict` already calls directly, and what `kumo_core.training.workflow.TrainingWorkflowRequest.from_spec` uses internally for `fit`. No materialized `GraphSnapshot` needed @@ -126,22 +150,21 @@ def _parsed(self) -> ValidatedNativeQuery: `ValidatedNativeQuery` are all plain `str`; `get_task_type` below wraps `task_type` into the shared, public `kumoapi.task.TaskType`. - Reuses `self.graph.native_graph` (itself a `cached_property`) - rather than building its own `graph_config`/`NativeGraph` -- one - canonical conversion per `Graph`, shared by `Graph.validate()` and - every `PredictiveQuery` bound to it. + For a `Graph`, reuses its cached `native_graph`, shared by + `Graph.validate()` and every `PredictiveQuery` bound to it. """ - return parse_and_validate_query(self.query, self.graph.native_graph) + return parse_and_validate_query(self.query, self._native_graph) def _require_matching_snapshot(self, snapshot: GraphSnapshot) -> None: # snapshot.id is 'graph_key@data_version_prefix' (see - # Graph.snapshot()); self.graph.id is the graph_key alone, so + # Graph.snapshot()); self._graph_id is the graph_key alone, so # membership is a prefix check, not equality. - if not snapshot.id.startswith(f'{self.graph.id}@'): + if snapshot.id != self.graph.id and not snapshot.id.startswith( + f'{self._graph_id}@'): raise ValueError( f"snapshot {snapshot.id!r} was not materialized from " - f"self.graph (id {self.graph.id!r}) -- call " - "self.graph.snapshot() and pass its result.") + f"this query's graph (id {self._graph_id!r}) -- pass a " + "snapshot of that graph.") def validate(self, verbose: bool = True) -> ValidationResponse: r"""Return validation warnings and information for a valid query. @@ -173,7 +196,7 @@ def get_task_type(self) -> TaskType: def _model_plan_validation_facts(self, run_mode: RunMode) -> dict[str, Any]: parsed = self._parsed - graph = self.graph.native_graph + graph = self._native_graph entity_tables = [parsed.entity_table] if parsed.rhs_node_type is not None: entity_tables.append(parsed.rhs_node_type) @@ -195,7 +218,7 @@ def suggest_model_plan( train_table_spec: Optional[TrainingTableSpec] = None, ) -> ModelPlan: r"""Suggests a modeling plan given this query and a materialized - `snapshot` of `self.graph`. Mirrors the cloud SDK's + `snapshot` of this query's graph. Mirrors the cloud SDK's `PredictiveQuery.suggest_model_plan`. Uses resolved graph/query metadata without generating a training table @@ -223,7 +246,7 @@ def generate_training_table( mlops: Optional[Any] = None, ) -> TrainingTable: r"""Compiles this query into a `TrainingTable` against a materialized - `snapshot` of `self.graph`. Mirrors the cloud SDK's + `snapshot` of this query's graph. Mirrors the cloud SDK's `PredictiveQuery.generate_training_table`. `Trainer.fit` consumes the returned `TrainingTable` directly. Its @@ -231,7 +254,7 @@ def generate_training_table( settings are retained without keeping a graph-bound task plan. Args: - snapshot: a snapshot of `self.graph` to compile against. + snapshot: a snapshot of this query's graph to compile against. plan: if not given, uses `kumo_core`'s generation defaults, matching the local backend's existing `suggest_training_table_plan` default. @@ -330,7 +353,7 @@ def generate_prediction_table( plan: Optional[TrainingTableGenerationPlan] = None, ) -> PredictionTable: r"""Compiles this query into a `PredictionTable` against a - materialized `snapshot` of `self.graph`. Mirrors the cloud SDK's + materialized `snapshot` of this query's graph. Mirrors the cloud SDK's `PredictiveQuery.generate_prediction_table`, and parallels `generate_training_table` above. @@ -342,7 +365,7 @@ def generate_prediction_table( docstring). Args: - snapshot: a snapshot of `self.graph` to compile against. + snapshot: a snapshot of this query's graph to compile against. prediction_time: the anchor timestamp for a temporal query; when omitted, uses the latest source-time maximum across the query's target aggregations in the snapshot. diff --git a/packages/kumo-engine/test/test_query_from_snapshot.py b/packages/kumo-engine/test/test_query_from_snapshot.py new file mode 100644 index 000000000..8c2d4ee23 --- /dev/null +++ b/packages/kumo-engine/test/test_query_from_snapshot.py @@ -0,0 +1,54 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from dataclasses import replace + +import pandas as pd +import pytest +from kumoapi.task import TaskType + +import kumo.engine as kumo + + +def _scores(trainer, table, snapshot): + frame = trainer.predict(table, snapshot=snapshot, + binary_classification_threshold=0.5).frame + return frame.sort_values('ENTITY').reset_index(drop=True) + + +def test_snapshot_bound_query_predicts_like_graph_bound(fitted, tmp_path): + trainer, graph_table, snapshot = fitted + query = kumo.PredictiveQuery.from_snapshot(snapshot, graph_table.query) + expected = _scores(trainer, graph_table, snapshot) + + assert query.get_task_type() == TaskType.BINARY_CLASSIFICATION + table = query.generate_prediction_table( + snapshot, plan=trainer.training_table_generation_plan) + assert table.timestamps == graph_table.timestamps + pd.testing.assert_frame_equal(_scores(trainer, table, snapshot), expected) + + anchor = pd.Timestamp(graph_table.timestamps[0]) + pd.DataFrame({ + 'ENTITY1': range(50), + 'TIMESTAMP': [anchor] * 50, + }).to_parquet(tmp_path / 'rows.parquet') + supplied = kumo.PredictionTable.from_source_table( + kumo.FileConnector(root_dir=str(tmp_path))['rows.parquet'], pq=query, + snapshot=snapshot) + pd.testing.assert_frame_equal(_scores(trainer, supplied, snapshot), + expected) + + +def test_snapshot_bound_query_refuses_other_graph(fitted): + _, graph_table, snapshot = fitted + query = kumo.PredictiveQuery.from_snapshot(snapshot, graph_table.query) + newer = replace(snapshot, id=snapshot.id.split('@')[0] + '@newer') + other = replace(snapshot, id='other@' + snapshot.id.split('@')[1]) + + query._require_matching_snapshot(newer) + plain = replace(snapshot, id='plain') + kumo.PredictiveQuery.from_snapshot(plain, '')._require_matching_snapshot( + plain) + with pytest.raises(ValueError, match='was not materialized'): + query.generate_prediction_table(other) From 5256d2c881580c3e8e82fe0f7ebca193a0767af4 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 08:08:23 +0000 Subject: [PATCH 02/10] Expose Trainer.training_query and keep the graph-bound snapshot check A job scoring a loaded model had no public way to get the query it was trained on, so binding it with PredictiveQuery.from_snapshot meant copying the text out of training-spec.json. Trainer.training_query now returns it. The snapshot check accepted a snapshot whose id equals a Graph's bare id; equality now applies only to a snapshot-bound query, so a Graph-bound query keeps the original prefix check. Verified with test_query_from_snapshot.py. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/pquery/predictive_query.py | 6 ++++-- packages/kumo-engine/kumo/engine/trainer/trainer.py | 6 ++++++ packages/kumo-engine/test/test_query_from_snapshot.py | 3 ++- 3 files changed, 12 insertions(+), 3 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/pquery/predictive_query.py b/packages/kumo-engine/kumo/engine/pquery/predictive_query.py index a967774f1..f7cd0f5c8 100644 --- a/packages/kumo-engine/kumo/engine/pquery/predictive_query.py +++ b/packages/kumo-engine/kumo/engine/pquery/predictive_query.py @@ -159,8 +159,10 @@ def _require_matching_snapshot(self, snapshot: GraphSnapshot) -> None: # snapshot.id is 'graph_key@data_version_prefix' (see # Graph.snapshot()); self._graph_id is the graph_key alone, so # membership is a prefix check, not equality. - if snapshot.id != self.graph.id and not snapshot.id.startswith( - f'{self._graph_id}@'): + if isinstance(self.graph, GraphSnapshot) and (snapshot.id + == self.graph.id): + return + if not snapshot.id.startswith(f'{self._graph_id}@'): raise ValueError( f"snapshot {snapshot.id!r} was not materialized from " f"this query's graph (id {self._graph_id!r}) -- pass a " diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 1acc5e100..c3d6d04d6 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -344,6 +344,12 @@ def training_table_generation_plan( """Effective table-generation settings saved by the last fit/load.""" return self._training_table_generation_plan + @property + def training_query(self) -> Optional[str]: + """Predictive query of the last fit/load, for + :meth:`PredictiveQuery.from_snapshot`.""" + return self._training_query + @property def scratch_dir(self) -> Path: return trainer_scratch_dir(self._id) diff --git a/packages/kumo-engine/test/test_query_from_snapshot.py b/packages/kumo-engine/test/test_query_from_snapshot.py index 8c2d4ee23..1e90285ea 100644 --- a/packages/kumo-engine/test/test_query_from_snapshot.py +++ b/packages/kumo-engine/test/test_query_from_snapshot.py @@ -19,7 +19,8 @@ def _scores(trainer, table, snapshot): def test_snapshot_bound_query_predicts_like_graph_bound(fitted, tmp_path): trainer, graph_table, snapshot = fitted - query = kumo.PredictiveQuery.from_snapshot(snapshot, graph_table.query) + assert trainer.training_query == graph_table.query + query = kumo.PredictiveQuery.from_snapshot(snapshot, trainer.training_query) expected = _scores(trainer, graph_table, snapshot) assert query.get_task_type() == TaskType.BINARY_CLASSIFICATION From 81fb21c075407a967448120a35407845b1b5939c Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 08:19:53 +0000 Subject: [PATCH 03/10] Trim comments and implementation-detail tests The from_snapshot change carried docstring prose that restated the code, and its graph-check test called a private helper. Shorten the module note and training_query docstring, drop the restating _parsed paragraph, and exercise the snapshot check through the public suggest_model_plan instead of _require_matching_snapshot. test_query_from_snapshot.py: 2 passed. Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/pquery/predictive_query.py | 7 +------ packages/kumo-engine/kumo/engine/trainer/trainer.py | 3 +-- .../kumo-engine/test/test_query_from_snapshot.py | 12 ++++++------ 3 files changed, 8 insertions(+), 14 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/pquery/predictive_query.py b/packages/kumo-engine/kumo/engine/pquery/predictive_query.py index f7cd0f5c8..300fa1410 100644 --- a/packages/kumo-engine/kumo/engine/pquery/predictive_query.py +++ b/packages/kumo-engine/kumo/engine/pquery/predictive_query.py @@ -7,9 +7,7 @@ `PredictiveQuery(graph, query)` binds a `Graph` at construction, matching the cloud SDK's `kumo.PredictiveQuery(graph=graph, query=QUERY)` -- not a `GraphSnapshot`. `PredictiveQuery.from_snapshot(snapshot, query)` binds a -snapshot instead, for callers that load a snapshot without its `Graph`; the -query then parses against `snapshot.graph_config`, the same portable config -the source `Graph` produces. +snapshot instead and parses against `snapshot.graph_config`. Parsing (`kumo_core.pquery.parse_and_validate_query` against `self._native_graph`, needing no materialized index) happens lazily, on first @@ -149,9 +147,6 @@ def _parsed(self) -> ValidatedNativeQuery: `task_type`/`task_family`/`query_type` on the returned `ValidatedNativeQuery` are all plain `str`; `get_task_type` below wraps `task_type` into the shared, public `kumoapi.task.TaskType`. - - For a `Graph`, reuses its cached `native_graph`, shared by - `Graph.validate()` and every `PredictiveQuery` bound to it. """ return parse_and_validate_query(self.query, self._native_graph) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index c3d6d04d6..4396b8731 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -346,8 +346,7 @@ def training_table_generation_plan( @property def training_query(self) -> Optional[str]: - """Predictive query of the last fit/load, for - :meth:`PredictiveQuery.from_snapshot`.""" + """Predictive query of the last fit/load.""" return self._training_query @property diff --git a/packages/kumo-engine/test/test_query_from_snapshot.py b/packages/kumo-engine/test/test_query_from_snapshot.py index 1e90285ea..b106248aa 100644 --- a/packages/kumo-engine/test/test_query_from_snapshot.py +++ b/packages/kumo-engine/test/test_query_from_snapshot.py @@ -41,15 +41,15 @@ def test_snapshot_bound_query_predicts_like_graph_bound(fitted, tmp_path): expected) -def test_snapshot_bound_query_refuses_other_graph(fitted): +def test_snapshot_bound_query_checks_snapshot_graph(fitted): _, graph_table, snapshot = fitted query = kumo.PredictiveQuery.from_snapshot(snapshot, graph_table.query) newer = replace(snapshot, id=snapshot.id.split('@')[0] + '@newer') other = replace(snapshot, id='other@' + snapshot.id.split('@')[1]) - - query._require_matching_snapshot(newer) plain = replace(snapshot, id='plain') - kumo.PredictiveQuery.from_snapshot(plain, '')._require_matching_snapshot( - plain) + + query.suggest_model_plan(newer) + kumo.PredictiveQuery.from_snapshot(plain, graph_table.query) \ + .suggest_model_plan(plain) with pytest.raises(ValueError, match='was not materialized'): - query.generate_prediction_table(other) + query.suggest_model_plan(other) From 4ffea50945c1348fdc463294aee854f0d8cc5a8c Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 18:47:20 +0000 Subject: [PATCH 04/10] Read the native graph and graph id from GraphSnapshot PredictiveQuery kept its own normalized-graph cache and parsed the graph id out of snapshot ids. GraphSnapshot now offers native_graph (cached, like Graph.native_graph), and the snapshot check compares graph_key. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/graph.py | 6 +++ .../kumo/engine/pquery/predictive_query.py | 38 ++++++------------- 2 files changed, 18 insertions(+), 26 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index 3c182605a..4174ca94d 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -326,6 +326,12 @@ def _data_version(self) -> str: from; `id` carries its first 16 characters.""" return Path(self.index_path).name + @cached_property + def native_graph(self) -> NativeGraph: + """`graph_config` in `kumo_core`'s normalized form, as + `Graph.native_graph`.""" + return normalize_graph(self.graph_config) + def export(self, output_config: 'SnapshotOutputConfig') -> str: r"""Copy the index to `//` and return that location. Any earlier export there is replaced; it is not diff --git a/packages/kumo-engine/kumo/engine/pquery/predictive_query.py b/packages/kumo-engine/kumo/engine/pquery/predictive_query.py index 300fa1410..0eee72f01 100644 --- a/packages/kumo-engine/kumo/engine/pquery/predictive_query.py +++ b/packages/kumo-engine/kumo/engine/pquery/predictive_query.py @@ -10,7 +10,7 @@ snapshot instead and parses against `snapshot.graph_config`. Parsing (`kumo_core.pquery.parse_and_validate_query` against -`self._native_graph`, needing no materialized index) happens lazily, on first +`self.graph.native_graph`, needing no materialized index) happens lazily, on first access to the `_parsed` `cached_property`, and only once: `validate()`/`get_task_type()` both read that same property rather than re-parsing on every call. Confirmed against the existing local backend's own @@ -59,9 +59,7 @@ ) from kumo_core.model_plan_conversion import public_model_plan_from_internal from kumo_core.pquery import ( - NativeGraph, ValidatedNativeQuery, - normalize_graph, parse_and_validate_query, query_is_temporal, resolve_query_split, @@ -112,21 +110,9 @@ def from_snapshot(cls, snapshot: GraphSnapshot, r"""Bind `query` to `snapshot`'s graph without its `Graph` object.""" return cls(graph=snapshot, query=query) - @cached_property - def _native_graph(self) -> NativeGraph: - if isinstance(self.graph, GraphSnapshot): - return normalize_graph(self.graph.graph_config) - return self.graph.native_graph - - @property - def _graph_id(self) -> str: - if isinstance(self.graph, GraphSnapshot): - return self.graph.id.partition('@')[0] - return self.graph.id - @cached_property def _parsed(self) -> ValidatedNativeQuery: - r"""Parses and validates `self.query` against `self._native_graph`, + r"""Parses and validates `self.query` against `self.graph.native_graph`, via `kumo_core.pquery.parse_and_validate_query` -- the same parser `Trainer._predict` already calls directly, and what `kumo_core.training.workflow.TrainingWorkflowRequest.from_spec` @@ -148,19 +134,19 @@ def _parsed(self) -> ValidatedNativeQuery: `ValidatedNativeQuery` are all plain `str`; `get_task_type` below wraps `task_type` into the shared, public `kumoapi.task.TaskType`. """ - return parse_and_validate_query(self.query, self._native_graph) + return parse_and_validate_query(self.query, self.graph.native_graph) def _require_matching_snapshot(self, snapshot: GraphSnapshot) -> None: - # snapshot.id is 'graph_key@data_version_prefix' (see - # Graph.snapshot()); self._graph_id is the graph_key alone, so - # membership is a prefix check, not equality. - if isinstance(self.graph, GraphSnapshot) and (snapshot.id - == self.graph.id): - return - if not snapshot.id.startswith(f'{self._graph_id}@'): + if isinstance(self.graph, GraphSnapshot): + if snapshot.id == self.graph.id: + return + graph_key = self.graph.graph_key + else: + graph_key = self.graph.id + if snapshot.graph_key != graph_key: raise ValueError( f"snapshot {snapshot.id!r} was not materialized from " - f"this query's graph (id {self._graph_id!r}) -- pass a " + f"this query's graph (id {graph_key!r}) -- pass a " "snapshot of that graph.") def validate(self, verbose: bool = True) -> ValidationResponse: @@ -193,7 +179,7 @@ def get_task_type(self) -> TaskType: def _model_plan_validation_facts(self, run_mode: RunMode) -> dict[str, Any]: parsed = self._parsed - graph = self._native_graph + graph = self.graph.native_graph entity_tables = [parsed.entity_table] if parsed.rhs_node_type is not None: entity_tables.append(parsed.rhs_node_type) From 76da3dc9010bb98a86efc05c5020fa1427883b4f Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 21:18:29 +0000 Subject: [PATCH 05/10] Add Trainer.predict_snapshot for one-call graph-free scoring A graph-free scoring job had to bind the saved query to the snapshot, pass the saved generation plan to generate_prediction_table and then call predict; each caller (MLflow model, job scripts) repeated that. Trainer.predict_snapshot(snapshot, *, entities=None, prediction_time=None, ...) does it in one call with predict's options. The docs now say the query's graph check only ties a snapshot to the query's graph, not to a trained model. Verified with the query-from-snapshot and predictive-query tests. Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/pquery/predictive_query.py | 8 +++- .../kumo/engine/trainer/trainer.py | 42 +++++++++++++++++++ .../test/test_query_from_snapshot.py | 20 +++++++++ 3 files changed, 69 insertions(+), 1 deletion(-) diff --git a/packages/kumo-engine/kumo/engine/pquery/predictive_query.py b/packages/kumo-engine/kumo/engine/pquery/predictive_query.py index 0eee72f01..d36621e05 100644 --- a/packages/kumo-engine/kumo/engine/pquery/predictive_query.py +++ b/packages/kumo-engine/kumo/engine/pquery/predictive_query.py @@ -107,7 +107,13 @@ class PredictiveQuery: @classmethod def from_snapshot(cls, snapshot: GraphSnapshot, query: str) -> 'PredictiveQuery': - r"""Bind `query` to `snapshot`'s graph without its `Graph` object.""" + r"""Bind `query` to `snapshot`'s graph without its `Graph` object. + + Snapshots passed to this query's methods must come from the same + graph (same `graph_key`). That only checks the snapshot belongs to + this query's graph; it says nothing about compatibility with a + trained model. + """ return cls(graph=snapshot, query=query) @cached_property diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 4396b8731..16a508a85 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -55,6 +55,7 @@ from kumo.engine._manifest import ModelManifest from kumo.engine._scratch import new_run_id, trainer_scratch_dir from kumo.engine.connector.base import Connector +from kumo.engine.connector.source_table import SourceTable from kumo.engine.graph import GraphSnapshot from kumo.engine.pquery.prediction_table import PredictionTable from kumo.engine.pquery.training_table import TrainingTable @@ -685,6 +686,47 @@ def _finish_fit( return training_result + def predict_snapshot( + self, + snapshot: GraphSnapshot, + *, + entities: Optional[SourceTable] = None, + prediction_time: Optional[datetime] = None, + output_config: Optional[PredictionOutputConfig] = None, + num_classes_to_return: Optional[int] = None, + num_workers: int = 1, + binary_classification_threshold: Optional[float] = None, + ) -> PredictionResult: + r"""Score ``snapshot`` with the query and table-generation plan saved + by the last fit/load. + + Without ``entities``, scores the query's population at + ``prediction_time`` (as :meth:`PredictiveQuery.generate_prediction_table`). + ``entities`` is a source table of ``ENTITY1`` (and ``TIMESTAMP`` where + the task needs one) rows to score instead. Other arguments are as in + :meth:`predict`. + """ + from kumo.engine.pquery.predictive_query import PredictiveQuery + + if self._training_query is None: + raise RuntimeError( + "Trainer.predict_snapshot() requires a prior Trainer.fit() or " + "Trainer.load() call on this same Trainer instance.") + if entities is not None and prediction_time is not None: + raise ValueError("Pass prediction_time or entities, not both; " + "put anchor times in the entities' TIMESTAMP.") + query = PredictiveQuery.from_snapshot(snapshot, self._training_query) + plan = self._training_table_generation_plan + table = (query.generate_prediction_table( + snapshot, prediction_time=prediction_time, plan=plan) + if entities is None else PredictionTable.from_source_table( + entities, pq=query, snapshot=snapshot, plan=plan)) + return self.predict( + table, snapshot=snapshot, output_config=output_config, + num_classes_to_return=num_classes_to_return, + num_workers=num_workers, + binary_classification_threshold=binary_classification_threshold) + @_without_torch_scatter_hint def predict( self, diff --git a/packages/kumo-engine/test/test_query_from_snapshot.py b/packages/kumo-engine/test/test_query_from_snapshot.py index b106248aa..65629f9b4 100644 --- a/packages/kumo-engine/test/test_query_from_snapshot.py +++ b/packages/kumo-engine/test/test_query_from_snapshot.py @@ -41,6 +41,26 @@ def test_snapshot_bound_query_predicts_like_graph_bound(fitted, tmp_path): expected) +def test_predict_snapshot_scores_with_the_saved_query(fitted, tmp_path): + trainer, graph_table, snapshot = fitted + expected = _scores(trainer, graph_table, snapshot) + + def scores(**options): + return trainer.predict_snapshot( + snapshot, binary_classification_threshold=0.5, + **options).frame.sort_values('ENTITY').reset_index(drop=True) + + pd.testing.assert_frame_equal(scores(), expected) + pd.DataFrame({ + 'ENTITY1': range(50), + 'TIMESTAMP': [pd.Timestamp(graph_table.timestamps[0])] * 50, + }).to_parquet(tmp_path / 'rows.parquet') + rows = kumo.FileConnector(root_dir=str(tmp_path))['rows.parquet'] + pd.testing.assert_frame_equal(scores(entities=rows), expected) + with pytest.raises(ValueError, match='not both'): + scores(entities=rows, prediction_time=graph_table.timestamps[0]) + + def test_snapshot_bound_query_checks_snapshot_graph(fitted): _, graph_table, snapshot = fitted query = kumo.PredictiveQuery.from_snapshot(snapshot, graph_table.query) From aab9d02ab40e25533b552cfd4caf2a51561ec8fe Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 22:36:39 +0000 Subject: [PATCH 06/10] Test graph-free scoring of an exported snapshot end to end Exports the fitted snapshot, loads it into a fresh scratch directory without a Graph, reloads the model from its bundle and scores with predict_snapshot and with PredictiveQuery.from_snapshot; both match the graph-bound scores. Co-Authored-By: Claude Opus 5.5 --- .../test/test_query_from_snapshot.py | 27 +++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/packages/kumo-engine/test/test_query_from_snapshot.py b/packages/kumo-engine/test/test_query_from_snapshot.py index 65629f9b4..87866165c 100644 --- a/packages/kumo-engine/test/test_query_from_snapshot.py +++ b/packages/kumo-engine/test/test_query_from_snapshot.py @@ -9,6 +9,7 @@ from kumoapi.task import TaskType import kumo.engine as kumo +from kumo.engine._scratch import get_scratch_dir def _scores(trainer, table, snapshot): @@ -61,6 +62,32 @@ def scores(**options): scores(entities=rows, prediction_time=graph_table.timestamps[0]) +def test_exported_snapshot_scores_in_a_fresh_scratch_without_a_graph( + fitted, tmp_path): + trainer, graph_table, snapshot = fitted + expected = _scores(trainer, graph_table, snapshot) + exports = kumo.FileConnector(root_dir=str(tmp_path / 'exports')) + location = snapshot.export(exports, 'snapshots') + bundle = trainer._training_result.bundle_dir + previous = get_scratch_dir() + kumo.configure(scratch_dir=tmp_path / 'fresh') + try: + loaded = kumo.GraphSnapshot.load_from_connector(exports, path=location) + assert loaded.id == snapshot.id + assert loaded.index_path.is_relative_to(tmp_path / 'fresh') + model = kumo.Trainer.load(bundle) + scored = model.predict_snapshot( + loaded, binary_classification_threshold=0.5).frame + pd.testing.assert_frame_equal( + scored.sort_values('ENTITY').reset_index(drop=True), expected) + query = kumo.PredictiveQuery.from_snapshot(loaded, model.training_query) + table = query.generate_prediction_table( + loaded, plan=model.training_table_generation_plan) + pd.testing.assert_frame_equal(_scores(model, table, loaded), expected) + finally: + kumo.configure(scratch_dir=previous) + + def test_snapshot_bound_query_checks_snapshot_graph(fitted): _, graph_table, snapshot = fitted query = kumo.PredictiveQuery.from_snapshot(snapshot, graph_table.query) From 32e9cdda3cb6b6c24ded1497c2d3312c3f7f613b Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 22 Oct 2026 01:14:39 +0000 Subject: [PATCH 07/10] Warn when predict_snapshot scores another graph's snapshot predict_snapshot bound the saved query to whatever snapshot it was given, so scoring a snapshot of a different graph than the model was trained on went through silently. It now warns (and still scores) when the snapshot's graph_key differs from the graph key of the snapshot recorded at fit/load. Verified with test_query_from_snapshot.py. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/trainer/trainer.py | 9 ++++++++- packages/kumo-engine/test/test_query_from_snapshot.py | 6 ++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 16a508a85..6a5ba60b3 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -704,7 +704,8 @@ def predict_snapshot( ``prediction_time`` (as :meth:`PredictiveQuery.generate_prediction_table`). ``entities`` is a source table of ``ENTITY1`` (and ``TIMESTAMP`` where the task needs one) rows to score instead. Other arguments are as in - :meth:`predict`. + :meth:`predict`. Warns when ``snapshot`` was taken of a different + graph than the one the model was trained on. """ from kumo.engine.pquery.predictive_query import PredictiveQuery @@ -715,6 +716,12 @@ def predict_snapshot( if entities is not None and prediction_time is not None: raise ValueError("Pass prediction_time or entities, not both; " "put anchor times in the entities' TIMESTAMP.") + trained = (self._snapshot_id or '').partition('@')[0] + if trained and snapshot.graph_key != trained: + warnings.warn( + f"Snapshot {snapshot.id} was taken of graph " + f"{snapshot.graph_key}, but this model was trained on a " + f"snapshot of graph {trained}.", UserWarning, stacklevel=2) query = PredictiveQuery.from_snapshot(snapshot, self._training_query) plan = self._training_table_generation_plan table = (query.generate_prediction_table( diff --git a/packages/kumo-engine/test/test_query_from_snapshot.py b/packages/kumo-engine/test/test_query_from_snapshot.py index 87866165c..da611aa38 100644 --- a/packages/kumo-engine/test/test_query_from_snapshot.py +++ b/packages/kumo-engine/test/test_query_from_snapshot.py @@ -60,6 +60,12 @@ def scores(**options): pd.testing.assert_frame_equal(scores(entities=rows), expected) with pytest.raises(ValueError, match='not both'): scores(entities=rows, prediction_time=graph_table.timestamps[0]) + other = replace(snapshot, id='other@' + snapshot.id.split('@')[1]) + with pytest.warns(UserWarning, match='trained on a snapshot of graph'): + scored = trainer.predict_snapshot( + other, binary_classification_threshold=0.5).frame + pd.testing.assert_frame_equal( + scored.sort_values('ENTITY').reset_index(drop=True), expected) def test_exported_snapshot_scores_in_a_fresh_scratch_without_a_graph( From c5b4bc8fdf3676fb11a57f4086a9f919c2172724 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 24 Oct 2026 18:03:16 +0000 Subject: [PATCH 08/10] One predict path: predict(snapshot, ...); results predict too predict_snapshot sat next to predict, so scoring had two entry points with overlapping arguments. - Trainer.predict(snapshot, *, table=, entities=, prediction_time=, ...) scores a snapshot with the trained query and generation plan (entities or a generated table instead when given) and warns when the snapshot is of another graph. main's predict(prediction_table, snapshot=...) still works, positionally or by keyword. Argument checks run first, as before. - TrainingResult.predict(...) scores with the model that produced or loaded it (TrainingResult._trainer, set by fit and load). - predict_snapshot is gone (never released); Trainer.training_query and GraphSnapshot.native_graph are private. Docs move to docs/mlops.md. Verified: test_query_from_snapshot.py (snapshot, entities, table and legacy forms, other-graph warning, a fresh scratch scoring an exported snapshot through a loaded result), prediction workers, output columns, completion, preflight, registry, warm start, embedding and file connector suites. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/docs/mlops.md | 17 +++ packages/kumo-engine/kumo/engine/graph.py | 2 +- .../kumo/engine/pquery/predictive_query.py | 15 +- .../kumo/engine/trainer/trainer.py | 138 ++++++++++-------- .../test/test_query_from_snapshot.py | 21 ++- 5 files changed, 122 insertions(+), 71 deletions(-) diff --git a/packages/kumo-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index 8ebc092c0..bc2b84697 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -67,3 +67,20 @@ snapshot = GraphSnapshot.load_from_connector(S3Connector(), path=location) - **Destinations:** any destination whose sink holds files at the path works: S3, GCS, local files, or a volume or stage once those sinks are installed. A warehouse connector raises. With MLOps on (`graph.snapshot(mlops=...)` or a `kumo.mlops` block), each snapshot is exported under `/snapshots/`, unless a complete export of it is already there. + +### Scoring without the graph + +A scoring job does not need the training job's `Graph` code. It loads the snapshot the training job exported and the model, then scores: + +```python +snapshot = GraphSnapshot.load_from_connector(store, path=location) +trainer = Trainer.load_from_connector(store, path="models/") +result = trainer.predict(snapshot, binary_classification_threshold=0.5) +``` + +- `predict(snapshot)` reuses the trained query and table-generation plan, and scores its population at `prediction_time`. +- `entities=` scores given rows instead; `table=` scores a prediction table generated beforehand. +- It warns when the snapshot was taken of a different graph than the one the model was trained on. +- `predict(prediction_table, snapshot=...)` keeps working as before. +- `fit` results score too: `result.predict(snapshot, ...)`. +- `PredictiveQuery.from_snapshot(snapshot, query)` binds a query to a snapshot without its `Graph`. Its graph check only means the snapshot belongs to the query's graph, not that it is compatible with a trained model. diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index 4174ca94d..e82f13cc4 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -327,7 +327,7 @@ def _data_version(self) -> str: return Path(self.index_path).name @cached_property - def native_graph(self) -> NativeGraph: + def _native_graph(self) -> NativeGraph: """`graph_config` in `kumo_core`'s normalized form, as `Graph.native_graph`.""" return normalize_graph(self.graph_config) diff --git a/packages/kumo-engine/kumo/engine/pquery/predictive_query.py b/packages/kumo-engine/kumo/engine/pquery/predictive_query.py index d36621e05..7ac4a9d8d 100644 --- a/packages/kumo-engine/kumo/engine/pquery/predictive_query.py +++ b/packages/kumo-engine/kumo/engine/pquery/predictive_query.py @@ -59,6 +59,7 @@ ) from kumo_core.model_plan_conversion import public_model_plan_from_internal from kumo_core.pquery import ( + NativeGraph, ValidatedNativeQuery, parse_and_validate_query, query_is_temporal, @@ -140,16 +141,22 @@ def _parsed(self) -> ValidatedNativeQuery: `ValidatedNativeQuery` are all plain `str`; `get_task_type` below wraps `task_type` into the shared, public `kumoapi.task.TaskType`. """ - return parse_and_validate_query(self.query, self.graph.native_graph) + return parse_and_validate_query(self.query, self._native_graph) + + @property + def _native_graph(self) -> NativeGraph: + if isinstance(self.graph, GraphSnapshot): + return self.graph._native_graph + return self.graph.native_graph def _require_matching_snapshot(self, snapshot: GraphSnapshot) -> None: if isinstance(self.graph, GraphSnapshot): if snapshot.id == self.graph.id: return - graph_key = self.graph.graph_key + graph_key = self.graph._graph_key else: graph_key = self.graph.id - if snapshot.graph_key != graph_key: + if snapshot._graph_key != graph_key: raise ValueError( f"snapshot {snapshot.id!r} was not materialized from " f"this query's graph (id {graph_key!r}) -- pass a " @@ -185,7 +192,7 @@ def get_task_type(self) -> TaskType: def _model_plan_validation_facts(self, run_mode: RunMode) -> dict[str, Any]: parsed = self._parsed - graph = self.graph.native_graph + graph = self._native_graph entity_tables = [parsed.entity_table] if parsed.rhs_node_type is not None: entity_tables.append(parsed.rhs_node_type) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 6a5ba60b3..82612fe8f 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -10,7 +10,7 @@ import logging import time import warnings -from dataclasses import dataclass, replace +from dataclasses import dataclass, field, replace from datetime import datetime, timezone from pathlib import Path from typing import ( @@ -21,6 +21,7 @@ Mapping, Optional, TypeVar, + Union, cast, ) @@ -94,6 +95,8 @@ class TrainingResult: holdout_path: Optional[Path] = None holdout_truncated: Optional[bool] = False bundle_id: Optional[str] = None + _trainer: Optional['Trainer'] = field(default=None, repr=False, + compare=False) @property def id(self) -> Optional[str]: @@ -107,6 +110,13 @@ def manifest(self) -> Optional[ModelManifest]: return read_provenance(self.bundle_dir) + def predict(self, *args: Any, **kwargs: Any) -> 'PredictionResult': + """Score with this model, as :meth:`Trainer.predict`.""" + if self._trainer is None: + raise RuntimeError('This result has no model to predict with; ' + 'use Trainer.load on its bundle_dir.') + return self._trainer.predict(*args, **kwargs) + def holdout_df(self) -> pd.DataFrame: r"""Read the holdout predictions as a pandas DataFrame. @@ -345,11 +355,6 @@ def training_table_generation_plan( """Effective table-generation settings saved by the last fit/load.""" return self._training_table_generation_plan - @property - def training_query(self) -> Optional[str]: - """Predictive query of the last fit/load.""" - return self._training_query - @property def scratch_dir(self) -> Path: return trainer_scratch_dir(self._id) @@ -679,6 +684,7 @@ def _finish_fit( if settings is not None and settings.wants('model'): _store_model(training_result, settings) + training_result = replace(training_result, _trainer=self) self._training_result = training_result self._snapshot_id = snapshot_id self._training_query = predictive_query @@ -686,56 +692,98 @@ def _finish_fit( return training_result - def predict_snapshot( + def predict( self, - snapshot: GraphSnapshot, + source: Union[GraphSnapshot, PredictionTable, None] = None, + /, *, + snapshot: Optional[GraphSnapshot] = None, + table: Optional[PredictionTable] = None, entities: Optional[SourceTable] = None, prediction_time: Optional[datetime] = None, output_config: Optional[PredictionOutputConfig] = None, num_classes_to_return: Optional[int] = None, num_workers: int = 1, binary_classification_threshold: Optional[float] = None, + prediction_table: Optional[PredictionTable] = None, ) -> PredictionResult: - r"""Score ``snapshot`` with the query and table-generation plan saved - by the last fit/load. - - Without ``entities``, scores the query's population at - ``prediction_time`` (as :meth:`PredictiveQuery.generate_prediction_table`). - ``entities`` is a source table of ``ENTITY1`` (and ``TIMESTAMP`` where - the task needs one) rows to score instead. Other arguments are as in - :meth:`predict`. Warns when ``snapshot`` was taken of a different - graph than the one the model was trained on. + r"""Score a snapshot and return durable file-backed results. + + ``predict(snapshot)`` scores the trained query's population at + ``prediction_time`` (as :meth:`PredictiveQuery.generate_prediction_table`); + ``entities`` is a source table of ``ENTITY1`` (and ``TIMESTAMP`` + where the task needs one) rows to score instead, and ``table`` a + prediction table generated beforehand. It warns when the snapshot + was taken of another graph than the one the model was trained on. + ``predict(prediction_table, snapshot=...)`` is accepted as before. + + The trainer must contain fitted state from :meth:`fit`, :meth:`load`, + or :meth:`load_from_connector`. Without ``output_config``, + predictions remain on local disk. + + ``num_workers`` controls scorer processes, not DataLoader workers. + The default stays in-process; values above one spawn independent + scorers over disjoint partitions, capped by the prediction population. + CUDA execution requires one visible device per concurrent scorer. + + ``num_classes_to_return`` must be a positive integer when supplied. + Link predictions require an explicit count of at most 500. Embeddings + alone do not require a count. + + Binary classification requires an explicit threshold. Multilabel + classification defaults to 0.5. Both accept thresholds from 0 to 1; + thresholds change predicted labels without changing class scores. """ + from kumo_core.prediction_options import ( + validate_num_classes_to_return, + validate_prediction_workers, + ) + + validate_prediction_workers(num_workers) + validate_num_classes_to_return(num_classes_to_return) + if source is not None and not isinstance(source, GraphSnapshot): + source, prediction_table = None, source + snapshot = snapshot if snapshot is not None else source + if snapshot is None: + raise TypeError('predict() needs a snapshot') + table = table if table is not None else prediction_table + if table is None: + table = self._prediction_table(snapshot, entities, + prediction_time) + return self._predict( + table, snapshot=snapshot, output_config=output_config, + num_classes_to_return=num_classes_to_return, + num_workers=num_workers, + binary_classification_threshold=binary_classification_threshold) + + def _prediction_table(self, snapshot: GraphSnapshot, + entities: Optional[SourceTable], + prediction_time: Optional[datetime] + ) -> PredictionTable: from kumo.engine.pquery.predictive_query import PredictiveQuery if self._training_query is None: raise RuntimeError( - "Trainer.predict_snapshot() requires a prior Trainer.fit() or " + "Trainer.predict() requires a prior Trainer.fit() or " "Trainer.load() call on this same Trainer instance.") if entities is not None and prediction_time is not None: raise ValueError("Pass prediction_time or entities, not both; " "put anchor times in the entities' TIMESTAMP.") trained = (self._snapshot_id or '').partition('@')[0] - if trained and snapshot.graph_key != trained: + if trained and snapshot._graph_key != trained: warnings.warn( f"Snapshot {snapshot.id} was taken of graph " - f"{snapshot.graph_key}, but this model was trained on a " - f"snapshot of graph {trained}.", UserWarning, stacklevel=2) + f"{snapshot._graph_key}, but this model was trained on a " + f"snapshot of graph {trained}.", UserWarning, stacklevel=3) query = PredictiveQuery.from_snapshot(snapshot, self._training_query) plan = self._training_table_generation_plan - table = (query.generate_prediction_table( + return (query.generate_prediction_table( snapshot, prediction_time=prediction_time, plan=plan) - if entities is None else PredictionTable.from_source_table( - entities, pq=query, snapshot=snapshot, plan=plan)) - return self.predict( - table, snapshot=snapshot, output_config=output_config, - num_classes_to_return=num_classes_to_return, - num_workers=num_workers, - binary_classification_threshold=binary_classification_threshold) + if entities is None else PredictionTable.from_source_table( + entities, pq=query, snapshot=snapshot, plan=plan)) @_without_torch_scatter_hint - def predict( + def _predict( self, prediction_table: PredictionTable, *, @@ -745,33 +793,6 @@ def predict( num_workers: int = 1, binary_classification_threshold: Optional[float] = None, ) -> PredictionResult: - r"""Run batch prediction and return durable file-backed results. - - The trainer must contain fitted state from :meth:`fit`, :meth:`load`, - or :meth:`load_from_connector`. ``snapshot`` selects graph data and - features independently of the table's retained population and anchor - times. Without ``output_config``, predictions remain on local disk. - - ``num_workers`` controls scorer processes, not DataLoader workers. - The default stays in-process; values above one spawn independent - scorers over disjoint partitions, capped by the prediction population. - CUDA execution requires one visible device per concurrent scorer. - - ``num_classes_to_return`` must be a positive integer when supplied. - Link predictions require an explicit count of at most 500. Embeddings - alone do not require a count. - - Binary classification requires an explicit threshold. Multilabel - classification defaults to 0.5. Both accept thresholds from 0 to 1; - thresholds change predicted labels without changing class scores. - """ - from kumo_core.prediction_options import ( - validate_num_classes_to_return, - validate_prediction_workers, - ) - - validate_prediction_workers(num_workers) - validate_num_classes_to_return(num_classes_to_return) if self._training_result is None: raise RuntimeError( "Trainer.predict() requires a prior Trainer.fit() call on " @@ -1015,5 +1036,6 @@ def _load_bundle( trainer._training_table_generation_plan = from_json_dict( spec.training_table_generation, TrainingTableGenerationPlan) trainer._training_result = TrainingResult(metrics, paths, bundle_dir, - bundle_id=bundle_id) + bundle_id=bundle_id, + _trainer=trainer) return trainer diff --git a/packages/kumo-engine/test/test_query_from_snapshot.py b/packages/kumo-engine/test/test_query_from_snapshot.py index da611aa38..488d09d0d 100644 --- a/packages/kumo-engine/test/test_query_from_snapshot.py +++ b/packages/kumo-engine/test/test_query_from_snapshot.py @@ -20,8 +20,8 @@ def _scores(trainer, table, snapshot): def test_snapshot_bound_query_predicts_like_graph_bound(fitted, tmp_path): trainer, graph_table, snapshot = fitted - assert trainer.training_query == graph_table.query - query = kumo.PredictiveQuery.from_snapshot(snapshot, trainer.training_query) + assert trainer._training_query == graph_table.query + query = kumo.PredictiveQuery.from_snapshot(snapshot, trainer._training_query) expected = _scores(trainer, graph_table, snapshot) assert query.get_task_type() == TaskType.BINARY_CLASSIFICATION @@ -42,12 +42,12 @@ def test_snapshot_bound_query_predicts_like_graph_bound(fitted, tmp_path): expected) -def test_predict_snapshot_scores_with_the_saved_query(fitted, tmp_path): +def test_predict_scores_a_snapshot_with_the_saved_query(fitted, tmp_path): trainer, graph_table, snapshot = fitted expected = _scores(trainer, graph_table, snapshot) def scores(**options): - return trainer.predict_snapshot( + return trainer.predict( snapshot, binary_classification_threshold=0.5, **options).frame.sort_values('ENTITY').reset_index(drop=True) @@ -58,11 +58,16 @@ def scores(**options): }).to_parquet(tmp_path / 'rows.parquet') rows = kumo.FileConnector(root_dir=str(tmp_path))['rows.parquet'] pd.testing.assert_frame_equal(scores(entities=rows), expected) + pd.testing.assert_frame_equal(scores(table=graph_table), expected) + legacy = trainer.predict(prediction_table=graph_table, snapshot=snapshot, + binary_classification_threshold=0.5).frame + pd.testing.assert_frame_equal( + legacy.sort_values('ENTITY').reset_index(drop=True), expected) with pytest.raises(ValueError, match='not both'): scores(entities=rows, prediction_time=graph_table.timestamps[0]) other = replace(snapshot, id='other@' + snapshot.id.split('@')[1]) with pytest.warns(UserWarning, match='trained on a snapshot of graph'): - scored = trainer.predict_snapshot( + scored = trainer.predict( other, binary_classification_threshold=0.5).frame pd.testing.assert_frame_equal( scored.sort_values('ENTITY').reset_index(drop=True), expected) @@ -73,7 +78,7 @@ def test_exported_snapshot_scores_in_a_fresh_scratch_without_a_graph( trainer, graph_table, snapshot = fitted expected = _scores(trainer, graph_table, snapshot) exports = kumo.FileConnector(root_dir=str(tmp_path / 'exports')) - location = snapshot.export(exports, 'snapshots') + location = snapshot.export(kumo.SnapshotOutputConfig(exports, 'snapshots')) bundle = trainer._training_result.bundle_dir previous = get_scratch_dir() kumo.configure(scratch_dir=tmp_path / 'fresh') @@ -82,11 +87,11 @@ def test_exported_snapshot_scores_in_a_fresh_scratch_without_a_graph( assert loaded.id == snapshot.id assert loaded.index_path.is_relative_to(tmp_path / 'fresh') model = kumo.Trainer.load(bundle) - scored = model.predict_snapshot( + scored = model.result.predict( loaded, binary_classification_threshold=0.5).frame pd.testing.assert_frame_equal( scored.sort_values('ENTITY').reset_index(drop=True), expected) - query = kumo.PredictiveQuery.from_snapshot(loaded, model.training_query) + query = kumo.PredictiveQuery.from_snapshot(loaded, model._training_query) table = query.generate_prediction_table( loaded, plan=model.training_table_generation_plan) pd.testing.assert_frame_equal(_scores(model, table, loaded), expected) From 4fd38b9f7e9b7cb9334dfabb72584bbab811ae2d Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 24 Oct 2026 19:10:01 +0000 Subject: [PATCH 09/10] Hide the torch-scatter hint across all of predict() Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/trainer/trainer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 82612fe8f..6e903e98d 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -692,6 +692,7 @@ def _finish_fit( return training_result + @_without_torch_scatter_hint def predict( self, source: Union[GraphSnapshot, PredictionTable, None] = None, @@ -782,7 +783,6 @@ def _prediction_table(self, snapshot: GraphSnapshot, if entities is None else PredictionTable.from_source_table( entities, pq=query, snapshot=snapshot, plan=plan)) - @_without_torch_scatter_hint def _predict( self, prediction_table: PredictionTable, From 803cd82e2cdfa251e2a7cfa39f297243b6742816 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 24 Oct 2026 19:40:50 +0000 Subject: [PATCH 10/10] Compare reloaded-model scores within float tolerance Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/test/test_query_from_snapshot.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/packages/kumo-engine/test/test_query_from_snapshot.py b/packages/kumo-engine/test/test_query_from_snapshot.py index 488d09d0d..70cd9e66b 100644 --- a/packages/kumo-engine/test/test_query_from_snapshot.py +++ b/packages/kumo-engine/test/test_query_from_snapshot.py @@ -90,11 +90,13 @@ def test_exported_snapshot_scores_in_a_fresh_scratch_without_a_graph( scored = model.result.predict( loaded, binary_classification_threshold=0.5).frame pd.testing.assert_frame_equal( - scored.sort_values('ENTITY').reset_index(drop=True), expected) + scored.sort_values('ENTITY').reset_index(drop=True), expected, + atol=1e-5) query = kumo.PredictiveQuery.from_snapshot(loaded, model._training_query) table = query.generate_prediction_table( loaded, plan=model.training_table_generation_plan) - pd.testing.assert_frame_equal(_scores(model, table, loaded), expected) + pd.testing.assert_frame_equal(_scores(model, table, loaded), expected, + atol=1e-5) finally: kumo.configure(scratch_dir=previous)