Add TrainingResult.with_scoring_defaults and Trainer.export

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.
 

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"])

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")

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 `<root>/models/<bundle id>/`. 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:

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",

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",

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)

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()

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