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 3c182605a..e82f13cc4 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 5cbf4771c..7ac4a9d8d 100644 --- a/packages/kumo-engine/kumo/engine/pquery/predictive_query.py +++ b/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,7 +242,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 +250,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 +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. diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 1acc5e100..6e903e98d 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, ) @@ -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 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..70cd9e66b --- /dev/null +++ b/packages/kumo-engine/test/test_query_from_snapshot.py @@ -0,0 +1,115 @@ +# 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 +from kumo.engine._scratch import get_scratch_dir + + +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 + 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 + 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_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, 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) + 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( + 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( + 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(kumo.SnapshotOutputConfig(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.result.predict( + loaded, binary_classification_threshold=0.5).frame + pd.testing.assert_frame_equal( + 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, + atol=1e-5) + 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) + newer = replace(snapshot, id=snapshot.id.split('@')[0] + '@newer') + other = replace(snapshot, id='other@' + snapshot.id.split('@')[1]) + plain = replace(snapshot, id='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.suggest_model_plan(other)