diff --git a/packages/kumo-core/src/kumo_core/artifacts.py b/packages/kumo-core/src/kumo_core/artifacts.py index 838f4c0c4..c8e45e4a4 100644 --- a/packages/kumo-core/src/kumo_core/artifacts.py +++ b/packages/kumo-core/src/kumo_core/artifacts.py @@ -13,7 +13,14 @@ import tempfile from dataclasses import dataclass, field from pathlib import Path -from typing import Any, ClassVar, Mapping, Protocol, runtime_checkable +from typing import ( + Any, + ClassVar, + Iterable, + Mapping, + Protocol, + runtime_checkable, +) from .io.filesystem import atomic_write_bytes, fsync_directory, local_path from .io.json import load_json, save_json @@ -248,14 +255,17 @@ def _sha256(path: Path) -> str: return digest.hexdigest() +#: Earlier name of the bundle manifest; also kept out of the bundle id. +_OLD_PROVENANCE_FILE = 'provenance.json' + + def verify_model_bundle(bundle: Path) -> str: """Verify fetched artifacts before parsing or staging a model bundle. - Preserves the local SDK import verifier's error and content-identity - contract. Paths must stay inside the bundle after symlink resolution; - declared sizes and SHA-256 digests must match. Older manifests may omit - sizes or leave training-spec.json unlisted; its digest still participates - in the returned content identity. + Returns the bundle id, a SHA-256 over every verified file except + ``provenance.json``. Paths must stay inside the bundle after symlink + resolution; declared sizes and SHA-256 digests must match. Older + manifests may omit sizes or leave training-spec.json unlisted. """ from .sources import layout from .sources.errors import PermanentSourceError @@ -273,7 +283,7 @@ def verify_model_bundle(bundle: Path) -> str: artifacts = manifest.get('artifacts') problems: list[str] = [] - verified: list[tuple[str, str, str]] = [] + verified: list[tuple[str, str]] = [] seen = set() claimed = set() if not artifacts: @@ -325,7 +335,7 @@ def verify_model_bundle(bundle: Path) -> str: if VitalArtifactPublicationPolicy._sha256(path) != expected: problems.append(f'{relative} does not match its digest') continue - verified.append((role, safe, expected)) + verified.append((safe, expected)) if problems: raise ValueError( f"The model bundle at {bundle} did not survive the transfer: " @@ -335,11 +345,47 @@ def verify_model_bundle(bundle: Path) -> str: if not spec_path.resolve().is_relative_to(bundle): raise ValueError('Training specification is not below the bundle') if 'training-spec.json' not in seen and spec_path.is_file(): - verified.append(('training_spec', 'training-spec.json', + verified.append(('training-spec.json', VitalArtifactPublicationPolicy._sha256(spec_path))) - payload = json.dumps(sorted(verified), sort_keys=True, - separators=(',', ':'), default=str) - return hashlib.sha256(payload.encode('utf-8')).hexdigest()[:16] + return _bundle_id(verified) + + +def bundle_id_from_manifest(bundle: Path) -> str: + """The bundle id from the digests a freshly published bundle lists.""" + manifest_path = Path(bundle) / VitalArtifactPublicationPolicy.MANIFEST_NAME + manifest = json.loads(manifest_path.read_text(encoding='utf-8')) + return _bundle_id((str(entry['relative_path']), str(entry['sha256'])) + for entry in manifest['artifacts']) + + +def _bundle_id(pairs: Iterable[tuple[str, str]]) -> str: + payload = json.dumps( + sorted(pair for pair in pairs + if pair[0] not in (PROVENANCE_FILE, _OLD_PROVENANCE_FILE)), + separators=(',', ':')) + return hashlib.sha256(payload.encode('utf-8')).hexdigest() + + +def add_bundle_file(bundle_dir: Path, role: str, relative_path: str) -> None: + """List one file already written in a bundle in its publication manifest.""" + bundle_dir = Path(bundle_dir) + path = bundle_dir / relative_path + manifest_path = bundle_dir / VitalArtifactPublicationPolicy.MANIFEST_NAME + manifest = json.loads(manifest_path.read_text(encoding='utf-8')) + manifest['artifacts'] = [ + entry for entry in manifest.get('artifacts') or [] + if entry.get('relative_path') != relative_path + and entry.get('role') != role + ] + [ + PublishedArtifact( + role=role, + relative_path=relative_path, + size_bytes=path.stat().st_size, + sha256=VitalArtifactPublicationPolicy._sha256(path), + ).to_dict() + ] + payload = json.dumps(manifest, indent=2, sort_keys=True) + '\n' + atomic_write_bytes(manifest_path, payload.encode('utf-8')) @runtime_checkable diff --git a/packages/kumo-core/tests/test_artifacts.py b/packages/kumo-core/tests/test_artifacts.py index 071f4339b..8828667ae 100644 --- a/packages/kumo-core/tests/test_artifacts.py +++ b/packages/kumo-core/tests/test_artifacts.py @@ -15,7 +15,10 @@ FilesystemVitalArtifactPublisher, VitalArtifactPublicationPolicy, VitalArtifactSources, + add_bundle_file, + bundle_id_from_manifest, resolve_vital_artifact_sources, + verify_model_bundle, ) from kumo_core.config import NativeRuntimeConfig from kumo_core.manifest import RunManifest, RunPaths @@ -237,3 +240,27 @@ def best_ckpt(self): assert sources.selected_checkpoint == checkpoints[selected_mode] assert sources.graph_manifest == bundle / "graph.json" assert json.loads(sources.metrics.read_text()) == {"test_accuracy": 0.8} + + +def test_bundle_id_ignores_provenance_only(tmp_path: Path) -> None: + sources, _ = _vital_sources(tmp_path) + bundle = tmp_path / "bundle" + VitalArtifactPublicationPolicy().publish(sources, bundle) + bundle_id = verify_model_bundle(bundle) + + for name, payload in (("_provenance.json", '{"a": 1}'), + ("_provenance.json", '{"a": 2}'), + ("provenance.json", '{"a": 3}')): + (bundle / name).write_text(payload) + add_bundle_file(bundle, "provenance", name) + assert verify_model_bundle(bundle) == bundle_id + assert bundle_id_from_manifest(bundle) == bundle_id + roles = [ + entry["role"] for entry in json.loads( + (bundle / "publication-manifest.json").read_text())["artifacts"] + ] + assert roles.count("provenance") == 1 + + (bundle / "metrics.json").write_text('{"auc": 0}') + add_bundle_file(bundle, "metrics", "metrics.json") + assert verify_model_bundle(bundle) != bundle_id diff --git a/packages/kumo-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index e1480f29e..5c93f0f1f 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -35,3 +35,14 @@ A generated or caller-supplied training table carries `_provenance.json` in its `TrainingTable.id` is deterministic: the same snapshot, query, plan, split seed and engine build give the same id. A caller-supplied source without a version gets a random id. An export to files (`TrainingTable.export`) writes the manifest next to the Parquet parts. `TrainingTable.load_from_connector(connector, path=..., pq=..., snapshot=...)` materializes an export again. With MLOps on, each table is stored once under `/training_tables//`. + +### Models + +Every fitted bundle carries `_provenance.json` (`ModelManifest`): +- its `id`, the bundle id (a content hash that leaves the manifest file out); +- when it was fitted and with which fit mode; +- its parents: the training table's manifest, embedded with that table's snapshot, and the warm-start model's manifest, if any. + +An embedded model keeps only the ids of its own model parents, so a long chain of warm-start retrains does not grow each manifest. `fit` returns a new `TrainingResult` per call (`.id`, `.manifest`); `Trainer.result` and `Trainer.manifest` describe the latest fit or the loaded bundle. + +With MLOps on, each bundle is stored under `/models//`. A complete copy that is already there is kept; a partial one is replaced. diff --git a/packages/kumo-engine/kumo/engine/__init__.py b/packages/kumo-engine/kumo/engine/__init__.py index 52b0540ae..e2d3c1500 100644 --- a/packages/kumo-engine/kumo/engine/__init__.py +++ b/packages/kumo-engine/kumo/engine/__init__.py @@ -25,7 +25,11 @@ from kumoapi.plugin import ModelRegistryLocation from kumoapi.table import Column -from kumo.engine._manifest import SnapshotManifest, TrainingTableManifest +from kumo.engine._manifest import ( + ModelManifest, + SnapshotManifest, + TrainingTableManifest, +) from kumo.engine._mlops import MLOps, mlops from kumo.engine._version import __version__ from kumo.engine.config import configure @@ -89,6 +93,7 @@ "mlops", "SnapshotManifest", "TrainingTableManifest", + "ModelManifest", "PredictiveQuery", "TrainingTable", "PredictionTable", diff --git a/packages/kumo-engine/kumo/engine/_manifest.py b/packages/kumo-engine/kumo/engine/_manifest.py index a47a0737f..81265178f 100644 --- a/packages/kumo-engine/kumo/engine/_manifest.py +++ b/packages/kumo-engine/kumo/engine/_manifest.py @@ -224,6 +224,32 @@ class TrainingTableManifest(Manifest): split_seed: Optional[int] = None +@dataclass(frozen=True, kw_only=True) +class ModelParents(Record): + training_table: Optional[TrainingTableManifest] = None + warm_start: Optional['ModelManifest'] = None + derived_from: Optional['ModelManifest'] = omit_none() + + +@dataclass(frozen=True, kw_only=True) +class FitTiming(Record): + started_at: str + finished_at: str + seconds: float + + +@dataclass(frozen=True, kw_only=True) +class ModelManifest(Manifest): + """How a model bundle was fitted: from which training table (embedded, + with its snapshot) and warm-start model, when, and with which fit + mode. Its ``id`` is the bundle id.""" + + kind: str = 'model' + parents: Optional[ModelParents] = None + fit: Optional[FitTiming] = None + fit_mode: Optional[str] = None + + def write_manifest(path: Path, manifest: Manifest) -> None: payload = json.dumps(manifest.to_dict(), indent=2, sort_keys=True) + '\n' atomic_write_bytes(path, payload.encode('utf-8')) diff --git a/packages/kumo-engine/kumo/engine/trainer/_provenance.py b/packages/kumo-engine/kumo/engine/trainer/_provenance.py new file mode 100644 index 000000000..4429186dd --- /dev/null +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -0,0 +1,87 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""A model bundle's ``_provenance.json``; excluded from the bundle id.""" + +from dataclasses import replace +from datetime import datetime, timezone +from pathlib import Path +from typing import Optional + +from kumo.engine._manifest import ( + FitTiming, + ModelManifest, + ModelParents, + embedded, + read_manifest, + stamp, + stub, + write_manifest, +) +from kumo.engine.pquery.training_table import TrainingTable +from kumo.engine.trainer._spec import PROVENANCE_ROLE, listed_file + + +def write_provenance( + bundle_dir: Path, + *, + bundle_id: str, + train_table: TrainingTable, + warm_start: Optional[ModelManifest], + started_at: datetime, + fit_mode: Optional[str], +) -> None: + from kumo_core.artifacts import PROVENANCE_FILE, add_bundle_file + + finished_at = datetime.now(timezone.utc) + table = train_table.manifest + write_manifest( + bundle_dir / PROVENANCE_FILE, + ModelManifest( + id=bundle_id, + parents=ModelParents( + training_table=embedded(table) if table is not None else + stub('training_table', train_table.id or train_table.tt_id), + warm_start=None + if warm_start is None else embedded_model(warm_start)), + fit=FitTiming( + started_at=started_at.isoformat(timespec='seconds'), + finished_at=finished_at.isoformat(timespec='seconds'), + seconds=round((finished_at - started_at).total_seconds(), 3), + ), + fit_mode=fit_mode, + **stamp(), + )) + add_bundle_file(bundle_dir, PROVENANCE_ROLE, PROVENANCE_FILE) + + +def embedded_model(manifest: ModelManifest) -> ModelManifest: + r"""``manifest`` as a child model embeds it: its own model parents cut + to stubs, so a chain of retrains does not grow each manifest.""" + copy = embedded(manifest) + parents = copy.parents + if parents is None: + return copy + + def cut(model: Optional[ModelManifest]) -> Optional[ModelManifest]: + return None if model is None else ModelManifest(id=model.id, + reconstructed=True) + + return replace(copy, parents=replace( + parents, warm_start=cut(parents.warm_start), + derived_from=cut(parents.derived_from))) + + +def parent_model(bundle_dir: Path, bundle_id: Optional[str] + ) -> Optional[ModelManifest]: + r"""The manifest of the bundle at ``bundle_dir``, or a stub.""" + manifest = read_provenance(bundle_dir) + if manifest is not None: + return manifest + return None if bundle_id is None else ModelManifest(id=bundle_id, + reconstructed=True) + + +def read_provenance(bundle_dir: Path) -> Optional[ModelManifest]: + path = listed_file(bundle_dir, PROVENANCE_ROLE) + return None if path is None else read_manifest(path, ModelManifest) diff --git a/packages/kumo-engine/kumo/engine/trainer/_spec.py b/packages/kumo-engine/kumo/engine/trainer/_spec.py index 147817c0e..2c755fab4 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_spec.py +++ b/packages/kumo-engine/kumo/engine/trainer/_spec.py @@ -9,11 +9,10 @@ separate and must be supplied through a compatible prediction snapshot. """ -import hashlib import json from dataclasses import asdict, dataclass from pathlib import Path -from typing import Any, Dict +from typing import Any, Dict, Optional, Tuple from kumo_core.io.filesystem import atomic_write_bytes from kumoapi.model_plan import ModelPlan, TrainingTableGenerationPlan @@ -23,10 +22,10 @@ SPEC_VERSION = 1 #: What `kumo_core`'s `VitalArtifactPublicationPolicy` names its own -#: manifest. `write_training_spec` amends this file's `artifacts` list so -#: `training-spec.json` is covered the same way every other bundle file is. +#: manifest; `training-spec.json` is listed in it like every other file. MANIFEST_FILE = 'publication-manifest.json' SPEC_ROLE = 'training_spec' +PROVENANCE_ROLE = 'provenance' @dataclass(frozen=True) @@ -45,12 +44,6 @@ class TrainingSpec: training_table_generation: Dict[str, Any] -def _sha256(path: Path) -> str: - digest = hashlib.sha256() - digest.update(path.read_bytes()) - return digest.hexdigest() - - def write_training_spec( bundle_dir: Path, *, @@ -60,6 +53,7 @@ def write_training_spec( training_table_generation: TrainingTableGenerationPlan, ) -> None: r"""Write an integrity-covered training spec into a completed bundle.""" + from kumo_core.artifacts import add_bundle_file from kumo_core.io.json import to_json_dict spec = TrainingSpec( @@ -75,23 +69,7 @@ def write_training_spec( (json.dumps(asdict(spec), indent=2, sort_keys=True) + '\n').encode('utf-8'), ) - - manifest_path = bundle_dir / MANIFEST_FILE - with open(manifest_path, encoding='utf-8') as handle: - manifest = json.load(handle) - manifest['artifacts'] = [ - entry for entry in manifest.get('artifacts') or [] - if entry.get('relative_path') != SPEC_FILE - ] + [{ - 'role': SPEC_ROLE, - 'relative_path': SPEC_FILE, - 'size_bytes': spec_path.stat().st_size, - 'sha256': _sha256(spec_path), - }] - atomic_write_bytes( - manifest_path, - (json.dumps(manifest, indent=2, sort_keys=True) + '\n').encode('utf-8'), - ) + add_bundle_file(bundle_dir, SPEC_ROLE, SPEC_FILE) def read_training_spec(bundle_dir: Path) -> TrainingSpec: @@ -125,11 +103,23 @@ def read_training_spec(bundle_dir: Path) -> TrainingSpec: ) -def read_bundle_artifact_paths(bundle_dir: Path) -> Dict[str, str]: - r"""Rebuild the verified ``role -> path`` artifact mapping.""" +def listed_file(bundle_dir: Path, role: str) -> Optional[Path]: + try: + manifest = json.loads( + (bundle_dir / MANIFEST_FILE).read_text(encoding='utf-8')) + entries = manifest.get('artifacts') or [] + except (OSError, ValueError, AttributeError): + return None + return next((bundle_dir / entry['relative_path'] + for entry in entries if entry.get('role') == role), None) + + +def read_bundle_artifact_paths( + bundle_dir: Path) -> Tuple[Dict[str, str], str]: + r"""Rebuild the verified ``role -> path`` mapping and the bundle id.""" from kumo_core.artifacts import verify_model_bundle - verify_model_bundle(bundle_dir) + bundle_id = verify_model_bundle(bundle_dir) manifest_path = bundle_dir / MANIFEST_FILE if not manifest_path.is_file(): raise ValueError( @@ -150,5 +140,5 @@ def read_bundle_artifact_paths(bundle_dir: Path) -> Dict[str, str]: return { entry['role']: str(bundle_dir / entry['relative_path']) for entry in artifacts - if entry['role'] != SPEC_ROLE - } + if entry['role'] not in (SPEC_ROLE, PROVENANCE_ROLE) + }, bundle_id diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 28c3156a3..1acc5e100 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -52,6 +52,7 @@ require_embedding_models, require_trained_embedding_models, ) +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.graph import GraphSnapshot @@ -91,6 +92,19 @@ class TrainingResult: published_model: Optional[PublishedModel] = None holdout_path: Optional[Path] = None holdout_truncated: Optional[bool] = False + bundle_id: Optional[str] = None + + @property + def id(self) -> Optional[str]: + """The model's id: its bundle id.""" + return self.bundle_id + + @property + def manifest(self) -> Optional[ModelManifest]: + """The bundle's ``_provenance.json``; None for older bundles.""" + from kumo.engine.trainer._provenance import read_provenance + + return read_provenance(self.bundle_dir) def holdout_df(self) -> pd.DataFrame: r"""Read the holdout predictions as a pandas DataFrame. @@ -188,6 +202,41 @@ def _find_holdout(log_dir: str) -> tuple[Optional[Path], Optional[bool]]: return path, truncated +def _store_model(result: TrainingResult, settings: Any) -> None: + r"""Copy the bundle to ``/models//`` unless a complete + copy is there; failures warn.""" + from kumo_core.object_store_sink import upload_bundle + from pyarrow.fs import FileSelector, FileType + + from kumo.engine.connector import file_store + + if result.bundle_id is None: + return + try: + filesystem, root = file_store( + settings.root.connector, + settings.root.location(f'models/{result.bundle_id}')) + local = { + path.relative_to(result.bundle_dir).as_posix(): path.stat().st_size + for path in Path(result.bundle_dir).rglob('*') if path.is_file() + } + remote = { + info.path[len(root) + 1:]: info.size + for info in filesystem.get_file_info( + FileSelector(root, recursive=True, allow_not_found=True)) + if info.type is FileType.File + } + if remote == local: + return + if remote: + filesystem.delete_dir_contents(root, missing_dir_ok=True) + upload_bundle(filesystem, root, str(result.bundle_dir)) + except Exception as error: + warnings.warn( + f"Could not store model {result.bundle_id} under the artifact " + f"root: {error}", UserWarning, stacklevel=4) + + def _published_holdout(trainer: Any) -> tuple[Optional[Path], Optional[bool]]: from kumo_core.training.config import FitMode @@ -299,6 +348,22 @@ def training_table_generation_plan( def scratch_dir(self) -> Path: return trainer_scratch_dir(self._id) + @property + def bundle_id(self) -> Optional[str]: + """Content id of the fitted or loaded model bundle.""" + return getattr(self._training_result, 'bundle_id', None) + + @property + def result(self) -> Optional[TrainingResult]: + """The last fit's result, or the loaded bundle's.""" + return self._training_result + + @property + def manifest(self) -> Optional[ModelManifest]: + """The fitted or loaded bundle's ``_provenance.json``, if any.""" + result = self._training_result + return None if result is None else result.manifest + @_without_torch_scatter_hint def fit( self, @@ -307,6 +372,7 @@ def fit( snapshot: GraphSnapshot, export_config: Optional[ModelOutputConfig] = None, warm_start: Optional[Self] = None, + mlops: Optional[Any] = None, ) -> TrainingResult: r"""Train a model against `train_table`. @@ -335,6 +401,11 @@ def fit( may contain an AutoML search space instead of the fitted architecture; they support prediction but are unreliable warm-start sources. + mlops: a `kumo.MLOps`; with MLOps on (this or a `kumo.mlops` + block), the bundle is also stored under the artifact root. + + Each call returns a new `TrainingResult`; `Trainer.result` holds the + latest. """ from kumo_core.automl.runtime import ( inherit_warm_start_model_plan, @@ -363,18 +434,25 @@ def fit( raise NotImplementedError( f'{type(export_config.connector).__name__} cannot publish ' f'models') + from kumo.engine._mlops import resolve + from kumo.engine.trainer._provenance import parent_model + + settings = resolve(mlops) + started_at = datetime.now(timezone.utc) run_id = f'{self._id}-fit-{new_run_id()}' run_dir = self.scratch_dir / 'runs' / run_id model_plan = _to_core_model_plan(self._model_plan) warm_start_artifacts = None + warm_start_bundle_id = None effective_plan = self._model_plan if warm_start is not None: if warm_start._training_result is None: raise ValueError( "warm_start requires a fitted or loaded Trainer") - warm_start_artifacts = read_bundle_artifact_paths( - warm_start._training_result.bundle_dir) + warm_start_artifacts, warm_start_bundle_id = ( + read_bundle_artifact_paths( + warm_start._training_result.bundle_dir)) inherit_warm_start_model_plan( model_plan, load_json(warm_start_artifacts['resolved_model_plan'], @@ -508,6 +586,8 @@ def fit( holdout_path, holdout_truncated = _published_holdout( result.best_trainer) + fit_mode = result.best_trainer.published_fit_mode() + fit_mode = None if fit_mode is None else fit_mode.value training_result = TrainingResult( metrics=dict(result.metrics), artifact_paths=dict(result.artifact_paths), @@ -524,6 +604,13 @@ def fit( training_table_generation=train_table.generation_plan, export_config=export_config, model_plan=effective_plan, + provenance=dict( + train_table=train_table, + warm_start=None if warm_start is None else parent_model( + warm_start._training_result.bundle_dir, + warm_start_bundle_id), + started_at=started_at, fit_mode=fit_mode), + settings=settings, ) def _finish_fit( @@ -535,8 +622,14 @@ def _finish_fit( training_table_generation: TrainingTableGenerationPlan, export_config: Optional[ModelOutputConfig], model_plan: Optional[ModelPlan] = None, + provenance: Optional[Mapping[str, Any]] = None, + settings: Optional[Any] = None, ) -> TrainingResult: """Commit durable fit outputs before making this trainer usable.""" + from kumo_core.artifacts import bundle_id_from_manifest + + from kumo.engine.trainer._provenance import write_provenance + effective_plan = model_plan if model_plan is not None else self._model_plan write_training_spec( training_result.bundle_dir, @@ -545,6 +638,22 @@ def _finish_fit( snapshot_id=snapshot_id, training_table_generation=training_table_generation, ) + try: + training_result = replace( + training_result, + bundle_id=bundle_id_from_manifest(training_result.bundle_dir)) + except Exception as error: + warnings.warn(f"Could not compute the bundle id: {error}", + UserWarning, stacklevel=3) + if provenance is not None and training_result.bundle_id is not None: + try: + write_provenance(training_result.bundle_dir, + bundle_id=training_result.bundle_id, + **provenance) + except Exception as error: + warnings.warn( + f"Publishing the model bundle without _provenance.json: " + f"{error}", UserWarning, stacklevel=3) if export_config is not None: spec = read_training_spec(training_result.bundle_dir) published = export_config.connector.sink_backend.publish_model( @@ -561,6 +670,8 @@ def _finish_fit( ) training_result = replace(training_result, published_model=published) + if settings is not None and settings.wants('model'): + _store_model(training_result, settings) self._training_result = training_result self._snapshot_id = snapshot_id @@ -837,7 +948,7 @@ def _load_bundle( accelerator: Literal['cpu', 'gpu', 'mps', 'auto'] = 'auto', ) -> Self: bundle_dir = bundle_dir.resolve() - paths = read_bundle_artifact_paths(bundle_dir) + paths, bundle_id = read_bundle_artifact_paths(bundle_dir) spec = read_training_spec(bundle_dir) from kumo_core.io.json import from_json_dict plan = from_json_dict(spec.model_plan, ModelPlan) @@ -849,5 +960,6 @@ def _load_bundle( trainer._training_query = spec.predictive_query trainer._training_table_generation_plan = from_json_dict( spec.training_table_generation, TrainingTableGenerationPlan) - trainer._training_result = TrainingResult(metrics, paths, bundle_dir) + trainer._training_result = TrainingResult(metrics, paths, bundle_dir, + bundle_id=bundle_id) return trainer diff --git a/packages/kumo-engine/test/test_bundle_provenance.py b/packages/kumo-engine/test/test_bundle_provenance.py new file mode 100644 index 000000000..5b44d3153 --- /dev/null +++ b/packages/kumo-engine/test/test_bundle_provenance.py @@ -0,0 +1,164 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json +import re +import shutil +from datetime import datetime, timezone +from types import SimpleNamespace + +import pandas as pd +import pytest +from kumo_core.artifacts import add_bundle_file +from kumoapi.model_plan import ModelPlan, TrainingTableGenerationPlan + +import kumo.engine as kumo +from kumo.engine._manifest import ( + ModelManifest, + ModelParents, + read_manifest, + write_manifest, +) +from kumo.engine.trainer._provenance import embedded_model +from kumo.engine.trainer._spec import MANIFEST_FILE +from kumo.engine.trainer.trainer import TrainingResult + + +def test_fit_records_a_self_contained_manifest(fitted): + trainer, _, snapshot = fitted + result = trainer.result + manifest = result.manifest + assert manifest.kind == 'model' + assert re.fullmatch('[0-9a-f]{64}', trainer.bundle_id) + assert trainer.bundle_id == result.id == manifest.id + assert trainer.manifest == manifest + table = manifest.parents.training_table + assert table.kind == 'training_table' and table.id.startswith('tt-') + assert table.parents.snapshot == snapshot.manifest + assert manifest.parents.warm_start is None + assert manifest.fit_mode == 'refit_trainval' + assert (result.bundle_dir / '_provenance.json').is_file() + loaded = kumo.Trainer.load(result.bundle_dir) + assert loaded.result.id == manifest.id + assert loaded.result.manifest == manifest + + +def test_manifest_round_trips_unchanged(fitted, tmp_path): + trainer, _, _ = fitted + newer = json.loads( + (trainer.result.bundle_dir / '_provenance.json').read_text()) + newer['future'] = True + newer['fit']['gpu_seconds'] = 1.5 + source = tmp_path / 'newer.json' + source.write_text(json.dumps(newer, indent=2, sort_keys=True) + '\n') + + manifest = read_manifest(source, ModelManifest) + write_manifest(tmp_path / 'out.json', manifest) + + assert (tmp_path / 'out.json').read_bytes() == source.read_bytes() + assert manifest.extra == {'future': True} + assert manifest.fit.extra == {'gpu_seconds': 1.5} + + +def test_bundle_without_a_manifest_loads_and_predicts_identically( + fitted, tmp_path): + trainer, prediction_table, snapshot = fitted + result = trainer.result + old = tmp_path / 'old' + shutil.copytree(result.bundle_dir, old) + manifest = json.loads((old / MANIFEST_FILE).read_text()) + manifest['artifacts'] = [ + entry for entry in manifest['artifacts'] + if entry['role'] != 'provenance' + ] + (old / MANIFEST_FILE).write_text(json.dumps(manifest)) + (old / '_provenance.json').unlink() + restored = kumo.Trainer.load(old) + assert restored.bundle_id == result.bundle_id + assert restored.manifest is None + + def predict(model): + return model.predict( + prediction_table, snapshot=snapshot, + binary_classification_threshold=.5).frame.sort_values( + 'ENTITY').reset_index(drop=True) + + pd.testing.assert_frame_equal(predict(restored), predict(trainer)) + + +def test_each_fit_returns_a_new_result(fitted): + trainer, _, _ = fitted + first = trainer.result + second = trainer._finish_fit( + TrainingResult({}, dict(first.artifact_paths), first.bundle_dir), + predictive_query='PREDICT x', snapshot_id='snap', export_config=None, + training_table_generation=TrainingTableGenerationPlan()) + assert second is not first and trainer.result is second + + +def _finish_stub_fit(bundle, trainer, warm_start=None, settings=None): + (bundle / 'model').mkdir(parents=True) + (bundle / 'model' / 'config.json').write_text('{}') + (bundle / MANIFEST_FILE).write_text(json.dumps({'artifacts': []})) + add_bundle_file(bundle, 'trainer_config', 'model/config.json') + table = SimpleNamespace(tt_id='tt-raw', id='tt-1', manifest=None) + return trainer._finish_fit( + TrainingResult({}, {}, bundle), predictive_query='PREDICT x', + snapshot_id='snap', export_config=None, + training_table_generation=TrainingTableGenerationPlan(), + provenance=dict(train_table=table, warm_start=warm_start, + started_at=datetime.now(timezone.utc), fit_mode=None), + settings=settings) + + +def test_missing_parents_are_stubs_and_retrain_chains_stay_flat(tmp_path): + first = _finish_stub_fit(tmp_path / 'one', kumo.Trainer(ModelPlan())) + assert first.manifest.parents.training_table.to_dict() == { + 'kind': 'training_table', 'schema_version': 1, 'id': 'tt-1', + 'created_at': None, 'environment': None, 'reconstructed': True, + 'parents': None, 'query': None, 'generation_plan': None, + 'task_type': None, 'split_seed': None} + second = _finish_stub_fit(tmp_path / 'two', kumo.Trainer(ModelPlan()), + warm_start=first.manifest) + third = _finish_stub_fit(tmp_path / 'three', kumo.Trainer(ModelPlan()), + warm_start=second.manifest) + parent = third.manifest.parents.warm_start + assert parent == embedded_model(second.manifest) + assert parent.parents.warm_start == ModelManifest(id=first.id, + reconstructed=True) + assert parent.parents == ModelParents( + training_table=second.manifest.parents.training_table, + warm_start=ModelManifest(id=first.id, reconstructed=True)) + + +def test_a_mlops_block_stores_each_model_once(tmp_path): + cfg = kumo.MLOps(artifact_root=str(tmp_path / 'artifacts')) + result = _finish_stub_fit(tmp_path / 'bundle', kumo.Trainer(ModelPlan()), + settings=cfg) + stored = tmp_path / 'artifacts' / 'models' / result.id + assert kumo.Trainer.load(stored).bundle_id == result.id + (stored / 'model' / 'config.json').write_text('{"damaged": 1}') + again = _finish_stub_fit(tmp_path / 'again', kumo.Trainer(ModelPlan()), + settings=cfg) + assert again.id == result.id + assert (stored / 'model' / 'config.json').read_text() == '{}' + off = kumo.MLOps(artifact_root=str(tmp_path / 'off'), export={'snapshot'}) + _finish_stub_fit(tmp_path / 'off-bundle', kumo.Trainer(ModelPlan()), + settings=off) + assert not (tmp_path / 'off').exists() + + +def test_manifest_failure_warns_and_still_publishes(tmp_path, monkeypatch): + from kumo.engine.trainer import _provenance + + def fail(*args, **kwargs): + raise RuntimeError('unreadable manifest') + + monkeypatch.setattr(_provenance, 'write_manifest', fail) + trainer = kumo.Trainer(ModelPlan()) + with pytest.warns(UserWarning, match='unreadable manifest'): + result = _finish_stub_fit(tmp_path / 'bundle', trainer) + assert result.manifest is None + assert re.fullmatch('[0-9a-f]{64}', result.bundle_id) + assert trainer.result is result diff --git a/packages/kumo-engine/test/test_semantic_carry_forward.py b/packages/kumo-engine/test/test_semantic_carry_forward.py index a1533b131..9481d14d1 100644 --- a/packages/kumo-engine/test/test_semantic_carry_forward.py +++ b/packages/kumo-engine/test/test_semantic_carry_forward.py @@ -874,10 +874,9 @@ def fake_train(training_input, **kwargs): both_workers.wait(timeout=5) time.sleep(.05) run_name = Path(training_input.home_dir).name - return SimpleNamespace( - metrics={'run_name': - run_name}, artifact_paths={}, explain_disabled=True, - best_trainer=_best_trainer(tmp_path)) + return SimpleNamespace(metrics={'run_name': run_name}, + artifact_paths={}, explain_disabled=True, + best_trainer=_best_trainer(tmp_path)) monkeypatch.setattr( 'kumo.engine.trainer.trainer.build_generated_training_task_dataset', diff --git a/packages/kumo-engine/test/test_trainer_registry.py b/packages/kumo-engine/test/test_trainer_registry.py index 9900d5ff0..e668544b2 100644 --- a/packages/kumo-engine/test/test_trainer_registry.py +++ b/packages/kumo-engine/test/test_trainer_registry.py @@ -210,13 +210,13 @@ def test_finish_fit_persists_spec_without_export_before_exposing_state( trainer = Trainer(ModelPlan()) result = TrainingResult({'auc': .9}, {}, bundle) - assert trainer._finish_fit(result, predictive_query='PREDICT COUNT(*)', - snapshot_id='snap-123', - training_table_generation=_generation_plan(), - export_config=None) is result + finished = trainer._finish_fit( + result, predictive_query='PREDICT COUNT(*)', snapshot_id='snap-123', + training_table_generation=_generation_plan(), export_config=None) assert (bundle / SPEC_FILE).is_file() - assert trainer._training_result is result + assert finished.bundle_id == Trainer.load(bundle).bundle_id + assert trainer._training_result is finished assert trainer._snapshot_id == 'snap-123' assert trainer._training_query == 'PREDICT COUNT(*)' assert result.published_model is None diff --git a/packages/kumo-engine/test/test_warm_start.py b/packages/kumo-engine/test/test_warm_start.py index e79130045..25a109489 100644 --- a/packages/kumo-engine/test/test_warm_start.py +++ b/packages/kumo-engine/test/test_warm_start.py @@ -182,6 +182,8 @@ def configure(self): patch.setattr(BaseLightningModule, 'configure_optimizers', configure) second = destination.fit(table, snapshot=snapshot, warm_start=source) assert restored_states and configured_states + assert second.manifest.parents.warm_start.id == source.bundle_id + assert second.manifest.parents.warm_start.parents.warm_start is None assert destination.model_plan == configured_plan assert _hash_bundle(durable) == original_hashes assert _hash_bundle(source_bundle) == original_hashes