Log prediction runs to MLflow and Snowflake through the tracking adapters

packages/kumo-engine/docs/mlops.md

@@ -203,23 +203,34 @@ Runs carry pointers, not lineage:
 
 | tag | on | value |
 |---|---|---|
-| `kumo.kind` | every run | `fit` or `trial` |
-| `kumo.run_id` | every run | the kumo fit run id (trials: `<fit run id>-trial-<n>-attempt-<m>`) |
-| `kumo.bundle_id` | a finished fit run | the model's bundle id (`models/<bundle id>/` under the artifact root) |
-| `kumo.parent_run`, `kumo.status`, `kumo.error` | trial runs; `kumo.error` also on a failed fit | outcome |
+| `kumo.kind` | every run | `fit`, `trial`, `prediction` or `monitoring` |
+| `kumo.run_id` | every run | the kumo run id (trials: `<fit run id>-trial-<n>-attempt-<m>`; predictions: the prediction id) |
+| `kumo.bundle_id` | a finished fit run, prediction and monitoring runs | the model's bundle id (`models/<bundle id>/` under the artifact root) |
+| `kumo.prediction_run` | monitoring runs | the prediction id it monitors |
+| `kumo.parent_run`, `kumo.status`, `kumo.error` | trial runs; `kumo.error` also on a failed fit or prediction | outcome |
 
 The model's `_provenance.json` records the tracked fit as `tracking.run_id`
 (the `kumo.run_id` tag) and is logged on the fit run with `query.pql` and
 `model_plan.json`.
 
+With tracking on, `predict` also records each prediction run as one run named
+by its prediction id, after its `output_config` export and its artifact-root
+copy: scoring options as params; row counts and score statistics (for example
+`predictions.True_PROB.q50`) as metrics; the scoring snapshot's source tables as
+inputs; and its `_provenance.json` as an artifact. If the export raises, the run
+is still recorded, ended as failed with `kumo.error`, and `predict` raises the
+export's error. It runs on a tracking thread for up to 30 s; tracking errors
+never change the prediction. MLflow nests it under an active run, like a fit.
+
 Other systems plug in through the `kumo.engine.tracking.Tracking` protocol,
 which receives ready-built `RunRecord`s so every tracker records the same
 names: `start(record)` with the fit run, `on_trial(record, event)` with a trial
 attempt's run after each of its events, and `finish(record, error)` with the
 fit run's final record (`None` when tracking stopped early). A tracker that
 raises is not called again for that fit and its run is ended as failed.
 Optional members: `prepare()`, `tags`, `log_model(result, snapshot, registry)` (see MLflow models below), and the best-effort hooks
-`log_prediction_run` and `log_monitoring`. Params are the resolved model plan
+`log_prediction_run(record)` and `log_monitoring(record)` (a record tagged
+`kumo.error` is a failed run). Params are the resolved model plan
 (fit) or searched hyperparameters and `attempt` (trial). `RunRecord.inputs`
 maps each source table to its `connector`, `location`, `version` and
 `version_kind`. A tracker's `tags=` are added as given, except `kumo.*` names

packages/kumo-engine/kumo/engine/_tracking.py

@@ -34,6 +34,7 @@
 if TYPE_CHECKING:
     from kumoapi.plugin import ModelRegistryLocation, PublishedModel
 
+    from kumo.engine._manifest import PredictionManifest
     from kumo.engine.graph import GraphSnapshot
     from kumo.engine.trainer.trainer import TrainingResult
 
@@ -81,14 +82,17 @@ class RunRecord:
 
 
 def fit_inputs(manifest: Any) -> Dict[str, Dict[str, Optional[str]]]:
-    r"""Source tables of a training table or model manifest, read from its
-    embedded parents; stub parents add none."""
+    r"""Source tables of a training table or model manifest, or a prediction
+    run's scoring snapshot, read from its embedded parents; stub parents add
+    none."""
     try:
-        table = manifest
-        if getattr(manifest, 'kind', None) == 'model':
+        table, kind = manifest, getattr(manifest, 'kind', None)
+        if kind == 'model':
             table = getattr(_parents(manifest), 'training_table', None)
-        parents = _parents(table)
-        snapshot = _parents(getattr(parents, 'snapshot', None))
+        parents = None if kind == 'prediction' else _parents(table)
+        snapshot = _parents(
+            getattr(_parents(manifest) if kind == 'prediction' else parents,
+                    'snapshot', None))
         tables = dict(getattr(snapshot, 'source_tables', None) or {})
         source = getattr(parents, 'source_table', None)
         if source is not None:
@@ -215,6 +219,88 @@ def trial_record(context: FitContext, events: Sequence[Event]) -> RunRecord:
     )
 
 
+def prediction_record(manifest: 'PredictionManifest',
+                      error: Optional[BaseException] = None) -> RunRecord:
+    r"""A prediction run's record, from its manifest: scoring options as
+    params, row counts and score percentiles as metrics, the scoring
+    snapshot's source tables as inputs, pointer tags (``kumo.error`` when
+    ``error`` ended the run), and the manifest as an artifact."""
+    from kumo_core.instrumentation import format_error
+
+    table = manifest.prediction_table
+    params = {**manifest.scoring.to_dict(), 'anchors': table.anchors}
+    metrics: Dict[str, Any] = {'prediction_table.rows': table.rows}
+    for artifact, stats in (manifest.stats or {}).items():
+        flat = stats.to_dict()
+        flat.update(flat.pop('columns', None) or {})
+        for key, value in _flatten(flat, f'{artifact}.'):
+            if key.endswith('.percentiles'):
+                metrics.update(_logged_percentiles(key, value))
+            else:
+                metrics[key] = value
+    tags = {
+        'kumo.kind': 'prediction',
+        'kumo.run_id': manifest.id,
+        'kumo.bundle_id': manifest.parents.model.id,
+        'kumo.error': None if error is None else format_error(error),
+    }
+    return RunRecord(
+        name=manifest.id,
+        params={
+            key: _param_value(value)
+            for key, value in params.items() if value is not None
+        },
+        metrics=_finite_floats({
+            key: value
+            for key, value in metrics.items()
+            if isinstance(value, (int, float)) and not isinstance(value, bool)
+        }),
+        tags={
+            key: str(value)
+            for key, value in tags.items() if value is not None
+        },
+        artifacts={
+            PROVENANCE_FILE:
+            json.dumps(manifest.to_dict(), indent=2, sort_keys=True,
+                       default=str)
+        },
+        inputs=fit_inputs(manifest),
+    )
+
+
+def record_failure(record: RunRecord) -> Optional[Exception]:
+    r"""The error a run record says ended its run (its ``kumo.error``)."""
+    error = record.tags.get('kumo.error')
+    return None if error is None else RuntimeError(error)
+
+
+def track_prediction(manifest: Optional['PredictionManifest'], tracking: Any,
+                     error: Optional[BaseException] = None) -> None:
+    r"""Record a prediction run in ``tracking``, failed if ``error`` ended
+    it; never raises."""
+    if getattr(tracking, 'log_prediction_run', None) is None:
+        return
+    try:
+        if manifest is None:
+            raise ValueError('the run has no prediction manifest')
+        record = prediction_record(manifest, error)
+    except Exception:
+        logger.warning('Experiment tracking failed; the prediction run is '
+                       'not tracked', exc_info=True)
+        return
+    log_run(tracking, 'log_prediction_run', record)
+
+
+def _logged_percentiles(key: str, values: Any) -> Dict[str, Any]:
+    if not isinstance(values, list) or len(values) != 99:
+        return {}
+    column = key[:-len('percentiles')]
+    return {
+        f'{column}q{q:02d}': values[q - 1]
+        for q in (1, 5, 25, 50, 75, 95, 99)
+    }
+
+
 class Tracking(Protocol):
     r"""Records one `Trainer.fit` in an experiment tracker.
 

packages/kumo-engine/kumo/engine/integrations/mlflow.py

@@ -1,8 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES.
 # All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
-r"""MLflow tracking and models for `Trainer.fit`
-(``pip install 'kumo-engine[mlflow]'``)."""
+r"""MLflow tracking and models for `Trainer.fit`, `Trainer.predict` and
+monitoring (``pip install 'kumo-engine[mlflow]'``)."""
 
 import hashlib
 import json
@@ -30,7 +30,7 @@
 from kumoapi.plugin import ModelRegistryLocation, PublishedModel
 
 from kumo.engine._events import Event
-from kumo.engine._tracking import DROPPED_PARAMS, RunRecord
+from kumo.engine._tracking import DROPPED_PARAMS, RunRecord, record_failure
 from kumo.engine._version import __version__
 from kumo.engine.connector._extras import import_provider
 
@@ -63,6 +63,9 @@ class MlflowTracking:
     Catalog then links them to a model registered from the run), file and
     object-store sources by URI. Params MLflow rejects are counted in
     ``kumo.dropped_params``. Concurrent fits in one process are not supported.
+    As ``kumo.MLOps(tracking=...)``, it also records each prediction run and
+    each monitoring call as one run named after it, placed like a fit run;
+    ``run_name`` and ``tags`` apply to fit runs only.
 
     Args:
         experiment_name: experiment for the fit run, created if missing.
@@ -78,7 +81,8 @@ class MlflowTracking:
         serving_snapshot: exported snapshot URI the served model loads at
             startup (see `log_model`).
 
-    ``run_id`` is the MLflow id of the latest fit run, once it has started.
+    ``run_id`` is the MLflow id of the most recent tracked run, once it has
+    started.
     """
 
     def __init__(
@@ -153,6 +157,10 @@ def start(self, record: RunRecord) -> None:
         self._log_artifacts(record.artifacts)
         _log_inputs(mlflow, self._client, self._parent, record.inputs)
 
+    def log_prediction_run(self, record: RunRecord) -> None:
+        self.start(record)
+        self.finish(record, record_failure(record))
+
     def log_monitoring(self, record: RunRecord) -> None:
         self.start(record)
         experiment_id, run_id = self._experiment_id, self.run_id

packages/kumo-engine/kumo/engine/integrations/snowflake.py

@@ -1,8 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES.
 # All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
-r"""Snowflake Experiment Tracking for `Trainer.fit`
-(``pip install 'kumo-engine[snowflake]'``)."""
+r"""Snowflake Experiment Tracking for `Trainer.fit`, `Trainer.predict` and
+monitoring (``pip install 'kumo-engine[snowflake]'``)."""
 
 import logging
 import re
@@ -14,7 +14,12 @@
 from kumo_core.sources.sql import resolve_snowflake_identifier
 
 from kumo.engine._events import Event
-from kumo.engine._tracking import DROPPED_METRICS, DROPPED_PARAMS, RunRecord
+from kumo.engine._tracking import (
+    DROPPED_METRICS,
+    DROPPED_PARAMS,
+    RunRecord,
+    record_failure,
+)
 from kumo.engine.connector._extras import import_provider
 
 logger = logging.getLogger(__name__)
@@ -32,7 +37,9 @@ class SnowflakeTracking:
     trial attempts are logged in it under ``trial-<n>-attempt-<m>/`` names and
     tags are logged as params. Snowflake has no dataset inputs, so a fit's
     source tables are not recorded as inputs. The Snowpark session's
-    parameters are not changed.
+    parameters are not changed. As ``kumo.MLOps(tracking=...)``, it also
+    records each prediction run and each monitoring call as one run named
+    after it; ``run_name`` and ``tags`` apply to fit runs only.
 
     Args:
         session: Snowpark session; needs CREATE EXPERIMENT on the schema, or
@@ -47,8 +54,8 @@ class SnowflakeTracking:
         native_model_lineage: keep the run open as a context for the whole
             fit, so models registered meanwhile are linked to it.
 
-    ``run_id`` is the latest fit run's name (Snowflake runs are identified
-    by name), once it has started.
+    ``run_id`` is the most recent tracked run's name (Snowflake runs are
+    identified by name), once it has started.
     """
 
     def __init__(
@@ -116,6 +123,10 @@ def start(self, record: RunRecord) -> None:
         self._log_params({**record.params, **record.tags})
         self._log_artifacts(record.artifacts)
 
+    def log_prediction_run(self, record: RunRecord) -> None:
+        self.start(record)
+        self.finish(record, record_failure(record))
+
     def log_monitoring(self, record: RunRecord) -> None:
         self.start(record)
         self.finish(record, None)

packages/kumo-engine/kumo/engine/trainer/_run_manifest.py

@@ -51,26 +51,32 @@ def start(*, result: Any, trainer: Any, snapshot: Any, prediction_table: Any,
 def finish(result: Any, manifest: Optional[PredictionManifest],
            output_config: Any, settings: Any) -> None:
     r"""Export to ``output_config``, record the run (with the outputs
-    written, even if the export fails), then store it in the artifact
-    root."""
+    written, even if the export fails), store it in the artifact root and
+    track it."""
+    from kumo.engine._tracking import track_prediction
+
     records = [] if output_config is None else output_records(output_config)
     written: List[PredictionOutput] = []
-    status = 'failed'
+    error: Optional[BaseException] = None
     try:
         if output_config is not None:
             result._export(
                 output_config, None if manifest is None else replace(
                     manifest, outputs=records), records, written)
-        status = 'succeeded'
+    except BaseException as raised:
+        error = raised
+        raise
     finally:
         if manifest is not None:
             manifest = replace(
-                manifest, status=status,
+                manifest, status='failed' if error else 'succeeded',
                 outputs=None if records is None else written,
                 timings=replace(manifest.timings, finished_at=_now()))
             write(result, manifest)
-    if settings is not None and settings.wants('prediction'):
-        store(result, manifest, settings)
+        if settings is not None:
+            if error is None and settings.wants('prediction'):
+                store(result, manifest, settings)
+            track_prediction(manifest, settings.tracking, error)
 
 
 def store(result: Any, manifest: Optional[PredictionManifest],

packages/kumo-engine/test/test_mlflow_tracking.py

@@ -343,3 +343,30 @@ def test_a_monitoring_run_is_new_and_tags_its_prediction_run(tracking_uri):
     assert monitor.data.tags['kumo.prediction_run'] == 'pred-1'
     assert client.get_run(predicted).data.tags['kumo.monitored'] == (
         monitor.info.run_id)
+
+
+def test_prediction_run_nests_under_the_active_run(tracking_uri):
+    from dataclasses import replace
+
+    from test_tracking_contract import PREDICTION_MANIFEST
+
+    from kumo.engine._tracking import track_prediction
+    from kumo.engine.integrations.mlflow import MlflowTracking
+
+    client = mlflow.MlflowClient()
+    tracking = MlflowTracking(experiment_name='scores',
+                              log_system_metrics=False)
+    with mlflow.start_run() as outer:
+        track_prediction(PREDICTION_MANIFEST, tracking)
+        assert mlflow.active_run().info.run_id == outer.info.run_id
+    track_prediction(replace(PREDICTION_MANIFEST, id='trainer-predict-2'),
+                     tracking)
+
+    nested, top = sorted(
+        (run for run in _runs(client) if 'kumo.kind' in run.data.tags),
+        key=lambda run: 'mlflow.parentRunId' not in run.data.tags)
+    assert nested.data.tags['mlflow.parentRunId'] == outer.info.run_id
+    assert 'mlflow.parentRunId' not in top.data.tags
+    assert client.get_experiment(top.info.experiment_id).name == 'scores'
+    assert (top.info.run_name, top.info.status) == ('trainer-predict-2',
+                                                    'FINISHED')