Score a snapshot with Trainer.predict(snapshot) and TrainingResult.predict

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 `<root>/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/<bundle id>")
+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.

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 `<path>/<graph id>/<data version[:16]>` and
         return that location. Any earlier export there is replaced; it is not

packages/kumo-engine/kumo/engine/pquery/predictive_query.py

@@ -6,15 +6,17 @@
 
 `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 and parses against `snapshot.graph_config`.
+
+Parsing (`kumo_core.pquery.parse_and_validate_query` against
+`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
+`_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 +27,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,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,
@@ -97,15 +100,27 @@
 
 @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.
+
+        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
     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.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`
         uses internally for `fit`. No materialized `GraphSnapshot` needed
@@ -125,23 +140,27 @@ 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`.
-
-        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.
         """
-        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:
-        # 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 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"self.graph (id {self.graph.id!r}) -- call "
-                "self.graph.snapshot() and pass its result.")
+                f"this query's graph (id {graph_key!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 +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)
@@ -195,7 +214,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,15 +242,15 @@ 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
         generated artifacts, required split stages, and effective generation
         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 +349,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 +361,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.

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,
 )
 
@@ -55,6 +56,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
@@ -93,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]:
@@ -106,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.
 
@@ -673,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
@@ -683,20 +695,32 @@ def _finish_fit(
     @_without_torch_scatter_hint
     def predict(
         self,
-        prediction_table: PredictionTable,
+        source: Union[GraphSnapshot, PredictionTable, None] = None,
+        /,
         *,
-        snapshot: GraphSnapshot,
+        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"""Run batch prediction and return durable file-backed results.
+        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`. ``snapshot`` selects graph data and
-        features independently of the table's retained population and anchor
-        times. Without ``output_config``, predictions remain on local disk.
+        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
@@ -718,6 +742,57 @@ def predict(
 
         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() 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:
+            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=3)
+        query = PredictiveQuery.from_snapshot(snapshot, self._training_query)
+        plan = self._training_table_generation_plan
+        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))
+
+    def _predict(
+        self,
+        prediction_table: PredictionTable,
+        *,
+        snapshot: GraphSnapshot,
+        output_config: Optional[PredictionOutputConfig] = None,
+        num_classes_to_return: Optional[int] = None,
+        num_workers: int = 1,
+        binary_classification_threshold: Optional[float] = None,
+    ) -> PredictionResult:
         if self._training_result is None:
             raise RuntimeError(
                 "Trainer.predict() requires a prior Trainer.fit() call on "
@@ -961,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