From cca0d0028416949f9f3b25de4dce018be9bdb3ab Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 20 Oct 2026 20:49:35 +0000 Subject: [PATCH 01/15] Record bundle_id and fit provenance in every model bundle A fitted model bundle had no stable id and recorded no fit time, seed, trials, library versions or parent ids. verify_model_bundle already computed a 16-hex content digest that every caller discarded. verify_model_bundle now returns the 64-hex bundle_id: sha256 over the sorted [relative_path, sha256] pairs of the verified files (plus an unlisted training-spec.json, as before), excluding provenance.json. TrainingResult.bundle_id and Trainer.bundle_id carry it from fit and load; older bundles get an id at load without being rewritten. A new kumo-core add_bundle_file lists a written file in the publication manifest and replaces the engine's own training-spec amend. Fit writes provenance.json (role provenance, integrity-covered, outside the id): parents (training table id, snapshot id, warm-start bundle id), seed, fit times, selected trial, per-trial status/seconds/best epoch/val metrics/varying plan params, environment and scheduler. TrainingResult.provenance reads it only when listed. Provenance and id failures warn; the bundle still publishes. Verified with new tests for the id formula and exclusion, add_bundle_file, a two-trial fit's provenance, id stability across load and copy, a bundle without provenance loading and predicting identically, a stray unlisted provenance file, a failing provenance write and the warm-start parent id; plus the registry, compatibility, warm-start, holdout, snapshot, carry- forward, explanation, training-table export, Databricks and Snowflake registry suites. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-core/src/kumo_core/artifacts.py | 56 ++++- packages/kumo-core/tests/test_artifacts.py | 42 ++++ .../kumo/engine/trainer/_provenance.py | 146 +++++++++++++ .../kumo-engine/kumo/engine/trainer/_spec.py | 43 +--- .../kumo/engine/trainer/trainer.py | 53 ++++- .../test/test_bundle_provenance.py | 194 ++++++++++++++++++ .../kumo-engine/test/test_trainer_registry.py | 10 +- packages/kumo-engine/test/test_warm_start.py | 1 + 8 files changed, 493 insertions(+), 52 deletions(-) create mode 100644 packages/kumo-engine/kumo/engine/trainer/_provenance.py create mode 100644 packages/kumo-engine/test/test_bundle_provenance.py diff --git a/packages/kumo-core/src/kumo_core/artifacts.py b/packages/kumo-core/src/kumo_core/artifacts.py index 838f4c0c4..7ea69b635 100644 --- a/packages/kumo-core/src/kumo_core/artifacts.py +++ b/packages/kumo-core/src/kumo_core/artifacts.py @@ -248,14 +248,17 @@ def _sha256(path: Path) -> str: return digest.hexdigest() +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: the SHA-256 of the sorted ``[relative_path, + sha256]`` pairs of 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; its digest still participates in the id. """ from .sources import layout from .sources.errors import PermanentSourceError @@ -273,7 +276,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 +328,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 +338,42 @@ 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] + payload = json.dumps( + sorted(pair for pair in verified if pair[0] != 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() + ] + temporary = manifest_path.with_name(f'.{manifest_path.name}.tmp') + try: + with temporary.open('w', encoding='utf-8') as handle: + handle.write(json.dumps(manifest, indent=2, sort_keys=True) + '\n') + handle.flush() + os.fsync(handle.fileno()) + os.replace(temporary, manifest_path) + FilesystemArtifactStore._fsync_directory(bundle_dir) + finally: + temporary.unlink(missing_ok=True) @runtime_checkable diff --git a/packages/kumo-core/tests/test_artifacts.py b/packages/kumo-core/tests/test_artifacts.py index 071f4339b..e29622b96 100644 --- a/packages/kumo-core/tests/test_artifacts.py +++ b/packages/kumo-core/tests/test_artifacts.py @@ -15,7 +15,9 @@ FilesystemVitalArtifactPublisher, VitalArtifactPublicationPolicy, VitalArtifactSources, + add_bundle_file, resolve_vital_artifact_sources, + verify_model_bundle, ) from kumo_core.config import NativeRuntimeConfig from kumo_core.manifest import RunManifest, RunPaths @@ -237,3 +239,43 @@ 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_covers_listed_files_except_provenance( + tmp_path: Path) -> None: + sources, _ = _vital_sources(tmp_path) + bundle = tmp_path / "bundle" + VitalArtifactPublicationPolicy().publish(sources, bundle) + (bundle / "training-spec.json").write_text("{}") + entries = json.loads( + (bundle / "publication-manifest.json").read_text())["artifacts"] + pairs = sorted([entry["relative_path"], entry["sha256"]] + for entry in entries) + pairs = sorted(pairs + [[ + "training-spec.json", + hashlib.sha256(b"{}").hexdigest(), + ]]) + expected = hashlib.sha256( + json.dumps(pairs, separators=(",", ":")).encode()).hexdigest() + + assert verify_model_bundle(bundle) == expected + add_bundle_file(bundle, "training_spec", "training-spec.json") + assert verify_model_bundle(bundle) == expected + + (bundle / "provenance.json").write_text('{"a": 1}') + add_bundle_file(bundle, "provenance", "provenance.json") + assert verify_model_bundle(bundle) == expected + (bundle / "provenance.json").write_text('{"a": 2}') + with pytest.raises(ValueError, match="provenance.json"): + verify_model_bundle(bundle) + add_bundle_file(bundle, "provenance", "provenance.json") + assert verify_model_bundle(bundle) == expected + + 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) != expected 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..31740c462 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -0,0 +1,146 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""Descriptive fit record stored as ``provenance.json`` in a model bundle. + +It is integrity-covered by the publication manifest but excluded from the +bundle id, so its volatile fields (times, environment, scheduler) never change +which model a bundle is. +""" + +import csv +import json +import math +import re +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Dict, List, Optional + +from kumo.engine._manifest import ( + _flatten, + new_manifest, + read_manifest, + write_manifest, +) +from kumo.engine.pquery.training_table import TrainingTable +from kumo.engine.trainer._spec import ( + MANIFEST_FILE, + PROVENANCE_ROLE, + SPEC_VERSION, +) + + +def write_provenance( + bundle_dir: Path, + *, + train_table: TrainingTable, + snapshot_id: str, + warm_start_bundle_id: Optional[str], + started_at: datetime, + home_dir: Path, +) -> None: + from kumo_core.artifacts import PROVENANCE_FILE, add_bundle_file + + config = _read_json(bundle_dir / 'model' / 'config.json') or {} + finished_at = datetime.now(timezone.utc) + parents = { + 'training_table': train_table.tt_id, + 'training_table_content_sha256': + (train_table.manifest or {}).get('content_sha256'), + 'snapshot': snapshot_id, + 'warm_start': warm_start_bundle_id, + } + write_manifest( + bundle_dir / PROVENANCE_FILE, + new_manifest( + 'model_bundle', None, parents, + bundle_format_version=SPEC_VERSION, + seed=config.get('random_state_seed'), + fit={ + 'started_at': started_at.isoformat(timespec='seconds'), + 'finished_at': finished_at.isoformat(timespec='seconds'), + 'seconds': round((finished_at - started_at).total_seconds(), + 3), + }, + selected_trial=config.get('version'), + trials=_trials(home_dir, Path(config['root_dir'])), + )) + add_bundle_file(bundle_dir, PROVENANCE_ROLE, PROVENANCE_FILE) + + +def read_provenance(bundle_dir: Path) -> Optional[Dict[str, Any]]: + manifest = _read_json(bundle_dir / MANIFEST_FILE) or {} + for entry in manifest.get('artifacts') or []: + if entry.get('role') == PROVENANCE_ROLE: + return read_manifest(bundle_dir / entry['relative_path']) + return None + + +def _trials(home_dir: Path, automl_dir: Path) -> List[Dict[str, Any]]: + status = { + entry['trial']: entry + for entry in (_read_json(home_dir / 'status.json') or {}).get( + 'trials') or [] + } + versions = { + int(match.group(1)) + for path in automl_dir.iterdir() + if (match := re.fullmatch(r'version_(\d+)', path.name)) + } + trials, plans = [], {} + for trial in sorted(status.keys() | versions): + log_dir = automl_dir / f'version_{trial}' + config = _read_json(log_dir / 'config.json') or {} + progress = _read_json(log_dir / 'train_progress.json') or {} + best_epoch = (config.get('refit_cfg') or {}).get('best_val_epoch', -1) + succeeded = status.get(trial, {}).get('succeed') + trials.append({ + 'trial': trial, + 'status': (None if succeeded is None else + 'succeeded' if succeeded else 'failed'), + 'seconds': progress.get('elapsed_time'), + 'best_epoch': best_epoch if best_epoch >= 0 else None, + 'metrics': _val_metrics(log_dir / 'val_result.csv'), + }) + plans[trial] = { + key: json.dumps(value, sort_keys=True) + for key, value in _flatten( + _read_json(log_dir / 'model_plan.json') or {}).items() + } + varying = sorted( + key for key in {key + for plan in plans.values() for key in plan} + if len({plan.get(key) + for plan in plans.values()}) > 1) + for entry in trials: + plan = plans[entry['trial']] + entry['params'] = { + key: json.loads(plan[key]) if key in plan else None + for key in varying + } + return trials + + +def _val_metrics(path: Path) -> Dict[str, float]: + try: + with open(path, newline='', encoding='utf-8') as handle: + row = next(csv.DictReader(handle), None) or {} + except OSError: + return {} + metrics = {} + for key, value in row.items(): + try: + number = float(value) + except (TypeError, ValueError): + continue + if key.startswith('val_') and math.isfinite(number): + metrics[key] = number + return metrics + + +def _read_json(path: Path) -> Optional[Dict[str, Any]]: + try: + data = json.loads(path.read_text(encoding='utf-8')) + except (OSError, ValueError): + return None + return data if isinstance(data, dict) else None diff --git a/packages/kumo-engine/kumo/engine/trainer/_spec.py b/packages/kumo-engine/kumo/engine/trainer/_spec.py index 147817c0e..637714d45 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, 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,12 @@ 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 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 +129,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..bc0824125 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -83,6 +83,9 @@ class TrainingResult: model bundle. The file is capped at 1 GiB in memory; `holdout_truncated` is true when rows were dropped to fit, and `None` when the size record could not be read. + + `bundle_id` is the content id of the model bundle's scoring files, + computed when the bundle is written or loaded. """ metrics: Dict[str, float] @@ -91,6 +94,14 @@ class TrainingResult: published_model: Optional[PublishedModel] = None holdout_path: Optional[Path] = None holdout_truncated: Optional[bool] = False + bundle_id: Optional[str] = None + + @property + def provenance(self) -> Optional[Dict[str, Any]]: + """The bundle's recorded fit provenance; 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. @@ -299,6 +310,13 @@ 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.""" + if self._training_result is None: + return None + return self._training_result.bundle_id + @_without_torch_scatter_hint def fit( self, @@ -363,18 +381,21 @@ def fit( raise NotImplementedError( f'{type(export_config.connector).__name__} cannot publish ' f'models') + 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'], @@ -524,6 +545,10 @@ def fit( training_table_generation=train_table.generation_plan, export_config=export_config, model_plan=effective_plan, + provenance=dict(train_table=train_table, + snapshot_id=snapshot.id, + warm_start_bundle_id=warm_start_bundle_id, + started_at=started_at, home_dir=run_dir), ) def _finish_fit( @@ -535,8 +560,13 @@ def _finish_fit( training_table_generation: TrainingTableGenerationPlan, export_config: Optional[ModelOutputConfig], model_plan: Optional[ModelPlan] = None, + provenance: Optional[Mapping[str, Any]] = None, ) -> TrainingResult: """Commit durable fit outputs before making this trainer usable.""" + from kumo_core.artifacts import verify_model_bundle + + 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 +575,20 @@ def _finish_fit( snapshot_id=snapshot_id, training_table_generation=training_table_generation, ) + if provenance is not None: + try: + write_provenance(training_result.bundle_dir, **provenance) + except Exception as error: + warnings.warn( + f"Publishing the model bundle without provenance.json: " + f"{error}", UserWarning, stacklevel=3) + try: + training_result = replace( + training_result, + bundle_id=verify_model_bundle(training_result.bundle_dir)) + except (OSError, ValueError) as error: + warnings.warn(f"Could not compute the bundle id: {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( @@ -837,7 +881,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 +893,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..443c90742 --- /dev/null +++ b/packages/kumo-engine/test/test_bundle_provenance.py @@ -0,0 +1,194 @@ +# 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 + +import numpy as np +import pandas as pd +import pytest +from kumo_core.artifacts import add_bundle_file +from kumoapi.graph import Edge +from kumoapi.model_plan import ModelPlan, TrainingTableGenerationPlan +from kumoapi.typing import Stype + +import kumo.engine as kumo +from kumo.engine._scratch import get_scratch_dir +from kumo.engine.trainer._spec import MANIFEST_FILE +from kumo.engine.trainer.trainer import TrainingResult + + +@pytest.fixture(scope='module') +def two_trials(tmp_path_factory): + root = tmp_path_factory.mktemp('provenance') + previous = get_scratch_dir() + kumo.configure(scratch_dir=root / 'scratch') + try: + rng = np.random.default_rng(0) + data = root / 'data' + data.mkdir() + pd.DataFrame({'ID': np.arange(50)}).to_parquet(data / 'ents.parquet') + pd.DataFrame({ + 'EID': + np.arange(3000), + 'ENT_ID': + rng.integers(0, 50, 3000), + 'TS': + pd.Timestamp('2024-01-01') + + pd.to_timedelta(rng.integers(0, 120, 3000), unit='D'), + }).to_parquet(data / 'events.parquet') + connector = kumo.FileConnector(root_dir=str(data)) + events = kumo.Table.from_source_table( + connector['events.parquet'], primary_key='EID', + time_column='TS').infer_metadata().with_column( + 'ENT_ID', stype=Stype.ID) + graph = kumo.Graph( + tables={ + 'ents': + kumo.Table.from_source_table(connector['ents.parquet'], + primary_key='ID').infer_metadata(), + 'events': + events, + }, + edges=[Edge(src_table='events', fkey='ENT_ID', dst_table='ents')]) + snapshot = graph.snapshot() + pquery = kumo.PredictiveQuery( + graph=graph, query='PREDICT COUNT(events.*, 0, 7, days) > 0 ' + 'FOR EACH ents.ID') + plan = ModelPlan() + plan.training_job.disable_explain = True + plan.training_job.num_experiments = 2 + plan.optimization.max_epochs = 1 + plan.optimization.max_steps_per_epoch = 2 + table = pquery.generate_training_table(snapshot) + trainer = kumo.Trainer(model_plan=plan) + result = trainer.fit(table, snapshot=snapshot) + yield (trainer, result, table, snapshot, + pquery.generate_prediction_table(snapshot)) + finally: + kumo.configure(scratch_dir=previous) + + +def _predict(trainer, table, snapshot): + result = trainer.predict(table, snapshot=snapshot, + binary_classification_threshold=.5) + return result.frame.sort_values('ENTITY').reset_index(drop=True) + + +def _roles(bundle): + manifest = json.loads((bundle / MANIFEST_FILE).read_text()) + return {entry['role']: entry for entry in manifest['artifacts']} + + +def _strip_provenance(bundle): + manifest = json.loads((bundle / MANIFEST_FILE).read_text()) + manifest['artifacts'] = [ + entry for entry in manifest['artifacts'] + if entry['role'] != 'provenance' + ] + (bundle / MANIFEST_FILE).write_text(json.dumps(manifest)) + (bundle / 'provenance.json').unlink() + + +def test_fit_records_listed_provenance(two_trials): + trainer, result, table, snapshot, _ = two_trials + assert _roles(result.bundle_dir)['provenance']['relative_path'] == ( + 'provenance.json') + provenance = result.provenance + assert provenance['kind'] == 'model_bundle' + assert provenance['bundle_format_version'] == 1 + assert provenance['parents'] == { + 'training_table': table.tt_id, + 'training_table_content_sha256': None, + 'snapshot': snapshot.id, + 'warm_start': None, + } + assert provenance['seed'] is None + assert provenance['fit']['seconds'] >= 0 + assert [trial['trial'] for trial in provenance['trials']] == [0, 1] + assert {trial['status'] for trial in provenance['trials']} == {'succeeded'} + assert provenance['selected_trial'] in (0, 1) + for trial in provenance['trials']: + assert trial['best_epoch'] == 1 + assert trial['seconds'] > 0 + assert 'val_auroc' in trial['metrics'] + params = [trial['params'] for trial in provenance['trials']] + assert params[0] and params[0].keys() == params[1].keys() + assert all(params[0][key] != params[1][key] for key in params[0]) + assert provenance['environment']['kumo_engine'] == kumo.__version__ + assert trainer.bundle_id == result.bundle_id + + +def test_bundle_id_is_content_identity(two_trials, tmp_path): + trainer, result, *_ = two_trials + assert re.fullmatch('[0-9a-f]{64}', result.bundle_id) + assert kumo.Trainer.load(result.bundle_dir).bundle_id == result.bundle_id + copy = tmp_path / 'copy' + shutil.copytree(result.bundle_dir, copy) + assert kumo.Trainer.load(copy).bundle_id == result.bundle_id + + provenance = json.loads((copy / 'provenance.json').read_text()) + provenance['fit']['seconds'] = -1 + (copy / 'provenance.json').write_text(json.dumps(provenance)) + add_bundle_file(copy, 'provenance', 'provenance.json') + reloaded = kumo.Trainer.load(copy) + assert reloaded.bundle_id == result.bundle_id + assert reloaded._training_result.provenance['fit']['seconds'] == -1 + + config = copy / 'model' / 'config.json' + config.write_text(config.read_text() + ' ') + add_bundle_file(copy, 'trainer_config', 'model/config.json') + assert kumo.Trainer.load(copy).bundle_id != result.bundle_id + + +def test_bundle_without_provenance_loads_and_predicts_identically( + two_trials, tmp_path): + trainer, result, _, snapshot, prediction_table = two_trials + old = tmp_path / 'old' + shutil.copytree(result.bundle_dir, old) + _strip_provenance(old) + + restored = kumo.Trainer.load(old) + assert restored.bundle_id == result.bundle_id + assert restored._training_result.provenance is None + pd.testing.assert_frame_equal(_predict(restored, prediction_table, + snapshot), + _predict(trainer, prediction_table, snapshot)) + + (old / 'provenance.json').write_text(json.dumps(result.provenance)) + stray = kumo.Trainer.load(old) + assert stray._training_result.provenance is None + assert stray.bundle_id == result.bundle_id + + +def test_provenance_failure_warns_and_still_publishes(tmp_path, monkeypatch): + from kumo.engine.trainer import _provenance + + bundle = tmp_path / 'bundle' + (bundle / 'model').mkdir(parents=True) + (bundle / 'model' / 'config.json').write_text('{"root_dir": "automl"}') + (bundle / MANIFEST_FILE).write_text(json.dumps({'artifacts': []})) + add_bundle_file(bundle, 'trainer_config', 'model/config.json') + + def fail(*args): + raise RuntimeError('unreadable trials') + + monkeypatch.setattr(_provenance, '_trials', fail) + table = type('Table', (), {'tt_id': 'tt-1', 'manifest': None})() + trainer = kumo.Trainer(ModelPlan()) + with pytest.warns(UserWarning, match='unreadable trials'): + result = trainer._finish_fit( + TrainingResult({}, {}, bundle), predictive_query='PREDICT x', + snapshot_id='snap', export_config=None, + training_table_generation=TrainingTableGenerationPlan(), + provenance=dict(train_table=table, snapshot_id='snap', + warm_start_bundle_id=None, + started_at=datetime.now(timezone.utc), + home_dir=tmp_path)) + assert 'provenance' not in _roles(bundle) + assert result.provenance is None + assert re.fullmatch('[0-9a-f]{64}', result.bundle_id) + assert trainer._training_result is result 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..d60a645cb 100644 --- a/packages/kumo-engine/test/test_warm_start.py +++ b/packages/kumo-engine/test/test_warm_start.py @@ -182,6 +182,7 @@ 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.provenance['parents']['warm_start'] == source.bundle_id assert destination.model_plan == configured_plan assert _hash_bundle(durable) == original_hashes assert _hash_bundle(source_bundle) == original_hashes From 1b676995074f2b6feaf7864edfc7a492510e38cc Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 20 Oct 2026 21:03:27 +0000 Subject: [PATCH 02/15] Address bundle provenance review feedback Any bundle-id failure now warns instead of only OSError/ValueError. Provenance drops the always-null training-table content digest and records a training_table object (id, derived_from, modified) copied from the training-table manifest. Missing trial logs yield an empty trials list so parents, seed, times and environment are still written. Single-element trial params are unwrapped. new_manifest accepts a null id. Verified with the bundle provenance (incl. a new missing-trial-logs test), trainer registry, warm start and kumo-core artifact suites. Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/trainer/_provenance.py | 26 ++++++++--- .../kumo/engine/trainer/trainer.py | 2 +- .../test/test_bundle_provenance.py | 45 ++++++++++++------- 3 files changed, 51 insertions(+), 22 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/_provenance.py b/packages/kumo-engine/kumo/engine/trainer/_provenance.py index 31740c462..6c6ae586e 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_provenance.py +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -43,10 +43,10 @@ def write_provenance( config = _read_json(bundle_dir / 'model' / 'config.json') or {} finished_at = datetime.now(timezone.utc) + tt_id = train_table.tt_id + table_manifest = train_table.manifest or {} parents = { - 'training_table': train_table.tt_id, - 'training_table_content_sha256': - (train_table.manifest or {}).get('content_sha256'), + 'training_table': tt_id, 'snapshot': snapshot_id, 'warm_start': warm_start_bundle_id, } @@ -62,8 +62,14 @@ def write_provenance( 'seconds': round((finished_at - started_at).total_seconds(), 3), }, + training_table={ + 'id': tt_id, + 'derived_from': + (table_manifest.get('parents') or {}).get('derived_from'), + 'modified': table_manifest.get('modified'), + }, selected_trial=config.get('version'), - trials=_trials(home_dir, Path(config['root_dir'])), + trials=_trials(home_dir, config.get('root_dir')), )) add_bundle_file(bundle_dir, PROVENANCE_ROLE, PROVENANCE_FILE) @@ -76,7 +82,11 @@ def read_provenance(bundle_dir: Path) -> Optional[Dict[str, Any]]: return None -def _trials(home_dir: Path, automl_dir: Path) -> List[Dict[str, Any]]: +def _trials(home_dir: Path, + root_dir: Optional[str]) -> List[Dict[str, Any]]: + automl_dir = Path(root_dir) if root_dir else None + if automl_dir is None or not automl_dir.is_dir(): + return [] status = { entry['trial']: entry for entry in (_read_json(home_dir / 'status.json') or {}).get( @@ -115,12 +125,16 @@ def _trials(home_dir: Path, automl_dir: Path) -> List[Dict[str, Any]]: for entry in trials: plan = plans[entry['trial']] entry['params'] = { - key: json.loads(plan[key]) if key in plan else None + key: _unwrap(json.loads(plan[key])) if key in plan else None for key in varying } return trials +def _unwrap(value: Any) -> Any: + return value[0] if isinstance(value, list) and len(value) == 1 else value + + def _val_metrics(path: Path) -> Dict[str, float]: try: with open(path, newline='', encoding='utf-8') as handle: diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index bc0824125..b7ec406c1 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -586,7 +586,7 @@ def _finish_fit( training_result = replace( training_result, bundle_id=verify_model_bundle(training_result.bundle_dir)) - except (OSError, ValueError) as error: + except Exception as error: warnings.warn(f"Could not compute the bundle id: {error}", UserWarning, stacklevel=3) if export_config is not None: diff --git a/packages/kumo-engine/test/test_bundle_provenance.py b/packages/kumo-engine/test/test_bundle_provenance.py index 443c90742..ba3e383de 100644 --- a/packages/kumo-engine/test/test_bundle_provenance.py +++ b/packages/kumo-engine/test/test_bundle_provenance.py @@ -102,10 +102,14 @@ def test_fit_records_listed_provenance(two_trials): assert provenance['bundle_format_version'] == 1 assert provenance['parents'] == { 'training_table': table.tt_id, - 'training_table_content_sha256': None, 'snapshot': snapshot.id, 'warm_start': None, } + assert provenance['training_table'] == { + 'id': table.tt_id, + 'derived_from': None, + 'modified': None, + } assert provenance['seed'] is None assert provenance['fit']['seconds'] >= 0 assert [trial['trial'] for trial in provenance['trials']] == [0, 1] @@ -118,6 +122,8 @@ def test_fit_records_listed_provenance(two_trials): params = [trial['params'] for trial in provenance['trials']] assert params[0] and params[0].keys() == params[1].keys() assert all(params[0][key] != params[1][key] for key in params[0]) + assert not any(isinstance(value, list) and len(value) == 1 + for value in params[0].values()) assert provenance['environment']['kumo_engine'] == kumo.__version__ assert trainer.bundle_id == result.bundle_id @@ -164,31 +170,40 @@ def test_bundle_without_provenance_loads_and_predicts_identically( assert stray.bundle_id == result.bundle_id -def test_provenance_failure_warns_and_still_publishes(tmp_path, monkeypatch): - from kumo.engine.trainer import _provenance - +def _finish_stub_fit(tmp_path, trainer): bundle = tmp_path / 'bundle' (bundle / 'model').mkdir(parents=True) - (bundle / 'model' / 'config.json').write_text('{"root_dir": "automl"}') + (bundle / 'model' / 'config.json').write_text('{}') (bundle / MANIFEST_FILE).write_text(json.dumps({'artifacts': []})) add_bundle_file(bundle, 'trainer_config', 'model/config.json') + table = type('Table', (), {'tt_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, snapshot_id='snap', + warm_start_bundle_id=None, + started_at=datetime.now(timezone.utc), + home_dir=tmp_path)) + + +def test_provenance_without_trial_logs_still_records_parents(tmp_path): + result = _finish_stub_fit(tmp_path, kumo.Trainer(ModelPlan())) + assert result.provenance['trials'] == [] + assert result.provenance['parents']['training_table'] == 'tt-1' + + +def test_provenance_failure_warns_and_still_publishes(tmp_path, monkeypatch): + from kumo.engine.trainer import _provenance def fail(*args): raise RuntimeError('unreadable trials') monkeypatch.setattr(_provenance, '_trials', fail) - table = type('Table', (), {'tt_id': 'tt-1', 'manifest': None})() trainer = kumo.Trainer(ModelPlan()) with pytest.warns(UserWarning, match='unreadable trials'): - result = trainer._finish_fit( - TrainingResult({}, {}, bundle), predictive_query='PREDICT x', - snapshot_id='snap', export_config=None, - training_table_generation=TrainingTableGenerationPlan(), - provenance=dict(train_table=table, snapshot_id='snap', - warm_start_bundle_id=None, - started_at=datetime.now(timezone.utc), - home_dir=tmp_path)) - assert 'provenance' not in _roles(bundle) + result = _finish_stub_fit(tmp_path, trainer) + assert 'provenance' not in _roles(result.bundle_dir) assert result.provenance is None assert re.fullmatch('[0-9a-f]{64}', result.bundle_id) assert trainer._training_result is result From d15d97e367407e73d1973219fafa53751ef2df18 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 20 Oct 2026 22:35:09 +0000 Subject: [PATCH 03/15] Read Trainer.bundle_id defensively Training results built without a bundle_id field (test doubles) made the property raise; return None for them. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/trainer/trainer.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index b7ec406c1..28596f615 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -313,9 +313,7 @@ def scratch_dir(self) -> Path: @property def bundle_id(self) -> Optional[str]: """Content id of the fitted or loaded model bundle.""" - if self._training_result is None: - return None - return self._training_result.bundle_id + return getattr(self._training_result, 'bundle_id', None) @_without_torch_scatter_hint def fit( From ffbf2065700eab97046efd42cc6071fb6c43db4c Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 20 Oct 2026 23:56:18 +0000 Subject: [PATCH 04/15] Record the published model's fit mode in provenance.json metrics.json holds test_* (the fit model) and test_refit_* (the refit_trainval model) and nothing says which, if any, belongs to the shipped checkpoint; a refit_full model has none of its own. Record published_model ('fit', 'refit_trainval' or 'refit_full') from the best trainer's refit flags, the order kumo-core publishes in; null when unknown. metrics.json is unchanged. Verified with a two-trial fit and each combination of refit flags. Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/trainer/_provenance.py | 2 ++ .../kumo/engine/trainer/trainer.py | 20 +++++++++++++++++-- .../test/test_bundle_provenance.py | 20 +++++++++++++++++-- 3 files changed, 38 insertions(+), 4 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/_provenance.py b/packages/kumo-engine/kumo/engine/trainer/_provenance.py index 6c6ae586e..222fc931c 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_provenance.py +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -38,6 +38,7 @@ def write_provenance( warm_start_bundle_id: Optional[str], started_at: datetime, home_dir: Path, + published_model: Optional[str], ) -> None: from kumo_core.artifacts import PROVENANCE_FILE, add_bundle_file @@ -69,6 +70,7 @@ def write_provenance( 'modified': table_manifest.get('modified'), }, selected_trial=config.get('version'), + published_model=published_model, trials=_trials(home_dir, config.get('root_dir')), )) add_bundle_file(bundle_dir, PROVENANCE_ROLE, PROVENANCE_FILE) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 28596f615..4794169ca 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -98,7 +98,11 @@ class TrainingResult: @property def provenance(self) -> Optional[Dict[str, Any]]: - """The bundle's recorded fit provenance; None for older bundles.""" + """The bundle's recorded fit provenance; None for older bundles. + + Its `published_model` is the shipped checkpoint's fit mode, unlike + the registry record in `TrainingResult.published_model`. + """ from kumo.engine.trainer._provenance import read_provenance return read_provenance(self.bundle_dir) @@ -210,6 +214,16 @@ def _published_holdout(trainer: Any) -> tuple[Optional[Path], Optional[bool]]: return _find_holdout(trainer.get_log_dir(fit_mode=published)) +def _published_fit_mode(trainer: Any) -> Optional[str]: + refit = getattr(trainer, 'refit_cfg', None) + for mode, flag in (('refit_full', 'refitted_full'), + ('refit_trainval', 'refitted_trainval'), + ('fit', 'fitted')): + if getattr(refit, flag, False) is True: + return mode + return None + + def _collect_garbage(stage: str) -> None: # Lightning links each trainer to its data module and model in reference # cycles; free a finished run now rather than at a later collection. @@ -527,6 +541,7 @@ def fit( holdout_path, holdout_truncated = _published_holdout( result.best_trainer) + published_model = _published_fit_mode(result.best_trainer) training_result = TrainingResult( metrics=dict(result.metrics), artifact_paths=dict(result.artifact_paths), @@ -546,7 +561,8 @@ def fit( provenance=dict(train_table=train_table, snapshot_id=snapshot.id, warm_start_bundle_id=warm_start_bundle_id, - started_at=started_at, home_dir=run_dir), + started_at=started_at, home_dir=run_dir, + published_model=published_model), ) def _finish_fit( diff --git a/packages/kumo-engine/test/test_bundle_provenance.py b/packages/kumo-engine/test/test_bundle_provenance.py index ba3e383de..9c4363303 100644 --- a/packages/kumo-engine/test/test_bundle_provenance.py +++ b/packages/kumo-engine/test/test_bundle_provenance.py @@ -6,6 +6,7 @@ import re import shutil from datetime import datetime, timezone +from types import SimpleNamespace import numpy as np import pandas as pd @@ -18,7 +19,7 @@ import kumo.engine as kumo from kumo.engine._scratch import get_scratch_dir from kumo.engine.trainer._spec import MANIFEST_FILE -from kumo.engine.trainer.trainer import TrainingResult +from kumo.engine.trainer.trainer import TrainingResult, _published_fit_mode @pytest.fixture(scope='module') @@ -115,6 +116,7 @@ def test_fit_records_listed_provenance(two_trials): assert [trial['trial'] for trial in provenance['trials']] == [0, 1] assert {trial['status'] for trial in provenance['trials']} == {'succeeded'} assert provenance['selected_trial'] in (0, 1) + assert provenance['published_model'] == 'refit_trainval' for trial in provenance['trials']: assert trial['best_epoch'] == 1 assert trial['seconds'] > 0 @@ -184,13 +186,27 @@ def _finish_stub_fit(tmp_path, trainer): provenance=dict(train_table=table, snapshot_id='snap', warm_start_bundle_id=None, started_at=datetime.now(timezone.utc), - home_dir=tmp_path)) + home_dir=tmp_path, published_model=None)) def test_provenance_without_trial_logs_still_records_parents(tmp_path): result = _finish_stub_fit(tmp_path, kumo.Trainer(ModelPlan())) assert result.provenance['trials'] == [] assert result.provenance['parents']['training_table'] == 'tt-1' + assert result.provenance['published_model'] is None + + +@pytest.mark.parametrize('flags,mode', [ + ({}, None), + ({'fitted': True}, 'fit'), + ({'fitted': True, 'refitted_trainval': True}, 'refit_trainval'), + ({'fitted': True, 'refitted_full': True}, 'refit_full'), + ({'fitted': True, 'refitted_trainval': True, 'refitted_full': True}, + 'refit_full'), +]) +def test_published_fit_mode_follows_refit_flags(flags, mode): + trainer = SimpleNamespace(refit_cfg=SimpleNamespace(**flags)) + assert _published_fit_mode(trainer) == mode def test_provenance_failure_warns_and_still_publishes(tmp_path, monkeypatch): From 2a6fea337ad914345344c037fc7ac0a95afca840 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 06:58:17 +0000 Subject: [PATCH 05/15] Mirror the training table's source in bundle provenance #55 no longer records derived_from or modified for a reloaded table, so provenance.json's training_table block stopped having anything to copy. The block now records the table id, whether it was generated (pql) or loaded (custom), the snapshot it was generated on or validated against, and, for custom tables, the location and data version they were loaded from. The bundle stays self-contained without the scratch staging directory; query and plans stay only in training-spec.json. Verified with test_bundle_provenance.py (pql fit, custom-table stub, manifest-less stub) and test_trainer_registry.py. Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/trainer/_provenance.py | 24 ++++++++++---- .../test/test_bundle_provenance.py | 32 ++++++++++++++++--- 2 files changed, 45 insertions(+), 11 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/_provenance.py b/packages/kumo-engine/kumo/engine/trainer/_provenance.py index 222fc931c..34bc53d9e 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_provenance.py +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -45,7 +45,6 @@ def write_provenance( config = _read_json(bundle_dir / 'model' / 'config.json') or {} finished_at = datetime.now(timezone.utc) tt_id = train_table.tt_id - table_manifest = train_table.manifest or {} parents = { 'training_table': tt_id, 'snapshot': snapshot_id, @@ -63,12 +62,7 @@ def write_provenance( 'seconds': round((finished_at - started_at).total_seconds(), 3), }, - training_table={ - 'id': tt_id, - 'derived_from': - (table_manifest.get('parents') or {}).get('derived_from'), - 'modified': table_manifest.get('modified'), - }, + training_table=_training_table(tt_id, train_table.manifest), selected_trial=config.get('version'), published_model=published_model, trials=_trials(home_dir, config.get('root_dir')), @@ -76,6 +70,22 @@ def write_provenance( add_bundle_file(bundle_dir, PROVENANCE_ROLE, PROVENANCE_FILE) +def _training_table(tt_id: str, + manifest: Optional[Dict[str, Any]]) -> Dict[str, Any]: + manifest = manifest or {} + parents = manifest.get('parents') or {} + record = { + 'id': tt_id, + 'source': manifest.get('source'), + 'snapshot': parents.get('snapshot'), + } + if record['source'] == 'custom': + source_table = parents.get('source_table') or {} + record.update(location=source_table.get('location'), + version=source_table.get('version')) + return record + + def read_provenance(bundle_dir: Path) -> Optional[Dict[str, Any]]: manifest = _read_json(bundle_dir / MANIFEST_FILE) or {} for entry in manifest.get('artifacts') or []: diff --git a/packages/kumo-engine/test/test_bundle_provenance.py b/packages/kumo-engine/test/test_bundle_provenance.py index 9c4363303..b90810eed 100644 --- a/packages/kumo-engine/test/test_bundle_provenance.py +++ b/packages/kumo-engine/test/test_bundle_provenance.py @@ -108,8 +108,8 @@ def test_fit_records_listed_provenance(two_trials): } assert provenance['training_table'] == { 'id': table.tt_id, - 'derived_from': None, - 'modified': None, + 'source': 'pql', + 'snapshot': snapshot.id, } assert provenance['seed'] is None assert provenance['fit']['seconds'] >= 0 @@ -172,13 +172,13 @@ def test_bundle_without_provenance_loads_and_predicts_identically( assert stray.bundle_id == result.bundle_id -def _finish_stub_fit(tmp_path, trainer): +def _finish_stub_fit(tmp_path, trainer, manifest=None): bundle = tmp_path / 'bundle' (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 = type('Table', (), {'tt_id': 'tt-1', 'manifest': None})() + table = type('Table', (), {'tt_id': 'tt-1', 'manifest': manifest})() return trainer._finish_fit( TrainingResult({}, {}, bundle), predictive_query='PREDICT x', snapshot_id='snap', export_config=None, @@ -193,9 +193,33 @@ def test_provenance_without_trial_logs_still_records_parents(tmp_path): result = _finish_stub_fit(tmp_path, kumo.Trainer(ModelPlan())) assert result.provenance['trials'] == [] assert result.provenance['parents']['training_table'] == 'tt-1' + assert result.provenance['training_table'] == { + 'id': 'tt-1', + 'source': None, + 'snapshot': None + } assert result.provenance['published_model'] is None +def test_provenance_records_custom_table_source(tmp_path): + source_table = {'location': 's3://bucket/exports/tt', 'version': '17'} + manifest = { + 'source': 'custom', + 'query': 'PREDICT x', + 'parents': { + 'snapshot': 'snap', + 'source_table': source_table + }, + } + result = _finish_stub_fit(tmp_path, kumo.Trainer(ModelPlan()), manifest) + assert result.provenance['training_table'] == { + 'id': 'tt-1', + 'source': 'custom', + 'snapshot': 'snap', + **source_table, + } + + @pytest.mark.parametrize('flags,mode', [ ({}, None), ({'fitted': True}, 'fit'), From 68e72b91d3a5f09356feb35b34a630e4fc3dddcd Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 07:47:45 +0000 Subject: [PATCH 06/15] Trim provenance.json to facts no other bundle file holds seed copied model/config.json, bundle_format_version restated training-spec.json's spec version, and trials scraped AutoML-internal status, metric and plan files that experiment tracking already records. All three are dropped; selected_trial stays, read from the bundle's own model/config.json, and write_provenance no longer needs the run's home directory. The training_table block of a custom table now copies connector, location and version from the table manifest's parents.source_table, the same keys the snapshot manifest uses for its sources. Verified with test_bundle_provenance.py and test_trainer_registry.py. Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/trainer/_provenance.py | 98 ++----------------- .../kumo/engine/trainer/trainer.py | 2 +- .../test/test_bundle_provenance.py | 31 +++--- 3 files changed, 19 insertions(+), 112 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/_provenance.py b/packages/kumo-engine/kumo/engine/trainer/_provenance.py index 34bc53d9e..4fbce3479 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_provenance.py +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -8,26 +8,14 @@ which model a bundle is. """ -import csv import json -import math -import re from datetime import datetime, timezone from pathlib import Path -from typing import Any, Dict, List, Optional +from typing import Any, Dict, Optional -from kumo.engine._manifest import ( - _flatten, - new_manifest, - read_manifest, - write_manifest, -) +from kumo.engine._manifest import new_manifest, read_manifest, write_manifest from kumo.engine.pquery.training_table import TrainingTable -from kumo.engine.trainer._spec import ( - MANIFEST_FILE, - PROVENANCE_ROLE, - SPEC_VERSION, -) +from kumo.engine.trainer._spec import MANIFEST_FILE, PROVENANCE_ROLE def write_provenance( @@ -37,7 +25,6 @@ def write_provenance( snapshot_id: str, warm_start_bundle_id: Optional[str], started_at: datetime, - home_dir: Path, published_model: Optional[str], ) -> None: from kumo_core.artifacts import PROVENANCE_FILE, add_bundle_file @@ -54,8 +41,6 @@ def write_provenance( bundle_dir / PROVENANCE_FILE, new_manifest( 'model_bundle', None, parents, - bundle_format_version=SPEC_VERSION, - seed=config.get('random_state_seed'), fit={ 'started_at': started_at.isoformat(timespec='seconds'), 'finished_at': finished_at.isoformat(timespec='seconds'), @@ -65,7 +50,6 @@ def write_provenance( training_table=_training_table(tt_id, train_table.manifest), selected_trial=config.get('version'), published_model=published_model, - trials=_trials(home_dir, config.get('root_dir')), )) add_bundle_file(bundle_dir, PROVENANCE_ROLE, PROVENANCE_FILE) @@ -81,8 +65,10 @@ def _training_table(tt_id: str, } if record['source'] == 'custom': source_table = parents.get('source_table') or {} - record.update(location=source_table.get('location'), - version=source_table.get('version')) + record.update({ + key: source_table.get(key) + for key in ('connector', 'location', 'version') + }) return record @@ -94,76 +80,6 @@ def read_provenance(bundle_dir: Path) -> Optional[Dict[str, Any]]: return None -def _trials(home_dir: Path, - root_dir: Optional[str]) -> List[Dict[str, Any]]: - automl_dir = Path(root_dir) if root_dir else None - if automl_dir is None or not automl_dir.is_dir(): - return [] - status = { - entry['trial']: entry - for entry in (_read_json(home_dir / 'status.json') or {}).get( - 'trials') or [] - } - versions = { - int(match.group(1)) - for path in automl_dir.iterdir() - if (match := re.fullmatch(r'version_(\d+)', path.name)) - } - trials, plans = [], {} - for trial in sorted(status.keys() | versions): - log_dir = automl_dir / f'version_{trial}' - config = _read_json(log_dir / 'config.json') or {} - progress = _read_json(log_dir / 'train_progress.json') or {} - best_epoch = (config.get('refit_cfg') or {}).get('best_val_epoch', -1) - succeeded = status.get(trial, {}).get('succeed') - trials.append({ - 'trial': trial, - 'status': (None if succeeded is None else - 'succeeded' if succeeded else 'failed'), - 'seconds': progress.get('elapsed_time'), - 'best_epoch': best_epoch if best_epoch >= 0 else None, - 'metrics': _val_metrics(log_dir / 'val_result.csv'), - }) - plans[trial] = { - key: json.dumps(value, sort_keys=True) - for key, value in _flatten( - _read_json(log_dir / 'model_plan.json') or {}).items() - } - varying = sorted( - key for key in {key - for plan in plans.values() for key in plan} - if len({plan.get(key) - for plan in plans.values()}) > 1) - for entry in trials: - plan = plans[entry['trial']] - entry['params'] = { - key: _unwrap(json.loads(plan[key])) if key in plan else None - for key in varying - } - return trials - - -def _unwrap(value: Any) -> Any: - return value[0] if isinstance(value, list) and len(value) == 1 else value - - -def _val_metrics(path: Path) -> Dict[str, float]: - try: - with open(path, newline='', encoding='utf-8') as handle: - row = next(csv.DictReader(handle), None) or {} - except OSError: - return {} - metrics = {} - for key, value in row.items(): - try: - number = float(value) - except (TypeError, ValueError): - continue - if key.startswith('val_') and math.isfinite(number): - metrics[key] = number - return metrics - - def _read_json(path: Path) -> Optional[Dict[str, Any]]: try: data = json.loads(path.read_text(encoding='utf-8')) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 4794169ca..07c7bbf86 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -561,7 +561,7 @@ def fit( provenance=dict(train_table=train_table, snapshot_id=snapshot.id, warm_start_bundle_id=warm_start_bundle_id, - started_at=started_at, home_dir=run_dir, + started_at=started_at, published_model=published_model), ) diff --git a/packages/kumo-engine/test/test_bundle_provenance.py b/packages/kumo-engine/test/test_bundle_provenance.py index b90810eed..999dc552e 100644 --- a/packages/kumo-engine/test/test_bundle_provenance.py +++ b/packages/kumo-engine/test/test_bundle_provenance.py @@ -100,7 +100,6 @@ def test_fit_records_listed_provenance(two_trials): 'provenance.json') provenance = result.provenance assert provenance['kind'] == 'model_bundle' - assert provenance['bundle_format_version'] == 1 assert provenance['parents'] == { 'training_table': table.tt_id, 'snapshot': snapshot.id, @@ -111,21 +110,10 @@ def test_fit_records_listed_provenance(two_trials): 'source': 'pql', 'snapshot': snapshot.id, } - assert provenance['seed'] is None + assert not {'seed', 'bundle_format_version', 'trials'} & provenance.keys() assert provenance['fit']['seconds'] >= 0 - assert [trial['trial'] for trial in provenance['trials']] == [0, 1] - assert {trial['status'] for trial in provenance['trials']} == {'succeeded'} assert provenance['selected_trial'] in (0, 1) assert provenance['published_model'] == 'refit_trainval' - for trial in provenance['trials']: - assert trial['best_epoch'] == 1 - assert trial['seconds'] > 0 - assert 'val_auroc' in trial['metrics'] - params = [trial['params'] for trial in provenance['trials']] - assert params[0] and params[0].keys() == params[1].keys() - assert all(params[0][key] != params[1][key] for key in params[0]) - assert not any(isinstance(value, list) and len(value) == 1 - for value in params[0].values()) assert provenance['environment']['kumo_engine'] == kumo.__version__ assert trainer.bundle_id == result.bundle_id @@ -186,12 +174,11 @@ def _finish_stub_fit(tmp_path, trainer, manifest=None): provenance=dict(train_table=table, snapshot_id='snap', warm_start_bundle_id=None, started_at=datetime.now(timezone.utc), - home_dir=tmp_path, published_model=None)) + published_model=None)) -def test_provenance_without_trial_logs_still_records_parents(tmp_path): +def test_provenance_without_a_table_manifest(tmp_path): result = _finish_stub_fit(tmp_path, kumo.Trainer(ModelPlan())) - assert result.provenance['trials'] == [] assert result.provenance['parents']['training_table'] == 'tt-1' assert result.provenance['training_table'] == { 'id': 'tt-1', @@ -202,7 +189,11 @@ def test_provenance_without_trial_logs_still_records_parents(tmp_path): def test_provenance_records_custom_table_source(tmp_path): - source_table = {'location': 's3://bucket/exports/tt', 'version': '17'} + source_table = { + 'connector': 'S3', + 'location': 's3://bucket/exports/tt', + 'version': '17' + } manifest = { 'source': 'custom', 'query': 'PREDICT x', @@ -237,11 +228,11 @@ def test_provenance_failure_warns_and_still_publishes(tmp_path, monkeypatch): from kumo.engine.trainer import _provenance def fail(*args): - raise RuntimeError('unreadable trials') + raise RuntimeError('unreadable manifest') - monkeypatch.setattr(_provenance, '_trials', fail) + monkeypatch.setattr(_provenance, '_training_table', fail) trainer = kumo.Trainer(ModelPlan()) - with pytest.warns(UserWarning, match='unreadable trials'): + with pytest.warns(UserWarning, match='unreadable manifest'): result = _finish_stub_fit(tmp_path, trainer) assert 'provenance' not in _roles(result.bundle_dir) assert result.provenance is None From c5fe930a858aa725190f595a9c5d011c58c3d3f2 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 08:14:48 +0000 Subject: [PATCH 07/15] Record fit_mode and source versions in provenance; expose it on Trainer published_model in provenance.json held the shipped checkpoint's fit mode while TrainingResult.published_model is the registry record; it is now fit_mode. selected_trial indexed AutoML log directories that the bundle does not ship, numbered differently from fit events, and had no reader; it is dropped, as is training_table.id, which repeated parents.training_table. The fit snapshot's parents.source_tables (connector, location, version per table) is copied into provenance.json, so a bundle still says which source versions it was trained on after the scratch index is gone. Trainer.provenance returns the record for a fitted or loaded trainer. Verified with test_bundle_provenance.py, test_trainer_registry.py and test_warm_start.py. Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/trainer/_provenance.py | 30 ++++++++++--------- .../kumo/engine/trainer/trainer.py | 18 ++++++----- .../test/test_bundle_provenance.py | 25 ++++++++-------- 3 files changed, 38 insertions(+), 35 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/_provenance.py b/packages/kumo-engine/kumo/engine/trainer/_provenance.py index 4fbce3479..745b3826c 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_provenance.py +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -1,12 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -r"""Descriptive fit record stored as ``provenance.json`` in a model bundle. - -It is integrity-covered by the publication manifest but excluded from the -bundle id, so its volatile fields (times, environment, scheduler) never change -which model a bundle is. -""" +"""Fit record in ``provenance.json``; excluded from the bundle id.""" import json from datetime import datetime, timezone @@ -14,6 +9,7 @@ from typing import Any, Dict, Optional from kumo.engine._manifest import new_manifest, read_manifest, write_manifest +from kumo.engine.graph import GraphSnapshot from kumo.engine.pquery.training_table import TrainingTable from kumo.engine.trainer._spec import MANIFEST_FILE, PROVENANCE_ROLE @@ -25,11 +21,11 @@ def write_provenance( snapshot_id: str, warm_start_bundle_id: Optional[str], started_at: datetime, - published_model: Optional[str], + fit_mode: Optional[str], + snapshot: Optional[GraphSnapshot] = None, ) -> None: from kumo_core.artifacts import PROVENANCE_FILE, add_bundle_file - config = _read_json(bundle_dir / 'model' / 'config.json') or {} finished_at = datetime.now(timezone.utc) tt_id = train_table.tt_id parents = { @@ -47,19 +43,25 @@ def write_provenance( 'seconds': round((finished_at - started_at).total_seconds(), 3), }, - training_table=_training_table(tt_id, train_table.manifest), - selected_trial=config.get('version'), - published_model=published_model, + training_table=_training_table(train_table.manifest), + source_tables=_source_tables(snapshot), + fit_mode=fit_mode, )) add_bundle_file(bundle_dir, PROVENANCE_ROLE, PROVENANCE_FILE) -def _training_table(tt_id: str, - manifest: Optional[Dict[str, Any]]) -> Dict[str, Any]: +def _source_tables( + snapshot: Optional[GraphSnapshot]) -> Optional[Dict[str, Any]]: + manifest = None if snapshot is None else snapshot.manifest + if not manifest or manifest.get('reconstructed'): + return None + return (manifest.get('parents') or {}).get('source_tables') + + +def _training_table(manifest: Optional[Dict[str, Any]]) -> Dict[str, Any]: manifest = manifest or {} parents = manifest.get('parents') or {} record = { - 'id': tt_id, 'source': manifest.get('source'), 'snapshot': parents.get('snapshot'), } diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 07c7bbf86..0423426ae 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -98,11 +98,7 @@ class TrainingResult: @property def provenance(self) -> Optional[Dict[str, Any]]: - """The bundle's recorded fit provenance; None for older bundles. - - Its `published_model` is the shipped checkpoint's fit mode, unlike - the registry record in `TrainingResult.published_model`. - """ + """The bundle's recorded fit provenance; None for older bundles.""" from kumo.engine.trainer._provenance import read_provenance return read_provenance(self.bundle_dir) @@ -329,6 +325,12 @@ 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 provenance(self) -> Optional[Dict[str, Any]]: + """Fit provenance of the fitted or loaded bundle, if recorded.""" + result = self._training_result + return None if result is None else result.provenance + @_without_torch_scatter_hint def fit( self, @@ -541,7 +543,7 @@ def fit( holdout_path, holdout_truncated = _published_holdout( result.best_trainer) - published_model = _published_fit_mode(result.best_trainer) + fit_mode = _published_fit_mode(result.best_trainer) training_result = TrainingResult( metrics=dict(result.metrics), artifact_paths=dict(result.artifact_paths), @@ -561,8 +563,8 @@ def fit( provenance=dict(train_table=train_table, snapshot_id=snapshot.id, warm_start_bundle_id=warm_start_bundle_id, - started_at=started_at, - published_model=published_model), + started_at=started_at, fit_mode=fit_mode, + snapshot=snapshot), ) def _finish_fit( diff --git a/packages/kumo-engine/test/test_bundle_provenance.py b/packages/kumo-engine/test/test_bundle_provenance.py index 999dc552e..9b4725c6e 100644 --- a/packages/kumo-engine/test/test_bundle_provenance.py +++ b/packages/kumo-engine/test/test_bundle_provenance.py @@ -106,14 +106,13 @@ def test_fit_records_listed_provenance(two_trials): 'warm_start': None, } assert provenance['training_table'] == { - 'id': table.tt_id, 'source': 'pql', - 'snapshot': snapshot.id, + 'snapshot': snapshot.id } - assert not {'seed', 'bundle_format_version', 'trials'} & provenance.keys() + assert provenance['source_tables'] == ( + snapshot.manifest['parents']['source_tables']) assert provenance['fit']['seconds'] >= 0 - assert provenance['selected_trial'] in (0, 1) - assert provenance['published_model'] == 'refit_trainval' + assert provenance['fit_mode'] == 'refit_trainval' assert provenance['environment']['kumo_engine'] == kumo.__version__ assert trainer.bundle_id == result.bundle_id @@ -121,7 +120,9 @@ def test_fit_records_listed_provenance(two_trials): def test_bundle_id_is_content_identity(two_trials, tmp_path): trainer, result, *_ = two_trials assert re.fullmatch('[0-9a-f]{64}', result.bundle_id) - assert kumo.Trainer.load(result.bundle_dir).bundle_id == result.bundle_id + loaded = kumo.Trainer.load(result.bundle_dir) + assert loaded.bundle_id == result.bundle_id + assert loaded.provenance == result.provenance copy = tmp_path / 'copy' shutil.copytree(result.bundle_dir, copy) assert kumo.Trainer.load(copy).bundle_id == result.bundle_id @@ -132,7 +133,7 @@ def test_bundle_id_is_content_identity(two_trials, tmp_path): add_bundle_file(copy, 'provenance', 'provenance.json') reloaded = kumo.Trainer.load(copy) assert reloaded.bundle_id == result.bundle_id - assert reloaded._training_result.provenance['fit']['seconds'] == -1 + assert reloaded.provenance['fit']['seconds'] == -1 config = copy / 'model' / 'config.json' config.write_text(config.read_text() + ' ') @@ -149,14 +150,14 @@ def test_bundle_without_provenance_loads_and_predicts_identically( restored = kumo.Trainer.load(old) assert restored.bundle_id == result.bundle_id - assert restored._training_result.provenance is None + assert restored.provenance is None pd.testing.assert_frame_equal(_predict(restored, prediction_table, snapshot), _predict(trainer, prediction_table, snapshot)) (old / 'provenance.json').write_text(json.dumps(result.provenance)) stray = kumo.Trainer.load(old) - assert stray._training_result.provenance is None + assert stray.provenance is None assert stray.bundle_id == result.bundle_id @@ -174,18 +175,17 @@ def _finish_stub_fit(tmp_path, trainer, manifest=None): provenance=dict(train_table=table, snapshot_id='snap', warm_start_bundle_id=None, started_at=datetime.now(timezone.utc), - published_model=None)) + fit_mode=None)) def test_provenance_without_a_table_manifest(tmp_path): result = _finish_stub_fit(tmp_path, kumo.Trainer(ModelPlan())) assert result.provenance['parents']['training_table'] == 'tt-1' assert result.provenance['training_table'] == { - 'id': 'tt-1', 'source': None, 'snapshot': None } - assert result.provenance['published_model'] is None + assert result.provenance['source_tables'] is None def test_provenance_records_custom_table_source(tmp_path): @@ -204,7 +204,6 @@ def test_provenance_records_custom_table_source(tmp_path): } result = _finish_stub_fit(tmp_path, kumo.Trainer(ModelPlan()), manifest) assert result.provenance['training_table'] == { - 'id': 'tt-1', 'source': 'custom', 'snapshot': 'snap', **source_table, From d3de46ac2b1815681a15736857da3143624bdbb1 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 09:35:07 +0000 Subject: [PATCH 08/15] Trim bundle provenance to essentials The provenance tests fitted their own copy of the shared fitted-model fixture, restated manifest fields, tested the private fit-mode helper case by case and spread one back-compat scenario over three tests. Reuse the conftest fitted fixture, merge the load/back-compat tests, fold the missing-manifest case into the custom-source stub test, reduce the kumo-core bundle-id test to the provenance-exclusion contract, drop the fit-mode helper test, find a listed bundle file through one listed_file helper, and shorten docstrings. Behaviour is unchanged. Verified: test_bundle_provenance, test_trainer_registry and kumo-core test_artifacts pass. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-core/src/kumo_core/artifacts.py | 9 +- packages/kumo-core/tests/test_artifacts.py | 36 +--- .../kumo/engine/trainer/_provenance.py | 21 +- .../kumo-engine/kumo/engine/trainer/_spec.py | 13 +- .../kumo/engine/trainer/trainer.py | 3 - .../test/test_bundle_provenance.py | 195 ++++-------------- 6 files changed, 69 insertions(+), 208 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/artifacts.py b/packages/kumo-core/src/kumo_core/artifacts.py index 7ea69b635..7fb91d1ef 100644 --- a/packages/kumo-core/src/kumo_core/artifacts.py +++ b/packages/kumo-core/src/kumo_core/artifacts.py @@ -254,11 +254,10 @@ def _sha256(path: Path) -> str: def verify_model_bundle(bundle: Path) -> str: """Verify fetched artifacts before parsing or staging a model bundle. - Returns the bundle id: the SHA-256 of the sorted ``[relative_path, - sha256]`` pairs of 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; its digest still participates in the id. + 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 diff --git a/packages/kumo-core/tests/test_artifacts.py b/packages/kumo-core/tests/test_artifacts.py index e29622b96..9061d7ee7 100644 --- a/packages/kumo-core/tests/test_artifacts.py +++ b/packages/kumo-core/tests/test_artifacts.py @@ -241,41 +241,23 @@ def best_ckpt(self): assert json.loads(sources.metrics.read_text()) == {"test_accuracy": 0.8} -def test_bundle_id_covers_listed_files_except_provenance( - tmp_path: Path) -> None: + +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 / "training-spec.json").write_text("{}") - entries = json.loads( - (bundle / "publication-manifest.json").read_text())["artifacts"] - pairs = sorted([entry["relative_path"], entry["sha256"]] - for entry in entries) - pairs = sorted(pairs + [[ - "training-spec.json", - hashlib.sha256(b"{}").hexdigest(), - ]]) - expected = hashlib.sha256( - json.dumps(pairs, separators=(",", ":")).encode()).hexdigest() - - assert verify_model_bundle(bundle) == expected - add_bundle_file(bundle, "training_spec", "training-spec.json") - assert verify_model_bundle(bundle) == expected - - (bundle / "provenance.json").write_text('{"a": 1}') - add_bundle_file(bundle, "provenance", "provenance.json") - assert verify_model_bundle(bundle) == expected - (bundle / "provenance.json").write_text('{"a": 2}') - with pytest.raises(ValueError, match="provenance.json"): - verify_model_bundle(bundle) - add_bundle_file(bundle, "provenance", "provenance.json") - assert verify_model_bundle(bundle) == expected + bundle_id = verify_model_bundle(bundle) + for payload in ('{"a": 1}', '{"a": 2}'): + (bundle / "provenance.json").write_text(payload) + add_bundle_file(bundle, "provenance", "provenance.json") + assert verify_model_bundle(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) != expected + assert verify_model_bundle(bundle) != bundle_id diff --git a/packages/kumo-engine/kumo/engine/trainer/_provenance.py b/packages/kumo-engine/kumo/engine/trainer/_provenance.py index 745b3826c..231ee13f8 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_provenance.py +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -3,7 +3,6 @@ # SPDX-License-Identifier: Apache-2.0 """Fit record in ``provenance.json``; excluded from the bundle id.""" -import json from datetime import datetime, timezone from pathlib import Path from typing import Any, Dict, Optional @@ -11,7 +10,7 @@ from kumo.engine._manifest import new_manifest, read_manifest, write_manifest from kumo.engine.graph import GraphSnapshot from kumo.engine.pquery.training_table import TrainingTable -from kumo.engine.trainer._spec import MANIFEST_FILE, PROVENANCE_ROLE +from kumo.engine.trainer._spec import PROVENANCE_ROLE, listed_file def write_provenance( @@ -27,9 +26,8 @@ def write_provenance( from kumo_core.artifacts import PROVENANCE_FILE, add_bundle_file finished_at = datetime.now(timezone.utc) - tt_id = train_table.tt_id parents = { - 'training_table': tt_id, + 'training_table': train_table.tt_id, 'snapshot': snapshot_id, 'warm_start': warm_start_bundle_id, } @@ -75,16 +73,5 @@ def _training_table(manifest: Optional[Dict[str, Any]]) -> Dict[str, Any]: def read_provenance(bundle_dir: Path) -> Optional[Dict[str, Any]]: - manifest = _read_json(bundle_dir / MANIFEST_FILE) or {} - for entry in manifest.get('artifacts') or []: - if entry.get('role') == PROVENANCE_ROLE: - return read_manifest(bundle_dir / entry['relative_path']) - return None - - -def _read_json(path: Path) -> Optional[Dict[str, Any]]: - try: - data = json.loads(path.read_text(encoding='utf-8')) - except (OSError, ValueError): - return None - return data if isinstance(data, dict) else None + path = listed_file(bundle_dir, PROVENANCE_ROLE) + return None if path is None else read_manifest(path) diff --git a/packages/kumo-engine/kumo/engine/trainer/_spec.py b/packages/kumo-engine/kumo/engine/trainer/_spec.py index 637714d45..2c755fab4 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_spec.py +++ b/packages/kumo-engine/kumo/engine/trainer/_spec.py @@ -12,7 +12,7 @@ import json from dataclasses import asdict, dataclass from pathlib import Path -from typing import Any, Dict, Tuple +from typing import Any, Dict, Optional, Tuple from kumo_core.io.filesystem import atomic_write_bytes from kumoapi.model_plan import ModelPlan, TrainingTableGenerationPlan @@ -103,6 +103,17 @@ def read_training_spec(bundle_dir: Path) -> TrainingSpec: ) +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.""" diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 0423426ae..623e7918a 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -83,9 +83,6 @@ class TrainingResult: model bundle. The file is capped at 1 GiB in memory; `holdout_truncated` is true when rows were dropped to fit, and `None` when the size record could not be read. - - `bundle_id` is the content id of the model bundle's scoring files, - computed when the bundle is written or loaded. """ metrics: Dict[str, float] diff --git a/packages/kumo-engine/test/test_bundle_provenance.py b/packages/kumo-engine/test/test_bundle_provenance.py index 9b4725c6e..fff92cd3f 100644 --- a/packages/kumo-engine/test/test_bundle_provenance.py +++ b/packages/kumo-engine/test/test_bundle_provenance.py @@ -8,166 +8,70 @@ from datetime import datetime, timezone from types import SimpleNamespace -import numpy as np import pandas as pd import pytest from kumo_core.artifacts import add_bundle_file -from kumoapi.graph import Edge from kumoapi.model_plan import ModelPlan, TrainingTableGenerationPlan -from kumoapi.typing import Stype import kumo.engine as kumo -from kumo.engine._scratch import get_scratch_dir from kumo.engine.trainer._spec import MANIFEST_FILE -from kumo.engine.trainer.trainer import TrainingResult, _published_fit_mode - - -@pytest.fixture(scope='module') -def two_trials(tmp_path_factory): - root = tmp_path_factory.mktemp('provenance') - previous = get_scratch_dir() - kumo.configure(scratch_dir=root / 'scratch') - try: - rng = np.random.default_rng(0) - data = root / 'data' - data.mkdir() - pd.DataFrame({'ID': np.arange(50)}).to_parquet(data / 'ents.parquet') - pd.DataFrame({ - 'EID': - np.arange(3000), - 'ENT_ID': - rng.integers(0, 50, 3000), - 'TS': - pd.Timestamp('2024-01-01') + - pd.to_timedelta(rng.integers(0, 120, 3000), unit='D'), - }).to_parquet(data / 'events.parquet') - connector = kumo.FileConnector(root_dir=str(data)) - events = kumo.Table.from_source_table( - connector['events.parquet'], primary_key='EID', - time_column='TS').infer_metadata().with_column( - 'ENT_ID', stype=Stype.ID) - graph = kumo.Graph( - tables={ - 'ents': - kumo.Table.from_source_table(connector['ents.parquet'], - primary_key='ID').infer_metadata(), - 'events': - events, - }, - edges=[Edge(src_table='events', fkey='ENT_ID', dst_table='ents')]) - snapshot = graph.snapshot() - pquery = kumo.PredictiveQuery( - graph=graph, query='PREDICT COUNT(events.*, 0, 7, days) > 0 ' - 'FOR EACH ents.ID') - plan = ModelPlan() - plan.training_job.disable_explain = True - plan.training_job.num_experiments = 2 - plan.optimization.max_epochs = 1 - plan.optimization.max_steps_per_epoch = 2 - table = pquery.generate_training_table(snapshot) - trainer = kumo.Trainer(model_plan=plan) - result = trainer.fit(table, snapshot=snapshot) - yield (trainer, result, table, snapshot, - pquery.generate_prediction_table(snapshot)) - finally: - kumo.configure(scratch_dir=previous) - - -def _predict(trainer, table, snapshot): - result = trainer.predict(table, snapshot=snapshot, - binary_classification_threshold=.5) - return result.frame.sort_values('ENTITY').reset_index(drop=True) - - -def _roles(bundle): - manifest = json.loads((bundle / MANIFEST_FILE).read_text()) - return {entry['role']: entry for entry in manifest['artifacts']} - - -def _strip_provenance(bundle): - manifest = json.loads((bundle / MANIFEST_FILE).read_text()) - manifest['artifacts'] = [ - entry for entry in manifest['artifacts'] - if entry['role'] != 'provenance' - ] - (bundle / MANIFEST_FILE).write_text(json.dumps(manifest)) - (bundle / 'provenance.json').unlink() +from kumo.engine.trainer.trainer import TrainingResult -def test_fit_records_listed_provenance(two_trials): - trainer, result, table, snapshot, _ = two_trials - assert _roles(result.bundle_dir)['provenance']['relative_path'] == ( - 'provenance.json') +def test_fit_records_provenance(fitted): + trainer, _, snapshot = fitted + result = trainer._training_result provenance = result.provenance - assert provenance['kind'] == 'model_bundle' - assert provenance['parents'] == { - 'training_table': table.tt_id, - 'snapshot': snapshot.id, - 'warm_start': None, - } + parents = provenance['parents'] + assert (parents['snapshot'], parents['warm_start']) == (snapshot.id, None) + assert parents['training_table'] assert provenance['training_table'] == { 'source': 'pql', 'snapshot': snapshot.id } assert provenance['source_tables'] == ( snapshot.manifest['parents']['source_tables']) - assert provenance['fit']['seconds'] >= 0 assert provenance['fit_mode'] == 'refit_trainval' - assert provenance['environment']['kumo_engine'] == kumo.__version__ + assert re.fullmatch('[0-9a-f]{64}', trainer.bundle_id) assert trainer.bundle_id == result.bundle_id -def test_bundle_id_is_content_identity(two_trials, tmp_path): - trainer, result, *_ = two_trials - assert re.fullmatch('[0-9a-f]{64}', result.bundle_id) +def test_bundle_without_provenance_loads_and_predicts_identically( + fitted, tmp_path): + trainer, prediction_table, snapshot = fitted + result = trainer._training_result loaded = kumo.Trainer.load(result.bundle_dir) - assert loaded.bundle_id == result.bundle_id - assert loaded.provenance == result.provenance - copy = tmp_path / 'copy' - shutil.copytree(result.bundle_dir, copy) - assert kumo.Trainer.load(copy).bundle_id == result.bundle_id - - provenance = json.loads((copy / 'provenance.json').read_text()) - provenance['fit']['seconds'] = -1 - (copy / 'provenance.json').write_text(json.dumps(provenance)) - add_bundle_file(copy, 'provenance', 'provenance.json') - reloaded = kumo.Trainer.load(copy) - assert reloaded.bundle_id == result.bundle_id - assert reloaded.provenance['fit']['seconds'] == -1 - - config = copy / 'model' / 'config.json' - config.write_text(config.read_text() + ' ') - add_bundle_file(copy, 'trainer_config', 'model/config.json') - assert kumo.Trainer.load(copy).bundle_id != result.bundle_id + assert (loaded.bundle_id, loaded.provenance) == (result.bundle_id, + result.provenance) - -def test_bundle_without_provenance_loads_and_predicts_identically( - two_trials, tmp_path): - trainer, result, _, snapshot, prediction_table = two_trials old = tmp_path / 'old' shutil.copytree(result.bundle_dir, old) - _strip_provenance(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.provenance is None - pd.testing.assert_frame_equal(_predict(restored, prediction_table, - snapshot), - _predict(trainer, prediction_table, snapshot)) - (old / 'provenance.json').write_text(json.dumps(result.provenance)) - stray = kumo.Trainer.load(old) - assert stray.provenance is None - assert stray.bundle_id == result.bundle_id + 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 _finish_stub_fit(tmp_path, trainer, manifest=None): - bundle = tmp_path / 'bundle' +def _finish_stub_fit(bundle, trainer, manifest=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 = type('Table', (), {'tt_id': 'tt-1', 'manifest': manifest})() + table = SimpleNamespace(tt_id='tt-1', manifest=manifest) return trainer._finish_fit( TrainingResult({}, {}, bundle), predictive_query='PREDICT x', snapshot_id='snap', export_config=None, @@ -178,17 +82,7 @@ def _finish_stub_fit(tmp_path, trainer, manifest=None): fit_mode=None)) -def test_provenance_without_a_table_manifest(tmp_path): - result = _finish_stub_fit(tmp_path, kumo.Trainer(ModelPlan())) - assert result.provenance['parents']['training_table'] == 'tt-1' - assert result.provenance['training_table'] == { - 'source': None, - 'snapshot': None - } - assert result.provenance['source_tables'] is None - - -def test_provenance_records_custom_table_source(tmp_path): +def test_provenance_records_custom_or_missing_table_source(tmp_path): source_table = { 'connector': 'S3', 'location': 's3://bucket/exports/tt', @@ -196,31 +90,23 @@ def test_provenance_records_custom_table_source(tmp_path): } manifest = { 'source': 'custom', - 'query': 'PREDICT x', 'parents': { 'snapshot': 'snap', 'source_table': source_table }, } - result = _finish_stub_fit(tmp_path, kumo.Trainer(ModelPlan()), manifest) - assert result.provenance['training_table'] == { + custom = _finish_stub_fit(tmp_path / 'custom', kumo.Trainer(ModelPlan()), + manifest) + assert custom.provenance['training_table'] == { 'source': 'custom', 'snapshot': 'snap', **source_table, } - - -@pytest.mark.parametrize('flags,mode', [ - ({}, None), - ({'fitted': True}, 'fit'), - ({'fitted': True, 'refitted_trainval': True}, 'refit_trainval'), - ({'fitted': True, 'refitted_full': True}, 'refit_full'), - ({'fitted': True, 'refitted_trainval': True, 'refitted_full': True}, - 'refit_full'), -]) -def test_published_fit_mode_follows_refit_flags(flags, mode): - trainer = SimpleNamespace(refit_cfg=SimpleNamespace(**flags)) - assert _published_fit_mode(trainer) == mode + older = _finish_stub_fit(tmp_path / 'older', kumo.Trainer(ModelPlan())) + assert older.provenance['training_table'] == { + 'source': None, + 'snapshot': None + } def test_provenance_failure_warns_and_still_publishes(tmp_path, monkeypatch): @@ -232,8 +118,7 @@ def fail(*args): monkeypatch.setattr(_provenance, '_training_table', fail) trainer = kumo.Trainer(ModelPlan()) with pytest.warns(UserWarning, match='unreadable manifest'): - result = _finish_stub_fit(tmp_path, trainer) - assert 'provenance' not in _roles(result.bundle_dir) + result = _finish_stub_fit(tmp_path / 'bundle', trainer) assert result.provenance is None assert re.fullmatch('[0-9a-f]{64}', result.bundle_id) assert trainer._training_result is result From efaeced434717e1c8a120b5adb13629c23f82019 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 18:29:44 +0000 Subject: [PATCH 09/15] Reuse shared helpers for bundle provenance add_bundle_file inlined its own atomic write, the engine re-walked the trainer's refit flags to name the published fit mode, and provenance.json repeated the fit snapshot that training-spec.json already records as snapshot_id. add_bundle_file now uses kumo_core.io.filesystem.atomic_write_bytes. kumo_core's Trainer gains published_fit_mode(), the order _set_fit_mode_to_use already used, and fit records it as fit_mode. provenance.json drops parents.snapshot; training_table.snapshot stays because a table may be generated on a different snapshot than the fit's. training_table.source is DiskGraph's source_mode (pquery or custom_tt). Verified with the bundle provenance, trainer registry, warm start, training-table export, snapshot manifest and kumo-core artifact/training suites. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-core/src/kumo_core/artifacts.py | 12 ++---------- .../kumo/engine/trainer/_provenance.py | 16 ++++++---------- .../kumo/engine/trainer/trainer.py | 14 ++------------ .../test/test_bundle_provenance.py | 19 +++++++++---------- .../test/test_semantic_carry_forward.py | 7 +++---- 5 files changed, 22 insertions(+), 46 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/artifacts.py b/packages/kumo-core/src/kumo_core/artifacts.py index 7fb91d1ef..68b22e2f8 100644 --- a/packages/kumo-core/src/kumo_core/artifacts.py +++ b/packages/kumo-core/src/kumo_core/artifacts.py @@ -363,16 +363,8 @@ def add_bundle_file(bundle_dir: Path, role: str, relative_path: str) -> None: sha256=VitalArtifactPublicationPolicy._sha256(path), ).to_dict() ] - temporary = manifest_path.with_name(f'.{manifest_path.name}.tmp') - try: - with temporary.open('w', encoding='utf-8') as handle: - handle.write(json.dumps(manifest, indent=2, sort_keys=True) + '\n') - handle.flush() - os.fsync(handle.fileno()) - os.replace(temporary, manifest_path) - FilesystemArtifactStore._fsync_directory(bundle_dir) - finally: - temporary.unlink(missing_ok=True) + 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-engine/kumo/engine/trainer/_provenance.py b/packages/kumo-engine/kumo/engine/trainer/_provenance.py index 231ee13f8..62826acf5 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_provenance.py +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -17,7 +17,6 @@ def write_provenance( bundle_dir: Path, *, train_table: TrainingTable, - snapshot_id: str, warm_start_bundle_id: Optional[str], started_at: datetime, fit_mode: Optional[str], @@ -28,7 +27,6 @@ def write_provenance( finished_at = datetime.now(timezone.utc) parents = { 'training_table': train_table.tt_id, - 'snapshot': snapshot_id, 'warm_start': warm_start_bundle_id, } write_manifest( @@ -41,7 +39,7 @@ def write_provenance( 'seconds': round((finished_at - started_at).total_seconds(), 3), }, - training_table=_training_table(train_table.manifest), + training_table=_training_table(train_table), source_tables=_source_tables(snapshot), fit_mode=fit_mode, )) @@ -56,17 +54,15 @@ def _source_tables( return (manifest.get('parents') or {}).get('source_tables') -def _training_table(manifest: Optional[Dict[str, Any]]) -> Dict[str, Any]: - manifest = manifest or {} - parents = manifest.get('parents') or {} +def _training_table(train_table: TrainingTable) -> Dict[str, Any]: + parents = (train_table.manifest or {}).get('parents') or {} record = { - 'source': manifest.get('source'), + 'source': train_table.source_mode, 'snapshot': parents.get('snapshot'), } - if record['source'] == 'custom': - source_table = parents.get('source_table') or {} + if 'source_table' in parents: record.update({ - key: source_table.get(key) + key: (parents['source_table'] or {}).get(key) for key in ('connector', 'location', 'version') }) return record diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 623e7918a..f2a3c95a8 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -207,16 +207,6 @@ def _published_holdout(trainer: Any) -> tuple[Optional[Path], Optional[bool]]: return _find_holdout(trainer.get_log_dir(fit_mode=published)) -def _published_fit_mode(trainer: Any) -> Optional[str]: - refit = getattr(trainer, 'refit_cfg', None) - for mode, flag in (('refit_full', 'refitted_full'), - ('refit_trainval', 'refitted_trainval'), - ('fit', 'fitted')): - if getattr(refit, flag, False) is True: - return mode - return None - - def _collect_garbage(stage: str) -> None: # Lightning links each trainer to its data module and model in reference # cycles; free a finished run now rather than at a later collection. @@ -540,7 +530,8 @@ def fit( holdout_path, holdout_truncated = _published_holdout( result.best_trainer) - fit_mode = _published_fit_mode(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), @@ -558,7 +549,6 @@ def fit( export_config=export_config, model_plan=effective_plan, provenance=dict(train_table=train_table, - snapshot_id=snapshot.id, warm_start_bundle_id=warm_start_bundle_id, started_at=started_at, fit_mode=fit_mode, snapshot=snapshot), diff --git a/packages/kumo-engine/test/test_bundle_provenance.py b/packages/kumo-engine/test/test_bundle_provenance.py index fff92cd3f..5c891c58c 100644 --- a/packages/kumo-engine/test/test_bundle_provenance.py +++ b/packages/kumo-engine/test/test_bundle_provenance.py @@ -23,10 +23,10 @@ def test_fit_records_provenance(fitted): result = trainer._training_result provenance = result.provenance parents = provenance['parents'] - assert (parents['snapshot'], parents['warm_start']) == (snapshot.id, None) - assert parents['training_table'] + assert parents.keys() == {'training_table', 'warm_start'} + assert parents['training_table'] and parents['warm_start'] is None assert provenance['training_table'] == { - 'source': 'pql', + 'source': 'pquery', 'snapshot': snapshot.id } assert provenance['source_tables'] == ( @@ -66,18 +66,18 @@ def predict(model): pd.testing.assert_frame_equal(predict(restored), predict(trainer)) -def _finish_stub_fit(bundle, trainer, manifest=None): +def _finish_stub_fit(bundle, trainer, manifest=None, source_mode=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-1', manifest=manifest) + table = SimpleNamespace(tt_id='tt-1', manifest=manifest, + source_mode=source_mode) 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, snapshot_id='snap', - warm_start_bundle_id=None, + provenance=dict(train_table=table, warm_start_bundle_id=None, started_at=datetime.now(timezone.utc), fit_mode=None)) @@ -89,16 +89,15 @@ def test_provenance_records_custom_or_missing_table_source(tmp_path): 'version': '17' } manifest = { - 'source': 'custom', 'parents': { 'snapshot': 'snap', 'source_table': source_table }, } custom = _finish_stub_fit(tmp_path / 'custom', kumo.Trainer(ModelPlan()), - manifest) + manifest, 'custom_tt') assert custom.provenance['training_table'] == { - 'source': 'custom', + 'source': 'custom_tt', 'snapshot': 'snap', **source_table, } 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', From c7288937b843b252786805b85ef156892ef098b9 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 18:31:54 +0000 Subject: [PATCH 10/15] Compute a fresh bundle's id from its publication manifest Fit re-read and re-hashed every bundle file to compute bundle_id right after publication and add_bundle_file had hashed them; on a bundle with a 1 GiB checkpoint that is about 0.76 s of SHA-256 per fit. kumo-core gains bundle_id_from_manifest, the same id formula over the digests the publication manifest already lists; verify_model_bundle uses the shared formula and still fully verifies bundles on load and fetch. Verified with the kumo-core artifact suite (the two ids agree) and the bundle provenance and warm-start suites. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-core/src/kumo_core/artifacts.py | 23 +++++++++++++++++-- packages/kumo-core/tests/test_artifacts.py | 2 ++ .../kumo/engine/trainer/trainer.py | 4 ++-- 3 files changed, 25 insertions(+), 4 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/artifacts.py b/packages/kumo-core/src/kumo_core/artifacts.py index 68b22e2f8..c5730c55f 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 @@ -339,8 +346,20 @@ def verify_model_bundle(bundle: Path) -> str: if 'training-spec.json' not in seen and spec_path.is_file(): verified.append(('training-spec.json', VitalArtifactPublicationPolicy._sha256(spec_path))) + 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 verified if pair[0] != PROVENANCE_FILE), + sorted(pair for pair in pairs if pair[0] != PROVENANCE_FILE), separators=(',', ':')) return hashlib.sha256(payload.encode('utf-8')).hexdigest() diff --git a/packages/kumo-core/tests/test_artifacts.py b/packages/kumo-core/tests/test_artifacts.py index 9061d7ee7..eadaf82e1 100644 --- a/packages/kumo-core/tests/test_artifacts.py +++ b/packages/kumo-core/tests/test_artifacts.py @@ -16,6 +16,7 @@ VitalArtifactPublicationPolicy, VitalArtifactSources, add_bundle_file, + bundle_id_from_manifest, resolve_vital_artifact_sources, verify_model_bundle, ) @@ -252,6 +253,7 @@ def test_bundle_id_ignores_provenance_only(tmp_path: Path) -> None: (bundle / "provenance.json").write_text(payload) add_bundle_file(bundle, "provenance", "provenance.json") 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"] diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index f2a3c95a8..7f3183295 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -566,7 +566,7 @@ def _finish_fit( provenance: Optional[Mapping[str, Any]] = None, ) -> TrainingResult: """Commit durable fit outputs before making this trainer usable.""" - from kumo_core.artifacts import verify_model_bundle + from kumo_core.artifacts import bundle_id_from_manifest from kumo.engine.trainer._provenance import write_provenance @@ -588,7 +588,7 @@ def _finish_fit( try: training_result = replace( training_result, - bundle_id=verify_model_bundle(training_result.bundle_dir)) + 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) From 58b80d391746c253aac29b4d51771fc461287855 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 21:10:20 +0000 Subject: [PATCH 11/15] Copy a custom table's full source entry into provenance provenance.json picked connector, location and version from the training table's source entry and dropped version_kind and Custom SQL dependencies that the training-table manifest now records. The whole entry is copied, so provenance mirrors the training-table manifest the way source_tables mirror the snapshot manifest. Verified with the bundle provenance and training-table export tests. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/trainer/_provenance.py | 5 +---- packages/kumo-engine/test/test_bundle_provenance.py | 3 ++- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/_provenance.py b/packages/kumo-engine/kumo/engine/trainer/_provenance.py index 62826acf5..12a15ade9 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_provenance.py +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -61,10 +61,7 @@ def _training_table(train_table: TrainingTable) -> Dict[str, Any]: 'snapshot': parents.get('snapshot'), } if 'source_table' in parents: - record.update({ - key: (parents['source_table'] or {}).get(key) - for key in ('connector', 'location', 'version') - }) + record.update(parents['source_table'] or {}) return record diff --git a/packages/kumo-engine/test/test_bundle_provenance.py b/packages/kumo-engine/test/test_bundle_provenance.py index 5c891c58c..d5c2c194c 100644 --- a/packages/kumo-engine/test/test_bundle_provenance.py +++ b/packages/kumo-engine/test/test_bundle_provenance.py @@ -86,7 +86,8 @@ def test_provenance_records_custom_or_missing_table_source(tmp_path): source_table = { 'connector': 'S3', 'location': 's3://bucket/exports/tt', - 'version': '17' + 'version': '17', + 'version_kind': 'timestamp', } manifest = { 'parents': { From f70ee4d5481797de3350ee26822f0ad00ba734ce Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 22:19:39 +0000 Subject: [PATCH 12/15] Store the bundle id in provenance.json provenance.json recorded "id": null, so a provenance file copied out of its bundle (logged to a tracker or registry) could not be matched to the bundle it describes. The id excludes provenance.json, so it is now computed first and written as the provenance id. Verified with test_bundle_provenance.py. Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/trainer/_provenance.py | 3 ++- .../kumo-engine/kumo/engine/trainer/trainer.py | 16 +++++++++------- .../kumo-engine/test/test_bundle_provenance.py | 3 ++- 3 files changed, 13 insertions(+), 9 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/_provenance.py b/packages/kumo-engine/kumo/engine/trainer/_provenance.py index 12a15ade9..f59d9cb98 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_provenance.py +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -16,6 +16,7 @@ def write_provenance( bundle_dir: Path, *, + bundle_id: Optional[str], train_table: TrainingTable, warm_start_bundle_id: Optional[str], started_at: datetime, @@ -32,7 +33,7 @@ def write_provenance( write_manifest( bundle_dir / PROVENANCE_FILE, new_manifest( - 'model_bundle', None, parents, + 'model_bundle', bundle_id, parents, fit={ 'started_at': started_at.isoformat(timespec='seconds'), 'finished_at': finished_at.isoformat(timespec='seconds'), diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 7f3183295..557ea6f23 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -578,13 +578,6 @@ def _finish_fit( snapshot_id=snapshot_id, training_table_generation=training_table_generation, ) - if provenance is not None: - try: - write_provenance(training_result.bundle_dir, **provenance) - except Exception as error: - warnings.warn( - f"Publishing the model bundle without provenance.json: " - f"{error}", UserWarning, stacklevel=3) try: training_result = replace( training_result, @@ -592,6 +585,15 @@ def _finish_fit( except Exception as error: warnings.warn(f"Could not compute the bundle id: {error}", UserWarning, stacklevel=3) + if provenance 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( diff --git a/packages/kumo-engine/test/test_bundle_provenance.py b/packages/kumo-engine/test/test_bundle_provenance.py index d5c2c194c..c20cd54dd 100644 --- a/packages/kumo-engine/test/test_bundle_provenance.py +++ b/packages/kumo-engine/test/test_bundle_provenance.py @@ -33,7 +33,8 @@ def test_fit_records_provenance(fitted): snapshot.manifest['parents']['source_tables']) assert provenance['fit_mode'] == 'refit_trainval' assert re.fullmatch('[0-9a-f]{64}', trainer.bundle_id) - assert trainer.bundle_id == result.bundle_id + assert trainer.bundle_id == result.bundle_id == provenance['id'] + assert kumo.Trainer.load(result.bundle_dir).bundle_id == provenance['id'] def test_bundle_without_provenance_loads_and_predicts_identically( From 42171445e48bb9a2a9ea52dbca01ef6acb989d73 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 22 Oct 2026 18:47:02 +0000 Subject: [PATCH 13/15] Make bundle provenance a frozen dataclass TrainingResult.provenance and Trainer.provenance returned a dict. They now return a BundleProvenance (parents, fit timing, the training table record and source tables) built on the shared manifest dataclasses. provenance.json on disk and the bundle id are unchanged; the training table record still writes its source table's keys inline. Verified with test_bundle_provenance.py and test_warm_start.py. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_manifest.py | 51 +++++++++++ .../kumo/engine/trainer/_provenance.py | 62 +++++++------ .../kumo/engine/trainer/trainer.py | 5 +- .../test/test_bundle_provenance.py | 88 +++++++++++-------- packages/kumo-engine/test/test_warm_start.py | 2 +- 5 files changed, 142 insertions(+), 66 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_manifest.py b/packages/kumo-engine/kumo/engine/_manifest.py index a47a0737f..8b1c7f860 100644 --- a/packages/kumo-engine/kumo/engine/_manifest.py +++ b/packages/kumo-engine/kumo/engine/_manifest.py @@ -16,6 +16,7 @@ List, Mapping, Optional, + Self, Type, TypeVar, Union, @@ -224,6 +225,56 @@ class TrainingTableManifest(Manifest): split_seed: Optional[int] = None +@dataclass(frozen=True, kw_only=True) +class BundleParents(Record): + training_table: str + warm_start: Optional[str] + + +@dataclass(frozen=True, kw_only=True) +class FitTiming(Record): + started_at: str + finished_at: str + seconds: float + + +@dataclass(frozen=True, kw_only=True) +class BundleTrainingTable(Record): + """Written with ``source_table``'s keys inline.""" + + source: Optional[str] + snapshot: Optional[str] + source_table: Optional[SourceTableEntry] = None + + def to_dict(self) -> Dict[str, Any]: + data = {**self.extra, 'source': self.source, 'snapshot': self.snapshot} + if self.source_table is not None: + data.update(self.source_table.to_dict()) + return data + + @classmethod + def from_dict(cls, data: Mapping[str, Any]) -> Self: + rest = { + key: value + for key, value in data.items() if key not in ('source', 'snapshot') + } + if 'version_kind' not in rest: + return cls(source=data.get('source'), snapshot=data.get('snapshot'), + extra=rest) + return cls(source=data.get('source'), snapshot=data.get('snapshot'), + source_table=SourceTableEntry.from_dict(rest)) + + +@dataclass(frozen=True, kw_only=True) +class BundleProvenance(Manifest): + kind: str = 'model_bundle' + parents: BundleParents + fit: FitTiming + training_table: BundleTrainingTable + source_tables: Optional[Dict[str, SourceTableEntry]] + fit_mode: Optional[str] + + 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 index f59d9cb98..3287db604 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_provenance.py +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -5,9 +5,18 @@ from datetime import datetime, timezone from pathlib import Path -from typing import Any, Dict, Optional +from typing import Dict, Optional -from kumo.engine._manifest import new_manifest, read_manifest, write_manifest +from kumo.engine._manifest import ( + BundleParents, + BundleProvenance, + BundleTrainingTable, + FitTiming, + SourceTableEntry, + read_manifest, + stamp, + write_manifest, +) from kumo.engine.graph import GraphSnapshot from kumo.engine.pquery.training_table import TrainingTable from kumo.engine.trainer._spec import PROVENANCE_ROLE, listed_file @@ -26,46 +35,43 @@ def write_provenance( from kumo_core.artifacts import PROVENANCE_FILE, add_bundle_file finished_at = datetime.now(timezone.utc) - parents = { - 'training_table': train_table.tt_id, - 'warm_start': warm_start_bundle_id, - } write_manifest( bundle_dir / PROVENANCE_FILE, - new_manifest( - 'model_bundle', bundle_id, parents, - fit={ - 'started_at': started_at.isoformat(timespec='seconds'), - 'finished_at': finished_at.isoformat(timespec='seconds'), - 'seconds': round((finished_at - started_at).total_seconds(), - 3), - }, + BundleProvenance( + id=bundle_id, + parents=BundleParents(training_table=train_table.tt_id, + warm_start=warm_start_bundle_id), + 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), + ), training_table=_training_table(train_table), source_tables=_source_tables(snapshot), fit_mode=fit_mode, + **stamp(), )) add_bundle_file(bundle_dir, PROVENANCE_ROLE, PROVENANCE_FILE) def _source_tables( - snapshot: Optional[GraphSnapshot]) -> Optional[Dict[str, Any]]: + snapshot: Optional[GraphSnapshot] +) -> Optional[Dict[str, SourceTableEntry]]: manifest = None if snapshot is None else snapshot.manifest - if not manifest or manifest.get('reconstructed'): + if manifest is None or manifest.reconstructed: return None - return (manifest.get('parents') or {}).get('source_tables') + return manifest.parents.source_tables -def _training_table(train_table: TrainingTable) -> Dict[str, Any]: - parents = (train_table.manifest or {}).get('parents') or {} - record = { - 'source': train_table.source_mode, - 'snapshot': parents.get('snapshot'), - } - if 'source_table' in parents: - record.update(parents['source_table'] or {}) - return record +def _training_table(train_table: TrainingTable) -> BundleTrainingTable: + manifest = train_table.manifest + return BundleTrainingTable( + source=train_table.source_mode, + snapshot=None if manifest is None else manifest.parents.snapshot, + source_table=None + if manifest is None else manifest.parents.source_table) -def read_provenance(bundle_dir: Path) -> Optional[Dict[str, Any]]: +def read_provenance(bundle_dir: Path) -> Optional[BundleProvenance]: path = listed_file(bundle_dir, PROVENANCE_ROLE) - return None if path is None else read_manifest(path) + return None if path is None else read_manifest(path, BundleProvenance) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 557ea6f23..c318590f0 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 BundleProvenance from kumo.engine._scratch import new_run_id, trainer_scratch_dir from kumo.engine.connector.base import Connector from kumo.engine.graph import GraphSnapshot @@ -94,7 +95,7 @@ class TrainingResult: bundle_id: Optional[str] = None @property - def provenance(self) -> Optional[Dict[str, Any]]: + def provenance(self) -> Optional[BundleProvenance]: """The bundle's recorded fit provenance; None for older bundles.""" from kumo.engine.trainer._provenance import read_provenance @@ -313,7 +314,7 @@ def bundle_id(self) -> Optional[str]: return getattr(self._training_result, 'bundle_id', None) @property - def provenance(self) -> Optional[Dict[str, Any]]: + def provenance(self) -> Optional[BundleProvenance]: """Fit provenance of the fitted or loaded bundle, if recorded.""" result = self._training_result return None if result is None else result.provenance diff --git a/packages/kumo-engine/test/test_bundle_provenance.py b/packages/kumo-engine/test/test_bundle_provenance.py index c20cd54dd..608564460 100644 --- a/packages/kumo-engine/test/test_bundle_provenance.py +++ b/packages/kumo-engine/test/test_bundle_provenance.py @@ -14,6 +14,14 @@ from kumoapi.model_plan import ModelPlan, TrainingTableGenerationPlan import kumo.engine as kumo +from kumo.engine._manifest import ( + BundleProvenance, + BundleTrainingTable, + SourceTableEntry, + TrainingTableParents, + read_manifest, + write_manifest, +) from kumo.engine.trainer._spec import MANIFEST_FILE from kumo.engine.trainer.trainer import TrainingResult @@ -22,19 +30,36 @@ def test_fit_records_provenance(fitted): trainer, _, snapshot = fitted result = trainer._training_result provenance = result.provenance - parents = provenance['parents'] - assert parents.keys() == {'training_table', 'warm_start'} - assert parents['training_table'] and parents['warm_start'] is None - assert provenance['training_table'] == { - 'source': 'pquery', - 'snapshot': snapshot.id - } - assert provenance['source_tables'] == ( - snapshot.manifest['parents']['source_tables']) - assert provenance['fit_mode'] == 'refit_trainval' + parents = provenance.parents + assert parents.training_table and parents.warm_start is None + assert parents.extra == {} + assert provenance.training_table == BundleTrainingTable( + source='pquery', snapshot=snapshot.id) + assert provenance.source_tables == snapshot.manifest.parents.source_tables + assert provenance.fit_mode == 'refit_trainval' assert re.fullmatch('[0-9a-f]{64}', trainer.bundle_id) - assert trainer.bundle_id == result.bundle_id == provenance['id'] - assert kumo.Trainer.load(result.bundle_dir).bundle_id == provenance['id'] + assert trainer.bundle_id == result.bundle_id == provenance.id + assert kumo.Trainer.load(result.bundle_dir).bundle_id == provenance.id + + +def test_provenance_round_trips_unchanged(fitted, tmp_path): + trainer, _, _ = fitted + newer = json.loads( + (trainer._training_result.bundle_dir / 'provenance.json').read_text()) + newer['future'] = True + newer['fit']['gpu_seconds'] = 1.5 + newer['training_table'].update(connector='S3', location='s3://b/tt', + version='1', version_kind='exact') + source = tmp_path / 'newer.json' + source.write_text(json.dumps(newer, indent=2, sort_keys=True) + '\n') + + provenance = read_manifest(source, BundleProvenance) + write_manifest(tmp_path / 'out.json', provenance) + + assert (tmp_path / 'out.json').read_bytes() == source.read_bytes() + assert provenance.extra == {'future': True} + assert provenance.fit.extra == {'gpu_seconds': 1.5} + assert provenance.training_table.source_table.location == 's3://b/tt' def test_bundle_without_provenance_loads_and_predicts_identically( @@ -79,35 +104,28 @@ def _finish_stub_fit(bundle, trainer, manifest=None, source_mode=None): snapshot_id='snap', export_config=None, training_table_generation=TrainingTableGenerationPlan(), provenance=dict(train_table=table, warm_start_bundle_id=None, - started_at=datetime.now(timezone.utc), - fit_mode=None)) + started_at=datetime.now(timezone.utc), fit_mode=None)) def test_provenance_records_custom_or_missing_table_source(tmp_path): - source_table = { - 'connector': 'S3', - 'location': 's3://bucket/exports/tt', - 'version': '17', - 'version_kind': 'timestamp', - } - manifest = { - 'parents': { - 'snapshot': 'snap', - 'source_table': source_table - }, - } + source_table = SourceTableEntry(connector='S3', + location='s3://bucket/exports/tt', + version='17', version_kind='timestamp') + manifest = SimpleNamespace(parents=TrainingTableParents( + snapshot='snap', source_table=source_table)) custom = _finish_stub_fit(tmp_path / 'custom', kumo.Trainer(ModelPlan()), manifest, 'custom_tt') - assert custom.provenance['training_table'] == { - 'source': 'custom_tt', - 'snapshot': 'snap', - **source_table, - } + assert custom.provenance.training_table == BundleTrainingTable( + source='custom_tt', snapshot='snap', source_table=source_table) + assert json.loads((tmp_path / 'custom' / + 'provenance.json').read_text())['training_table'] == { + 'source': 'custom_tt', + 'snapshot': 'snap', + **source_table.to_dict() + } older = _finish_stub_fit(tmp_path / 'older', kumo.Trainer(ModelPlan())) - assert older.provenance['training_table'] == { - 'source': None, - 'snapshot': None - } + assert older.provenance.training_table == BundleTrainingTable( + source=None, snapshot=None) def test_provenance_failure_warns_and_still_publishes(tmp_path, monkeypatch): diff --git a/packages/kumo-engine/test/test_warm_start.py b/packages/kumo-engine/test/test_warm_start.py index d60a645cb..3370f632f 100644 --- a/packages/kumo-engine/test/test_warm_start.py +++ b/packages/kumo-engine/test/test_warm_start.py @@ -182,7 +182,7 @@ 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.provenance['parents']['warm_start'] == source.bundle_id + assert second.provenance.parents.warm_start == source.bundle_id assert destination.model_plan == configured_plan assert _hash_bundle(durable) == original_hashes assert _hash_bundle(source_bundle) == original_hashes From 13209487d238adf9a2c97db97e7f687c32984851 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 24 Oct 2026 17:37:42 +0000 Subject: [PATCH 14/15] Embed the training table in a self-contained model manifest; store models Bundle provenance copied parts of the training table and snapshot into bespoke fields (training_table source/snapshot, source_tables) and named the warm-start model by id only. - ModelManifest (was BundleProvenance), kind model, file _provenance.json: parents.training_table embeds the training-table manifest (with its snapshot), parents.warm_start the warm-start model's manifest with its own model parents cut to stubs, so retrain chains stay flat; a missing parent is a stub. The bundle id still leaves the manifest out (old and new names). - fit(mlops=): with MLOps on, the bundle is stored under /models//; a complete copy is kept, a partial one replaced; failures warn. Each fit returns a new TrainingResult (.id, .manifest); Trainer.result and Trainer.manifest expose it. Verified: test_bundle_provenance.py (manifest contents and round trip, stubs, flat warm-start chains, new result per fit, root store and repair, write failure), warm start, registry, holdout and kumo-core artifact suites. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-core/src/kumo_core/artifacts.py | 6 +- packages/kumo-core/tests/test_artifacts.py | 9 +- packages/kumo-engine/docs/mlops.md | 11 ++ packages/kumo-engine/kumo/engine/__init__.py | 7 +- packages/kumo-engine/kumo/engine/_manifest.py | 51 ++---- .../kumo/engine/trainer/_provenance.py | 72 +++++---- .../kumo/engine/trainer/trainer.py | 85 ++++++++-- .../test/test_bundle_provenance.py | 153 ++++++++++-------- packages/kumo-engine/test/test_warm_start.py | 3 +- 9 files changed, 242 insertions(+), 155 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/artifacts.py b/packages/kumo-core/src/kumo_core/artifacts.py index c5730c55f..c8e45e4a4 100644 --- a/packages/kumo-core/src/kumo_core/artifacts.py +++ b/packages/kumo-core/src/kumo_core/artifacts.py @@ -255,7 +255,8 @@ def _sha256(path: Path) -> str: return digest.hexdigest() -PROVENANCE_FILE = 'provenance.json' +#: 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: @@ -359,7 +360,8 @@ def bundle_id_from_manifest(bundle: Path) -> str: def _bundle_id(pairs: Iterable[tuple[str, str]]) -> str: payload = json.dumps( - sorted(pair for pair in pairs if pair[0] != PROVENANCE_FILE), + 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() diff --git a/packages/kumo-core/tests/test_artifacts.py b/packages/kumo-core/tests/test_artifacts.py index eadaf82e1..8828667ae 100644 --- a/packages/kumo-core/tests/test_artifacts.py +++ b/packages/kumo-core/tests/test_artifacts.py @@ -242,16 +242,17 @@ def best_ckpt(self): 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 payload in ('{"a": 1}', '{"a": 2}'): - (bundle / "provenance.json").write_text(payload) - add_bundle_file(bundle, "provenance", "provenance.json") + 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 = [ 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 8b1c7f860..81265178f 100644 --- a/packages/kumo-engine/kumo/engine/_manifest.py +++ b/packages/kumo-engine/kumo/engine/_manifest.py @@ -16,7 +16,6 @@ List, Mapping, Optional, - Self, Type, TypeVar, Union, @@ -226,9 +225,10 @@ class TrainingTableManifest(Manifest): @dataclass(frozen=True, kw_only=True) -class BundleParents(Record): - training_table: str - warm_start: Optional[str] +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) @@ -239,40 +239,15 @@ class FitTiming(Record): @dataclass(frozen=True, kw_only=True) -class BundleTrainingTable(Record): - """Written with ``source_table``'s keys inline.""" - - source: Optional[str] - snapshot: Optional[str] - source_table: Optional[SourceTableEntry] = None - - def to_dict(self) -> Dict[str, Any]: - data = {**self.extra, 'source': self.source, 'snapshot': self.snapshot} - if self.source_table is not None: - data.update(self.source_table.to_dict()) - return data - - @classmethod - def from_dict(cls, data: Mapping[str, Any]) -> Self: - rest = { - key: value - for key, value in data.items() if key not in ('source', 'snapshot') - } - if 'version_kind' not in rest: - return cls(source=data.get('source'), snapshot=data.get('snapshot'), - extra=rest) - return cls(source=data.get('source'), snapshot=data.get('snapshot'), - source_table=SourceTableEntry.from_dict(rest)) - - -@dataclass(frozen=True, kw_only=True) -class BundleProvenance(Manifest): - kind: str = 'model_bundle' - parents: BundleParents - fit: FitTiming - training_table: BundleTrainingTable - source_tables: Optional[Dict[str, SourceTableEntry]] - fit_mode: Optional[str] +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: diff --git a/packages/kumo-engine/kumo/engine/trainer/_provenance.py b/packages/kumo-engine/kumo/engine/trainer/_provenance.py index 3287db604..4429186dd 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_provenance.py +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -1,23 +1,23 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Fit record in ``provenance.json``; excluded from the bundle id.""" +"""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 Dict, Optional +from typing import Optional from kumo.engine._manifest import ( - BundleParents, - BundleProvenance, - BundleTrainingTable, FitTiming, - SourceTableEntry, + ModelManifest, + ModelParents, + embedded, read_manifest, stamp, + stub, write_manifest, ) -from kumo.engine.graph import GraphSnapshot from kumo.engine.pquery.training_table import TrainingTable from kumo.engine.trainer._spec import PROVENANCE_ROLE, listed_file @@ -25,53 +25,63 @@ def write_provenance( bundle_dir: Path, *, - bundle_id: Optional[str], + bundle_id: str, train_table: TrainingTable, - warm_start_bundle_id: Optional[str], + warm_start: Optional[ModelManifest], started_at: datetime, fit_mode: Optional[str], - snapshot: Optional[GraphSnapshot] = None, ) -> 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, - BundleProvenance( + ModelManifest( id=bundle_id, - parents=BundleParents(training_table=train_table.tt_id, - warm_start=warm_start_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), ), - training_table=_training_table(train_table), - source_tables=_source_tables(snapshot), fit_mode=fit_mode, **stamp(), )) add_bundle_file(bundle_dir, PROVENANCE_ROLE, PROVENANCE_FILE) -def _source_tables( - snapshot: Optional[GraphSnapshot] -) -> Optional[Dict[str, SourceTableEntry]]: - manifest = None if snapshot is None else snapshot.manifest - if manifest is None or manifest.reconstructed: - return None - return manifest.parents.source_tables +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) -def _training_table(train_table: TrainingTable) -> BundleTrainingTable: - manifest = train_table.manifest - return BundleTrainingTable( - source=train_table.source_mode, - snapshot=None if manifest is None else manifest.parents.snapshot, - source_table=None - if manifest is None else manifest.parents.source_table) + return replace(copy, parents=replace( + parents, warm_start=cut(parents.warm_start), + derived_from=cut(parents.derived_from))) -def read_provenance(bundle_dir: Path) -> Optional[BundleProvenance]: +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, BundleProvenance) + return None if path is None else read_manifest(path, ModelManifest) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index c318590f0..bdf62243e 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -52,7 +52,7 @@ require_embedding_models, require_trained_embedding_models, ) -from kumo.engine._manifest import BundleProvenance +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 @@ -95,8 +95,13 @@ class TrainingResult: bundle_id: Optional[str] = None @property - def provenance(self) -> Optional[BundleProvenance]: - """The bundle's recorded fit provenance; None for older bundles.""" + 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) @@ -197,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 @@ -314,10 +354,15 @@ def bundle_id(self) -> Optional[str]: return getattr(self._training_result, 'bundle_id', None) @property - def provenance(self) -> Optional[BundleProvenance]: - """Fit provenance of the fitted or loaded bundle, if recorded.""" + 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.provenance + return None if result is None else result.manifest @_without_torch_scatter_hint def fit( @@ -327,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`. @@ -355,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, @@ -383,6 +434,10 @@ 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 @@ -549,10 +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_bundle_id=warm_start_bundle_id, - started_at=started_at, fit_mode=fit_mode, - snapshot=snapshot), + 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( @@ -565,6 +623,7 @@ def _finish_fit( 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 @@ -586,15 +645,17 @@ def _finish_fit( except Exception as error: warnings.warn(f"Could not compute the bundle id: {error}", UserWarning, stacklevel=3) - if provenance is not None: + 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"Publishing the model bundle without _provenance.json: " f"{error}", UserWarning, stacklevel=3) + if settings is not None and settings.wants('model'): + _store_model(training_result, settings) if export_config is not None: spec = read_training_spec(training_result.bundle_dir) published = export_config.connector.sink_backend.publish_model( diff --git a/packages/kumo-engine/test/test_bundle_provenance.py b/packages/kumo-engine/test/test_bundle_provenance.py index 608564460..5b44d3153 100644 --- a/packages/kumo-engine/test/test_bundle_provenance.py +++ b/packages/kumo-engine/test/test_bundle_provenance.py @@ -15,61 +15,56 @@ import kumo.engine as kumo from kumo.engine._manifest import ( - BundleProvenance, - BundleTrainingTable, - SourceTableEntry, - TrainingTableParents, + 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_provenance(fitted): +def test_fit_records_a_self_contained_manifest(fitted): trainer, _, snapshot = fitted - result = trainer._training_result - provenance = result.provenance - parents = provenance.parents - assert parents.training_table and parents.warm_start is None - assert parents.extra == {} - assert provenance.training_table == BundleTrainingTable( - source='pquery', snapshot=snapshot.id) - assert provenance.source_tables == snapshot.manifest.parents.source_tables - assert provenance.fit_mode == 'refit_trainval' + 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.bundle_id == provenance.id - assert kumo.Trainer.load(result.bundle_dir).bundle_id == provenance.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_provenance_round_trips_unchanged(fitted, tmp_path): +def test_manifest_round_trips_unchanged(fitted, tmp_path): trainer, _, _ = fitted newer = json.loads( - (trainer._training_result.bundle_dir / 'provenance.json').read_text()) + (trainer.result.bundle_dir / '_provenance.json').read_text()) newer['future'] = True newer['fit']['gpu_seconds'] = 1.5 - newer['training_table'].update(connector='S3', location='s3://b/tt', - version='1', version_kind='exact') source = tmp_path / 'newer.json' source.write_text(json.dumps(newer, indent=2, sort_keys=True) + '\n') - provenance = read_manifest(source, BundleProvenance) - write_manifest(tmp_path / 'out.json', provenance) + manifest = read_manifest(source, ModelManifest) + write_manifest(tmp_path / 'out.json', manifest) assert (tmp_path / 'out.json').read_bytes() == source.read_bytes() - assert provenance.extra == {'future': True} - assert provenance.fit.extra == {'gpu_seconds': 1.5} - assert provenance.training_table.source_table.location == 's3://b/tt' + assert manifest.extra == {'future': True} + assert manifest.fit.extra == {'gpu_seconds': 1.5} -def test_bundle_without_provenance_loads_and_predicts_identically( +def test_bundle_without_a_manifest_loads_and_predicts_identically( fitted, tmp_path): trainer, prediction_table, snapshot = fitted - result = trainer._training_result - loaded = kumo.Trainer.load(result.bundle_dir) - assert (loaded.bundle_id, loaded.provenance) == (result.bundle_id, - result.provenance) - + result = trainer.result old = tmp_path / 'old' shutil.copytree(result.bundle_dir, old) manifest = json.loads((old / MANIFEST_FILE).read_text()) @@ -78,10 +73,10 @@ def test_bundle_without_provenance_loads_and_predicts_identically( if entry['role'] != 'provenance' ] (old / MANIFEST_FILE).write_text(json.dumps(manifest)) - (old / 'provenance.json').unlink() + (old / '_provenance.json').unlink() restored = kumo.Trainer.load(old) assert restored.bundle_id == result.bundle_id - assert restored.provenance is None + assert restored.manifest is None def predict(model): return model.predict( @@ -92,52 +87,78 @@ def predict(model): pd.testing.assert_frame_equal(predict(restored), predict(trainer)) -def _finish_stub_fit(bundle, trainer, manifest=None, source_mode=None): +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-1', manifest=manifest, - source_mode=source_mode) + 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_bundle_id=None, - started_at=datetime.now(timezone.utc), fit_mode=None)) - - -def test_provenance_records_custom_or_missing_table_source(tmp_path): - source_table = SourceTableEntry(connector='S3', - location='s3://bucket/exports/tt', - version='17', version_kind='timestamp') - manifest = SimpleNamespace(parents=TrainingTableParents( - snapshot='snap', source_table=source_table)) - custom = _finish_stub_fit(tmp_path / 'custom', kumo.Trainer(ModelPlan()), - manifest, 'custom_tt') - assert custom.provenance.training_table == BundleTrainingTable( - source='custom_tt', snapshot='snap', source_table=source_table) - assert json.loads((tmp_path / 'custom' / - 'provenance.json').read_text())['training_table'] == { - 'source': 'custom_tt', - 'snapshot': 'snap', - **source_table.to_dict() - } - older = _finish_stub_fit(tmp_path / 'older', kumo.Trainer(ModelPlan())) - assert older.provenance.training_table == BundleTrainingTable( - source=None, snapshot=None) - - -def test_provenance_failure_warns_and_still_publishes(tmp_path, monkeypatch): + 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): + def fail(*args, **kwargs): raise RuntimeError('unreadable manifest') - monkeypatch.setattr(_provenance, '_training_table', fail) + 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.provenance is None + assert result.manifest is None assert re.fullmatch('[0-9a-f]{64}', result.bundle_id) - assert trainer._training_result is result + assert trainer.result is result diff --git a/packages/kumo-engine/test/test_warm_start.py b/packages/kumo-engine/test/test_warm_start.py index 3370f632f..25a109489 100644 --- a/packages/kumo-engine/test/test_warm_start.py +++ b/packages/kumo-engine/test/test_warm_start.py @@ -182,7 +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.provenance.parents.warm_start == source.bundle_id + 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 From 456b97d11103e5ceee1b47f3b82e8a43e7aee7ef Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 24 Oct 2026 17:41:09 +0000 Subject: [PATCH 15/15] Store the model under the root after the plain export Order per the series: a plain export_config publish first, then the artifact-root store. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/trainer/trainer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index bdf62243e..1acc5e100 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -654,8 +654,6 @@ def _finish_fit( warnings.warn( f"Publishing the model bundle without _provenance.json: " f"{error}", UserWarning, stacklevel=3) - if settings is not None and settings.wants('model'): - _store_model(training_result, settings) if export_config is not None: spec = read_training_spec(training_result.bundle_dir) published = export_config.connector.sink_backend.publish_model( @@ -672,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