From 206d88b0358362d1e4783b32e3a98710c927fe61 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 24 Oct 2026 18:17:11 +0000 Subject: [PATCH] Carry scoring defaults in a derived model bundle A threshold tuned on the holdout after fit has nowhere to live: every scheduled score job must pass it again, and a registry version cannot say which threshold it was validated with. TrainingResult.with_scoring_defaults(ScoringDefaults(...)) copies the bundle, adds scoring-defaults.json and returns a TrainingResult for the copy (its own bundle id). Its _provenance.json embeds the parent's ModelManifest in parents.derived_from (the parent's own model parents cut to ids). predict falls back to TrainingResult.scoring_defaults when a call omits the threshold or class count. Trainer.export(config) publishes a fitted or loaded bundle. Verified with test_scoring_defaults.py and kumo-core test_artifacts.py. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-core/src/kumo_core/artifacts.py | 41 +++++- .../kumo-core/src/kumo_core/automl/runtime.py | 5 +- packages/kumo-core/tests/test_artifacts.py | 8 ++ packages/kumo-engine/docs/mlops.md | 15 ++ packages/kumo-engine/kumo/engine/__init__.py | 2 + .../kumo/engine/trainer/__init__.py | 2 + .../kumo/engine/trainer/_provenance.py | 23 +++ .../kumo/engine/trainer/_scoring.py | 81 +++++++++++ .../kumo-engine/kumo/engine/trainer/_spec.py | 4 +- .../kumo/engine/trainer/trainer.py | 132 ++++++++++++++---- .../test/test_prediction_completion.py | 4 +- .../test/test_prediction_workers.py | 4 +- .../kumo-engine/test/test_scoring_defaults.py | 125 +++++++++++++++++ .../test/test_semantic_carry_forward.py | 15 +- 14 files changed, 420 insertions(+), 41 deletions(-) create mode 100644 packages/kumo-engine/kumo/engine/trainer/_scoring.py create mode 100644 packages/kumo-engine/test/test_scoring_defaults.py diff --git a/packages/kumo-core/src/kumo_core/artifacts.py b/packages/kumo-core/src/kumo_core/artifacts.py index c8e45e4a4..8e4f8cd64 100644 --- a/packages/kumo-core/src/kumo_core/artifacts.py +++ b/packages/kumo-core/src/kumo_core/artifacts.py @@ -14,6 +14,7 @@ from dataclasses import dataclass, field from pathlib import Path from typing import ( + TYPE_CHECKING, Any, ClassVar, Iterable, @@ -26,6 +27,9 @@ from .io.json import load_json, save_json from .manifest import RunManifest, RunPaths +if TYPE_CHECKING: + from .types import TaskType + #: Every engine artifact's manifest file; underscore-prefixed so Parquet #: readers skip it. PROVENANCE_FILE = '_provenance.json' @@ -384,8 +388,35 @@ def add_bundle_file(bundle_dir: Path, role: str, relative_path: str) -> None: sha256=VitalArtifactPublicationPolicy._sha256(path), ).to_dict() ] + _write_bundle_manifest(bundle_dir, manifest) + + +def remove_bundle_file(bundle_dir: Path, role: str) -> None: + """Delete the file listed under ``role`` and its manifest entry.""" + from .sources import layout + + bundle_dir = Path(bundle_dir) + manifest_path = bundle_dir / VitalArtifactPublicationPolicy.MANIFEST_NAME + manifest = json.loads(manifest_path.read_text(encoding='utf-8')) + entries = manifest.get('artifacts') or [] + removed = [entry for entry in entries if entry.get('role') == role] + if not removed: + return + manifest['artifacts'] = [ + entry for entry in entries if entry.get('role') != role + ] + _write_bundle_manifest(bundle_dir, manifest) + for entry in removed: + (bundle_dir / layout.confined(entry['relative_path'])).unlink( + missing_ok=True) + + +def _write_bundle_manifest(bundle_dir: Path, + manifest: Mapping[str, Any]) -> None: payload = json.dumps(manifest, indent=2, sort_keys=True) + '\n' - atomic_write_bytes(manifest_path, payload.encode('utf-8')) + atomic_write_bytes( + bundle_dir / VitalArtifactPublicationPolicy.MANIFEST_NAME, + payload.encode('utf-8')) @runtime_checkable @@ -477,6 +508,14 @@ def resolve_vital_artifact_sources( } +def bundle_task_type(artifacts: Mapping[str, str]) -> TaskType: + """The task type a bundle's prediction schema records.""" + from .types import TaskType + + schema = load_json(artifacts['prediction_schema'], dict) + return TaskType(schema['task_type']) + + def stage_model_bundle(artifacts: Mapping[str, str], dest: Path) -> Path: r"""Copies a trained run's artifacts into the layout scoring expects. diff --git a/packages/kumo-core/src/kumo_core/automl/runtime.py b/packages/kumo-core/src/kumo_core/automl/runtime.py index 209680edf..880246456 100644 --- a/packages/kumo-core/src/kumo_core/automl/runtime.py +++ b/packages/kumo-core/src/kumo_core/automl/runtime.py @@ -196,9 +196,10 @@ def resolve_warm_start( encoders; they do not retain hosted AutoML directory layouts. """ if artifacts is not None: - source_schema = load_json(artifacts["prediction_schema"], dict) + from kumo_core.artifacts import bundle_task_type + validate_warm_start_task_type(task.task_type, - TaskType(source_schema["task_type"])) + bundle_task_type(artifacts)) checkpoint = artifacts["selected_checkpoint"] plan_path = artifacts["resolved_model_plan"] encoder_dir = posixpath.dirname(artifacts["encoders"]) diff --git a/packages/kumo-core/tests/test_artifacts.py b/packages/kumo-core/tests/test_artifacts.py index 8828667ae..869b89ac2 100644 --- a/packages/kumo-core/tests/test_artifacts.py +++ b/packages/kumo-core/tests/test_artifacts.py @@ -17,6 +17,7 @@ VitalArtifactSources, add_bundle_file, bundle_id_from_manifest, + bundle_task_type, resolve_vital_artifact_sources, verify_model_bundle, ) @@ -264,3 +265,10 @@ def test_bundle_id_ignores_provenance_only(tmp_path: Path) -> None: (bundle / "metrics.json").write_text('{"auc": 0}') add_bundle_file(bundle, "metrics", "metrics.json") assert verify_model_bundle(bundle) != bundle_id + + +def test_bundle_task_type_reads_the_prediction_schema(tmp_path: Path) -> None: + schema = tmp_path / "prediction-schema.json" + schema.write_text('{"task_type": "regression"}') + assert bundle_task_type({"prediction_schema": str(schema)}).value == ( + "regression") diff --git a/packages/kumo-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index a71d07f71..8a12d1473 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -47,6 +47,21 @@ An embedded model keeps only the ids of its own model parents, so a long chain o With MLOps on, each bundle is stored under `/models//`. A complete copy that is already there is kept; a partial one is replaced. +### Scoring defaults + +A model can carry the prediction arguments `predict` uses when a call omits them: + +```python +from kumo.engine import ModelOutputConfig, ScoringDefaults + +tuned = result.with_scoring_defaults( + ScoringDefaults(binary_classification_threshold=0.37)) +tuned.predict(snapshot) # thresholds at 0.37 +tuned.scoring_defaults # ScoringDefaults(binary_classification_threshold=0.37, ...) +``` + +`with_scoring_defaults` returns a `TrainingResult` for a copy of the bundle with its own id. The copy's manifest embeds the original's in `parents.derived_from`. The threshold applies to binary and multilabel classification, `num_classes_to_return` to multiclass, multilabel and link prediction. Explicit `predict` arguments and a query's `TOP k` win; `ScoringDefaults()` removes the defaults. `Trainer.export(config)` publishes a fitted or loaded bundle like `fit(export_config=...)`. + ### Moving snapshots between machines To build an index on one machine and train or score on another, export it to a file destination and load it there: diff --git a/packages/kumo-engine/kumo/engine/__init__.py b/packages/kumo-engine/kumo/engine/__init__.py index db8c0743e..b9e467763 100644 --- a/packages/kumo-engine/kumo/engine/__init__.py +++ b/packages/kumo-engine/kumo/engine/__init__.py @@ -57,6 +57,7 @@ ModelOutputConfig, PredictionOutputConfig, PredictionResult, + ScoringDefaults, SnapshotOutputConfig, Trainer, TrainingResult, @@ -101,6 +102,7 @@ "Trainer", "TrainingResult", "PredictionResult", + "ScoringDefaults", "ModelOutputConfig", "PredictionOutputConfig", "TrainingTableOutputConfig", diff --git a/packages/kumo-engine/kumo/engine/trainer/__init__.py b/packages/kumo-engine/kumo/engine/trainer/__init__.py index 0f4c1cfff..abc11c7e3 100644 --- a/packages/kumo-engine/kumo/engine/trainer/__init__.py +++ b/packages/kumo-engine/kumo/engine/trainer/__init__.py @@ -2,6 +2,7 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from kumo.engine.trainer._scoring import ScoringDefaults from kumo.engine.trainer.output_config import ( ModelOutputConfig, PredictionOutputConfig, @@ -18,6 +19,7 @@ "Trainer", "TrainingResult", "PredictionResult", + "ScoringDefaults", "ModelOutputConfig", "PredictionOutputConfig", "SnapshotOutputConfig", diff --git a/packages/kumo-engine/kumo/engine/trainer/_provenance.py b/packages/kumo-engine/kumo/engine/trainer/_provenance.py index 28e23acf4..dfd6f151e 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_provenance.py +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -89,3 +89,26 @@ def parent_model(bundle_dir: Path, bundle_id: Optional[str] def read_provenance(bundle_dir: Path) -> Optional[ModelManifest]: path = listed_file(bundle_dir, PROVENANCE_ROLE) return None if path is None else read_manifest(path, ModelManifest) + + +def record_derivation(bundle_dir: Path, parent_dir: Path, + parent_id: Optional[str]) -> None: + r"""Writes the manifest of ``bundle_dir``, a copy of the bundle at + ``parent_dir`` with other scoring defaults.""" + from kumo_core.artifacts import ( + PROVENANCE_FILE, + add_bundle_file, + bundle_id_from_manifest, + ) + + parent = parent_model(parent_dir, parent_id) + write_manifest( + bundle_dir / PROVENANCE_FILE, + ModelManifest( + id=bundle_id_from_manifest(bundle_dir), + parents=ModelParents( + derived_from=None if parent is None else embedded_model( + parent)), + **stamp(), + )) + add_bundle_file(bundle_dir, PROVENANCE_ROLE, PROVENANCE_FILE) diff --git a/packages/kumo-engine/kumo/engine/trainer/_scoring.py b/packages/kumo-engine/kumo/engine/trainer/_scoring.py new file mode 100644 index 000000000..c6a1f0655 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/trainer/_scoring.py @@ -0,0 +1,81 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""Scoring defaults stored as ``scoring-defaults.json`` in a model bundle.""" + +import json +import warnings +from dataclasses import asdict, dataclass +from numbers import Integral, Real +from pathlib import Path +from typing import Any, Mapping, Optional + +from kumo.engine.trainer._spec import SCORING_DEFAULTS_ROLE, listed_file + +SCORING_DEFAULTS_FILE = 'scoring-defaults.json' + + +@dataclass(frozen=True) +class ScoringDefaults: + """Prediction arguments a bundle applies when ``predict`` omits them.""" + + binary_classification_threshold: Optional[float] = None + num_classes_to_return: Optional[int] = None + + +def normalized_scoring_defaults(defaults: ScoringDefaults, + task_type: Any) -> ScoringDefaults: + from kumo_core.prediction_options import ( + resolve_classification_threshold, + resolve_num_classes_to_return, + ) + + threshold = defaults.binary_classification_threshold + count = defaults.num_classes_to_return + if threshold is not None: + if isinstance(threshold, bool) or not isinstance(threshold, Real): + raise ValueError('binary_classification_threshold must be a ' + f'number, but got {threshold!r}') + threshold = float(threshold) + if resolve_classification_threshold(task_type, threshold) is None: + raise ValueError( + 'binary_classification_threshold applies only to binary and ' + f'multilabel classification, not {task_type.value}') + if isinstance(count, Integral) and not isinstance(count, bool): + count = int(count) + if count is not None: + resolve_num_classes_to_return(task_type, count) + return ScoringDefaults(threshold, count) + + +def write_scoring_defaults(bundle_dir: Path, defaults: ScoringDefaults) -> None: + from kumo_core.artifacts import add_bundle_file, remove_bundle_file + from kumo_core.io.filesystem import atomic_write_bytes + + if defaults == ScoringDefaults(): + remove_bundle_file(bundle_dir, SCORING_DEFAULTS_ROLE) + return + atomic_write_bytes( + bundle_dir / SCORING_DEFAULTS_FILE, + (json.dumps(asdict(defaults), indent=2, sort_keys=True) + + '\n').encode('utf-8')) + add_bundle_file(bundle_dir, SCORING_DEFAULTS_ROLE, SCORING_DEFAULTS_FILE) + + +def read_scoring_defaults(bundle_dir: Path, + artifact_paths: Mapping[str, str]) -> ScoringDefaults: + from kumo_core.artifacts import bundle_task_type + + path = listed_file(bundle_dir, SCORING_DEFAULTS_ROLE) + if path is None: + return ScoringDefaults() + try: + data = json.loads(path.read_text(encoding='utf-8')) + return normalized_scoring_defaults( + ScoringDefaults(data.get('binary_classification_threshold'), + data.get('num_classes_to_return')), + bundle_task_type(artifact_paths)) + except Exception as error: + warnings.warn(f'Ignoring the scoring defaults in {bundle_dir}: {error}', + UserWarning, stacklevel=3) + return ScoringDefaults() diff --git a/packages/kumo-engine/kumo/engine/trainer/_spec.py b/packages/kumo-engine/kumo/engine/trainer/_spec.py index 2c755fab4..757390902 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_spec.py +++ b/packages/kumo-engine/kumo/engine/trainer/_spec.py @@ -26,6 +26,7 @@ MANIFEST_FILE = 'publication-manifest.json' SPEC_ROLE = 'training_spec' PROVENANCE_ROLE = 'provenance' +SCORING_DEFAULTS_ROLE = 'scoring_defaults' @dataclass(frozen=True) @@ -140,5 +141,6 @@ def read_bundle_artifact_paths( return { entry['role']: str(bundle_dir / entry['relative_path']) for entry in artifacts - if entry['role'] not in (SPEC_ROLE, PROVENANCE_ROLE) + if entry['role'] not in (SPEC_ROLE, PROVENANCE_ROLE, + SCORING_DEFAULTS_ROLE) }, bundle_id diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 6efd488c4..fc35048c0 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -8,6 +8,8 @@ import hashlib import json import logging +import os +import shutil import time import warnings from dataclasses import dataclass, field, replace @@ -64,6 +66,7 @@ from kumo.engine.graph import GraphSnapshot from kumo.engine.pquery.prediction_table import PredictionTable from kumo.engine.pquery.training_table import TrainingTable +from kumo.engine.trainer._scoring import ScoringDefaults, read_scoring_defaults from kumo.engine.trainer._spec import ( read_bundle_artifact_paths, read_training_spec, @@ -99,6 +102,7 @@ class TrainingResult: holdout_path: Optional[Path] = None holdout_truncated: Optional[bool] = False bundle_id: Optional[str] = None + scoring_defaults: ScoringDefaults = field(default_factory=ScoringDefaults) _trainer: Optional['Trainer'] = field(default=None, repr=False, compare=False) @@ -121,6 +125,39 @@ def predict(self, *args: Any, **kwargs: Any) -> 'PredictionResult': 'use Trainer.load on its bundle_dir.') return self._trainer.predict(*args, **kwargs) + def with_scoring_defaults(self, + defaults: ScoringDefaults) -> 'TrainingResult': + """A copy of this model whose ``predict`` applies ``defaults`` when a + call omits them; a new bundle, whose manifest names this one in + ``parents.derived_from``.""" + from kumo_core.artifacts import bundle_task_type + + from kumo.engine.trainer._provenance import record_derivation + from kumo.engine.trainer._scoring import ( + normalized_scoring_defaults, + write_scoring_defaults, + ) + + if self._trainer is None: + raise RuntimeError('This result has no model; use Trainer.load on ' + 'its bundle_dir.') + defaults = normalized_scoring_defaults( + defaults, bundle_task_type(self.artifact_paths)) + bundle_dir = trainer_scratch_dir(new_run_id()) / 'derived-model' + shutil.copytree(self.bundle_dir, bundle_dir, + copy_function=_link_or_copy) + write_scoring_defaults(bundle_dir, defaults) + try: + record_derivation(bundle_dir, self.bundle_dir, self.bundle_id) + except Exception as error: + warnings.warn( + f"Could not record the parent model in _provenance.json: " + f"{error}", UserWarning, stacklevel=2) + result = Trainer._load_bundle(bundle_dir, + self._trainer._embedding_models).result + assert result is not None + return result + def holdout_df(self) -> pd.DataFrame: r"""Read the holdout predictions as a pandas DataFrame. @@ -280,6 +317,38 @@ def _to_core_model_plan(model_plan: ModelPlan) -> CoreModelPlan: return internal_model_plan_from_public(model_plan) +def _link_or_copy(source: str, destination: str) -> None: + try: + os.link(source, destination) + except OSError: + shutil.copy2(source, destination) + + +def _require_publisher(export_config: ModelOutputConfig) -> None: + if not callable( + getattr(export_config.connector.sink_backend, 'publish_model', + None)): + raise NotImplementedError( + f'{type(export_config.connector).__name__} cannot publish ' + f'models') + + +def _publish(training_result: TrainingResult, + export_config: ModelOutputConfig) -> PublishedModel: + spec = read_training_spec(training_result.bundle_dir) + return export_config.connector.sink_backend.publish_model( + export_config.connector, + str(training_result.bundle_dir), + path=export_config.path, + location=export_config.location, + metadata={ + name: str(getattr(spec, name)) + for name in ('spec_version', 'predictive_query', 'snapshot_id') + }, + metrics=training_result.metrics, + ) + + def _stage_model_bundle(artifacts: Mapping[str, str], dest: Path) -> Path: from kumo_core.artifacts import stage_model_bundle return stage_model_bundle(artifacts, dest) @@ -379,6 +448,19 @@ def manifest(self) -> Optional[ModelManifest]: result = self._training_result return None if result is None else result.manifest + def _fitted(self, method: str) -> TrainingResult: + if self._training_result is None: + raise RuntimeError( + f"Trainer.{method}() requires a prior Trainer.fit() call on " + "this same Trainer instance.") + return self._training_result + + def export(self, config: ModelOutputConfig) -> PublishedModel: + """Publish the fitted or loaded model bundle.""" + training_result = self._fitted('export') + _require_publisher(config) + return _publish(training_result, config) + @_without_torch_scatter_hint def fit( self, @@ -427,12 +509,8 @@ def fit( Each call returns a new `TrainingResult`; `Trainer.result` holds the latest. """ - if export_config is not None and not callable( - getattr(export_config.connector.sink_backend, 'publish_model', - None)): - raise NotImplementedError( - f'{type(export_config.connector).__name__} cannot publish ' - f'models') + if export_config is not None: + _require_publisher(export_config) from kumo.engine._mlops import resolve settings = resolve(mlops) @@ -539,11 +617,11 @@ def _fit( search=AutoMLSearchOptions(), ) if warm_start_artifacts is not None: - source_schema = json.loads( - Path(warm_start_artifacts['prediction_schema']).read_text()) + from kumo_core.artifacts import bundle_task_type + validate_warm_start_task_type( TaskType(workflow.query.task_type), - TaskType(source_schema['task_type']), + bundle_task_type(warm_start_artifacts), ) from kumo_core.pquery import ( query_is_temporal, @@ -727,21 +805,9 @@ def _finish_fit( 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( - export_config.connector, - str(training_result.bundle_dir), - path=export_config.path, - location=export_config.location, - metadata={ - name: str(getattr(spec, name)) - for name in ('spec_version', 'predictive_query', - 'snapshot_id') - }, - metrics=training_result.metrics, - ) - training_result = replace(training_result, - published_model=published) + training_result = replace( + training_result, + published_model=_publish(training_result, export_config)) if settings is not None and settings.wants('model'): _store_model(training_result, settings) @@ -795,6 +861,8 @@ def predict( Binary classification requires an explicit threshold. Multilabel classification defaults to 0.5. Both accept thresholds from 0 to 1; thresholds change predicted labels without changing class scores. + + Omitted arguments fall back to the bundle's :attr:`scoring_defaults`. """ from kumo_core.prediction_options import ( validate_num_classes_to_return, @@ -854,10 +922,10 @@ def _predict( num_workers: int = 1, binary_classification_threshold: Optional[float] = None, ) -> PredictionResult: - if self._training_result is None: - raise RuntimeError( - "Trainer.predict() requires a prior Trainer.fit() call on " - "this same Trainer instance.") + training_result = self._fitted('predict') + if binary_classification_threshold is None: + binary_classification_threshold = ( + training_result.scoring_defaults.binary_classification_threshold) from kumo_core.pquery import normalize_graph, parse_and_validate_query from kumo_core.prediction_options import ( prediction_query_contract, @@ -891,6 +959,8 @@ def _predict( # population and anchors still come from the immutable PredictionTable. native_graph = normalize_graph(graph_config) parsed = parse_and_validate_query(prediction_table.query, native_graph) + if num_classes_to_return is None and parsed.query.top_k is None: + num_classes_to_return = training_result.scoring_defaults.num_classes_to_return num_classes_to_return = resolve_num_classes_to_return( task_type, num_classes_to_return, @@ -917,7 +987,7 @@ def _predict( from kumo_core.io.json import load_json # Loaded once: the embedding check and the in-process scorer share it. - encoders_path = self._training_result.artifact_paths.get('encoders') + encoders_path = training_result.artifact_paths.get('encoders') encoders = (load_json(encoders_path, EncodersConfig) if encoders_path and Path(encoders_path).is_file() else None) require_trained_embedding_models(self._embedding_models, encoders) @@ -981,7 +1051,7 @@ def _predict( model_dir = str( _stage_model_bundle( - self._training_result.artifact_paths, + training_result.artifact_paths, run_dir / 'model_bundle', )) @@ -1098,5 +1168,7 @@ def _load_bundle( spec.training_table_generation, TrainingTableGenerationPlan) trainer._training_result = TrainingResult(metrics, paths, bundle_dir, bundle_id=bundle_id, + scoring_defaults=read_scoring_defaults( + bundle_dir, paths), _trainer=trainer) return trainer diff --git a/packages/kumo-engine/test/test_prediction_completion.py b/packages/kumo-engine/test/test_prediction_completion.py index 526d87a3e..255ba9a7a 100644 --- a/packages/kumo-engine/test/test_prediction_completion.py +++ b/packages/kumo-engine/test/test_prediction_completion.py @@ -13,6 +13,7 @@ from kumoapi.model_plan import ModelPlan from kumo.engine.pquery.prediction_table import PredictionTable +from kumo.engine.trainer._scoring import ScoringDefaults from kumo.engine.trainer.trainer import Trainer @@ -33,7 +34,8 @@ def prepared_prediction(snapshot, tmp_path, monkeypatch): from kumo_core.types import TaskType trainer = Trainer(ModelPlan()) - trainer._training_result = SimpleNamespace(artifact_paths={}) + trainer._training_result = SimpleNamespace( + artifact_paths={}, scoring_defaults=ScoringDefaults()) trainer._snapshot_id = snapshot.id trainer._training_query = 'PREDICT users.label FOR EACH users.id' table = PredictionTable(trainer._training_query, tmp_path, tmp_path, 8, diff --git a/packages/kumo-engine/test/test_prediction_workers.py b/packages/kumo-engine/test/test_prediction_workers.py index 2b9e7b7c7..d06ee9c92 100644 --- a/packages/kumo-engine/test/test_prediction_workers.py +++ b/packages/kumo-engine/test/test_prediction_workers.py @@ -14,6 +14,7 @@ from kumo_core.types.enums import PredictionArtifactType from kumoapi.model_plan import ModelPlan +from kumo.engine.trainer._scoring import ScoringDefaults from kumo.engine.trainer.trainer import Trainer @@ -80,7 +81,8 @@ def test_bad_class_count_fails_before_prediction_work(count, monkeypatch): def test_valid_class_count_reaches_scorer(task_type, count, tmp_path, monkeypatch): trainer = Trainer(ModelPlan(), accelerator='cpu') - trainer._training_result = SimpleNamespace(artifact_paths={}) + trainer._training_result = SimpleNamespace( + artifact_paths={}, scoring_defaults=ScoringDefaults()) trainer._training_query = 'trained query' monkeypatch.setattr(Trainer, 'scratch_dir', property(lambda self: tmp_path)) monkeypatch.setattr( diff --git a/packages/kumo-engine/test/test_scoring_defaults.py b/packages/kumo-engine/test/test_scoring_defaults.py new file mode 100644 index 000000000..497204d17 --- /dev/null +++ b/packages/kumo-engine/test/test_scoring_defaults.py @@ -0,0 +1,125 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import hashlib +import shutil + +import numpy as np +import pandas as pd +import pytest +from kumo_core.artifacts import add_bundle_file, remove_bundle_file +from kumo_core.types import TaskType + +import kumo.engine as kumo +from kumo.engine.trainer._scoring import normalized_scoring_defaults + + +def _frame(model, table, snapshot, **kwargs): + result = model.predict(table, snapshot=snapshot, **kwargs) + return result.frame.sort_values('ENTITY').reset_index(drop=True) + + +def _digest(directory): + return { + str(path.relative_to(directory)): + hashlib.sha256(path.read_bytes()).hexdigest() + for path in sorted(directory.rglob('*')) if path.is_file() + } + + +def _copy(trainer, destination): + shutil.copytree(trainer.result.bundle_dir, destination) + return destination + + +def test_derived_model_thresholds_unless_explicit(fitted): + trainer, table, snapshot = fitted + parent = trainer.result + original = _digest(parent.bundle_dir) + probs = _frame(parent, table, snapshot, + binary_classification_threshold=.5)['True_PROB'] + threshold = float(probs.median()) + child = parent.with_scoring_defaults( + kumo.ScoringDefaults(binary_classification_threshold=threshold)) + + implicit = _frame(child, table, snapshot) + pd.testing.assert_frame_equal( + implicit, + _frame(parent, table, snapshot, + binary_classification_threshold=threshold)) + assert implicit['TARGET_PRED'].any() + assert not _frame(child, table, snapshot, + binary_classification_threshold=1.0)['TARGET_PRED'].any() + + assert _digest(parent.bundle_dir) == original + assert parent.scoring_defaults == kumo.ScoringDefaults() + assert child.id != parent.id + assert child.manifest.id == child.id + assert child.manifest.parents.training_table is None + derived_from = child.manifest.parents.derived_from + assert derived_from.id == parent.id + assert derived_from.parents.training_table == ( + parent.manifest.parents.training_table) + cleared = child.with_scoring_defaults(kumo.ScoringDefaults()) + assert cleared.id == parent.id + assert cleared.manifest.parents.derived_from.parents.derived_from == ( + kumo.ModelManifest(id=parent.id, reconstructed=True)) + + +def test_equal_defaults_give_equal_bundle_ids(fitted): + trainer, *_ = fitted + children = [ + trainer.result.with_scoring_defaults( + kumo.ScoringDefaults(binary_classification_threshold=value)) + for value in (1, np.float32(1)) + ] + assert children[0].id == children[1].id + assert children[0].scoring_defaults.binary_classification_threshold == 1.0 + + +def test_parent_without_a_manifest_reads_as_a_stub(fitted, tmp_path): + trainer, *_ = fitted + old = _copy(trainer, tmp_path / 'old') + remove_bundle_file(old, 'provenance') + child = kumo.Trainer.load(old).result.with_scoring_defaults( + kumo.ScoringDefaults(binary_classification_threshold=.5)) + assert child.manifest.parents.derived_from == kumo.ModelManifest( + id=trainer.bundle_id, reconstructed=True) + + +def test_defaults_survive_export_and_connector_load(fitted, tmp_path): + trainer, *_ = fitted + child = trainer.result.with_scoring_defaults( + kumo.ScoringDefaults(binary_classification_threshold=.4)) + registry = kumo.FileConnector(root_dir=str(tmp_path)) + child._trainer.export(kumo.ModelOutputConfig(registry, path='m1')) + loaded = kumo.Trainer.load_from_connector(registry, path='m1') + assert loaded.result.scoring_defaults == child.scoring_defaults + assert loaded.bundle_id == child.id + + +def test_defaults_are_validated_against_the_task(fitted): + trainer, *_ = fitted + with pytest.raises(ValueError, match='num_classes_to_return'): + trainer.result.with_scoring_defaults( + kumo.ScoringDefaults(num_classes_to_return=2)) + assert normalized_scoring_defaults( + kumo.ScoringDefaults(.3, np.int64(2)), + TaskType.MULTILABEL_CLASSIFICATION) == kumo.ScoringDefaults(.3, 2) + + +def test_unlisted_or_unreadable_defaults_are_ignored(fitted, tmp_path): + trainer, table, snapshot = fitted + bundle = _copy(trainer, tmp_path / 'old') + (bundle / 'scoring-defaults.json').write_text('not json') + old = kumo.Trainer.load(bundle) + assert old.result.scoring_defaults == kumo.ScoringDefaults() + assert old.bundle_id == trainer.bundle_id + + add_bundle_file(bundle, 'scoring_defaults', 'scoring-defaults.json') + with pytest.warns(UserWarning, match='Ignoring the scoring defaults'): + ignored = kumo.Trainer.load(bundle) + assert ignored.result.scoring_defaults == kumo.ScoringDefaults() + with pytest.raises(ValueError, match='must be specified'): + ignored.predict(table, snapshot=snapshot) diff --git a/packages/kumo-engine/test/test_semantic_carry_forward.py b/packages/kumo-engine/test/test_semantic_carry_forward.py index 9481d14d1..64c6d66e5 100644 --- a/packages/kumo-engine/test/test_semantic_carry_forward.py +++ b/packages/kumo-engine/test/test_semantic_carry_forward.py @@ -24,6 +24,7 @@ PredictionTable, _compile_prediction_table, ) +from kumo.engine.trainer._scoring import ScoringDefaults from kumo.engine.trainer.trainer import Trainer, _stage_model_bundle @@ -134,7 +135,7 @@ def test_prediction_admission_before_materialization(snapshot, tmp_path, query, monkeypatch): from kumoapi.model_plan import ModelPlan trainer = Trainer(ModelPlan()) - trainer._training_result = SimpleNamespace() + trainer._training_result = SimpleNamespace(scoring_defaults=ScoringDefaults()) trainer._snapshot_id = snapshot.id trainer._training_query = 'PREDICT users.label FOR EACH users.id' materialize = Mock(side_effect=AssertionError('must not materialize')) @@ -162,7 +163,8 @@ def test_binary_predictions_preserve_probabilities_and_metadata_contract( _stub_generated_task_dataset_builder(monkeypatch) monkeypatch.setattr(type(trainer), 'scratch_dir', property(lambda _: tmp_path)) - trainer._training_result = SimpleNamespace(artifact_paths={}) + trainer._training_result = SimpleNamespace( + artifact_paths={}, scoring_defaults=ScoringDefaults()) trainer._snapshot_id = snapshot.id trainer._training_query = 'PREDICT users.label FOR EACH users.id' table = PredictionTable(trainer._training_query, tmp_path, tmp_path, 0, @@ -271,7 +273,8 @@ def stats(index_path, table_name, mode): trainer = Trainer(ModelPlan()) monkeypatch.setattr(type(trainer), 'scratch_dir', property(lambda _: tmp_path)) - trainer._training_result = SimpleNamespace(artifact_paths={}) + trainer._training_result = SimpleNamespace( + artifact_paths={}, scoring_defaults=ScoringDefaults()) trainer._snapshot_id = snapshot.id trainer._training_query = query @@ -706,7 +709,8 @@ def run_predict() -> None: trainer = Trainer(ModelPlan()) monkeypatch.setattr(type(trainer), 'scratch_dir', property(lambda _: tmp_path)) - trainer._training_result = SimpleNamespace(artifact_paths={}) + trainer._training_result = SimpleNamespace( + artifact_paths={}, scoring_defaults=ScoringDefaults()) trainer._snapshot_id = snapshot.id trainer._training_query = query trainer.predict(table, snapshot=snapshot, output_config=output, @@ -718,7 +722,8 @@ def run_predict() -> None: same_trainer = Trainer(ModelPlan()) monkeypatch.setattr(type(same_trainer), 'scratch_dir', property(lambda _: tmp_path)) - same_trainer._training_result = SimpleNamespace(artifact_paths={}) + same_trainer._training_result = SimpleNamespace( + artifact_paths={}, scoring_defaults=ScoringDefaults()) same_trainer._snapshot_id = snapshot.id same_trainer._training_query = query same_trainer.predict(table, snapshot=snapshot, output_config=output,