diff --git a/packages/kumo-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index 625e99dcc..9ed7898d8 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -203,15 +203,25 @@ 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: `-trial--attempt-`) | -| `kumo.bundle_id` | a finished fit run | the model's bundle id (`models//` 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: `-trial--attempt-`; predictions: the prediction id) | +| `kumo.bundle_id` | a finished fit run, prediction and monitoring runs | the model's bundle id (`models//` 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 @@ -219,7 +229,8 @@ 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 diff --git a/packages/kumo-engine/kumo/engine/_tracking.py b/packages/kumo-engine/kumo/engine/_tracking.py index ab249a5b0..09c1cc3b5 100644 --- a/packages/kumo-engine/kumo/engine/_tracking.py +++ b/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. diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index 0b4c6ee5d..a893d375d 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/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 diff --git a/packages/kumo-engine/kumo/engine/integrations/snowflake.py b/packages/kumo-engine/kumo/engine/integrations/snowflake.py index 75ba30e71..592d94dca 100644 --- a/packages/kumo-engine/kumo/engine/integrations/snowflake.py +++ b/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--attempt-/`` 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) diff --git a/packages/kumo-engine/kumo/engine/trainer/_run_manifest.py b/packages/kumo-engine/kumo/engine/trainer/_run_manifest.py index 6a60bdb19..dd275b8ba 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_run_manifest.py +++ b/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], diff --git a/packages/kumo-engine/test/test_mlflow_tracking.py b/packages/kumo-engine/test/test_mlflow_tracking.py index 1931480d0..5be98f8f8 100644 --- a/packages/kumo-engine/test/test_mlflow_tracking.py +++ b/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') diff --git a/packages/kumo-engine/test/test_prediction_manifest.py b/packages/kumo-engine/test/test_prediction_manifest.py index acd70cdd8..2a418f2fe 100644 --- a/packages/kumo-engine/test/test_prediction_manifest.py +++ b/packages/kumo-engine/test/test_prediction_manifest.py @@ -311,8 +311,27 @@ def fail(*args, **kwargs): bare_exported) -def test_failed_export_still_writes_manifest(graph, classifier, tmp_path, - monkeypatch): +class _Recorder: + def __init__(self, raises=False): + self.records, self.raises = [], raises + + def start(self, record): + pass + + def on_trial(self, record, event): + pass + + def finish(self, record, error): + pass + + def log_prediction_run(self, record): + if self.raises: + raise RuntimeError('tracker down') + self.records.append(record) + + +def test_failed_export_still_writes_and_tracks_the_run(graph, classifier, + tmp_path, monkeypatch): _, snapshot = graph trainer, table = classifier connector = kumo.FileConnector(root_dir=str(tmp_path)) @@ -328,9 +347,11 @@ def fail_second(self, connector, name, *args, **kwargs): output = PredictionOutputConfig( connector=connector, output_paths={'predictions': 'preds/', 'embeddings': 'embs/'}) + recorder = _Recorder() + cfg = kumo.MLOps(artifact_root=str(tmp_path / 'root'), tracking=recorder) with pytest.raises(OSError, match='disk full'): trainer.predict(table, snapshot=snapshot, output_config=output, - binary_classification_threshold=0.5) + binary_classification_threshold=0.5, mlops=cfg) manifest, = [ json.loads(path.read_text()) for path in trainer.scratch_dir.glob( @@ -343,3 +364,32 @@ def fail_second(self, connector, name, *args, **kwargs): 'location': str(tmp_path / 'preds')}, }] assert manifest['timings']['finished_at'] is not None + assert not (tmp_path / 'root' / 'predictions').exists() + record, = recorder.records + assert record.tags['kumo.error'] == 'OSError: disk full' + assert record.tags['kumo.run_id'] == manifest['id'] + + +def test_tracking_records_the_run_and_never_changes_it(graph, classifier, + tmp_path, caplog): + from kumo.engine._tracking import prediction_record + + _, snapshot = graph + trainer, table = classifier + recorder = _Recorder() + + def predict(tracking): + cfg = kumo.MLOps(artifact_root=str(tmp_path / 'root'), + tracking=tracking, export=set()) + return trainer.predict(table, snapshot=snapshot, mlops=cfg, + binary_classification_threshold=0.4) + + untracked = trainer.predict(table, snapshot=snapshot, + binary_classification_threshold=0.4) + tracked = predict(recorder) + failed = predict(_Recorder(raises=True)) + + for result in (tracked, failed): + pd.testing.assert_frame_equal(result.frame, untracked.frame) + assert recorder.records == [prediction_record(tracked.manifest)] + assert 'tracker down' in caplog.text diff --git a/packages/kumo-engine/test/test_snowflake_tracking.py b/packages/kumo-engine/test/test_snowflake_tracking.py index 25a1ce057..17a29fa80 100644 --- a/packages/kumo-engine/test/test_snowflake_tracking.py +++ b/packages/kumo-engine/test/test_snowflake_tracking.py @@ -13,11 +13,13 @@ from test_tracking_contract import ( CONTEXT, EVENTS, + PREDICTION_MANIFEST, + PREDICTION_RUN, RESULT, join_tracking_threads, ) -from kumo.engine._tracking import TrackingSession +from kumo.engine._tracking import TrackingSession, prediction_record FAKED = ('__init__', 'set_experiment', 'start_run', 'end_run', 'log_params', 'log_metrics', 'log_artifact') @@ -117,10 +119,12 @@ def test_an_instance_in_use_is_left_alone(fake, busy, caplog): calls = len(user.calls) _replay(_tracking(session)) + _tracking(session).log_prediction_run( + prediction_record(PREDICTION_MANIFEST)) assert user.calls[calls:] == [] assert user.warnings_as_errors is True - assert caplog.text.count('not tracked in Snowflake') == 1 + assert caplog.text.count('not tracked in Snowflake') == 2 def test_location_defaults_to_the_session_not_the_last_fit(fake): @@ -214,3 +218,17 @@ def test_a_hung_snowflake_call_does_not_stall_the_fit(fake, monkeypatch, assert 'did not finish within' in caplog.text run, = fake._instance.runs.values() assert run['status'] == 'FINISHED' + + +def test_prediction_runs_are_never_resumed(fake, caplog): + tracking = _tracking() + record = prediction_record(PREDICTION_MANIFEST) + tracking.log_prediction_run(record) + tracking.log_prediction_run(record) + + assert {name: run['status'] + for name, run in fake._instance.runs.items()} == { + PREDICTION_RUN: 'FINISHED', + f'{PREDICTION_RUN}_1': 'FINISHED', + } + assert 'already exists' in caplog.text diff --git a/packages/kumo-engine/test/test_tracking_contract.py b/packages/kumo-engine/test/test_tracking_contract.py index 2efe3cb35..17c5ee5dc 100644 --- a/packages/kumo-engine/test/test_tracking_contract.py +++ b/packages/kumo-engine/test/test_tracking_contract.py @@ -1,12 +1,13 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Every tracking adapter records the same run records as a recorded fit.""" +"""Every tracking adapter records the same run records as a recorded fit +and a recorded prediction run.""" import json import threading import time -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from pathlib import Path from types import SimpleNamespace from typing import Any, Callable, Dict, List, Mapping, Optional, Tuple @@ -14,12 +15,20 @@ import pytest from kumo_core.artifacts import PROVENANCE_FILE +from kumo.engine._manifest import ( + ArtifactStats, + PredictionManifest, + PredictionParents, + stub, +) from kumo.engine._tracking import ( DROPPED_METRICS, DROPPED_PARAMS, FitContext, TrackingSession, + prediction_record, run_record, + track_prediction, ) from kumo.engine.tracking import Event, RunRecord @@ -81,6 +90,53 @@ def _epochs(trial, *values): id=BUNDLE, bundle_dir=Path('no-bundle')) +PREDICTION_RUN = 'trainer-predict-1' +PREDICTION_MANIFEST = PredictionManifest.from_dict({ + 'kind': 'prediction', + 'schema_version': 1, + 'id': PREDICTION_RUN, + 'parents': { + 'model': {'kind': 'model', 'schema_version': 1, 'id': BUNDLE, + 'reconstructed': True}, + 'snapshot': { + 'kind': 'snapshot', 'schema_version': 1, 'id': 'g@d', + 'parents': {'source_tables': INPUTS}, + }, + }, + 'created_at': '2026-10-21T00:00:00+00:00', + 'environment': None, + 'prediction_table': { + 'query_sha256': '69dab700', + 'source': 'pql', + 'anchors': ['2024-04-29T00:00:00+00:00'], + 'rows': 50, + }, + 'scoring': { + 'binary_classification_threshold': 0.4, + 'threshold_source': 'argument', + 'num_classes_to_return': None, + 'count_source': 'none', + 'num_workers': None, + }, + 'outputs': [], + 'status': 'succeeded', + 'timings': {'started_at': '2026-10-21T00:00:00+00:00', + 'finished_at': '2026-10-21T00:01:00+00:00'}, + 'stats': { + 'predictions': { + 'rows': 50, + 'columns': { + 'True_PROB': { + 'count': 50, + 'percentiles': [q / 100 for q in range(1, 100)], + } + }, + 'predicted_positive_rate': 0.62, + }, + }, +}) + + class _Recorder: def __init__(self, tracker=None): self.calls: List[Tuple[str, Any]] = [] @@ -251,7 +307,7 @@ def snowflake_adapter(monkeypatch): def observe(): observed = {} - for run in fake._instance.runs.values(): + for run_name, run in fake._instance.runs.items(): names: Dict[Optional[str], Dict[str, Dict[str, Any]]] = {} def bucket(key): @@ -274,7 +330,7 @@ def bucket(key): (step, value)) for prefix, runs in names.items(): observed[runs['tags']['kumo.run_id']] = Observed( - name=prefix, + name=prefix or run_name, artifacts=run['artifacts'] if prefix is None else {}, **runs) return observed @@ -398,3 +454,69 @@ def start(self, record): assert elapsed < 1.5 assert 'did not finish within' in caplog.text + + +def test_prediction_record_of_the_recorded_run(): + record = prediction_record(PREDICTION_MANIFEST) + + assert record.name == PREDICTION_RUN + assert record.params == { + 'anchors': '["2024-04-29T00:00:00+00:00"]', + 'binary_classification_threshold': '0.4', + 'threshold_source': 'argument', + 'count_source': 'none', + } + assert record.metrics == { + 'prediction_table.rows': 50.0, + 'predictions.rows': 50.0, + 'predictions.True_PROB.count': 50.0, + **{ + f'predictions.True_PROB.q{q:02d}': q / 100 + for q in (1, 5, 25, 50, 75, 95, 99) + }, + 'predictions.predicted_positive_rate': 0.62, + } + assert record.tags == {'kumo.kind': 'prediction', + 'kumo.run_id': PREDICTION_RUN, + 'kumo.bundle_id': BUNDLE} + assert record.inputs == INPUTS + assert json.loads(record.artifacts[PROVENANCE_FILE]) == ( + PREDICTION_MANIFEST.to_dict()) + + +def test_prediction_record_of_a_partial_manifest(): + record = prediction_record( + replace(PREDICTION_MANIFEST, + parents=PredictionParents(model=stub('model', BUNDLE), + snapshot=stub('snapshot', 'g@d')), + scoring=replace(PREDICTION_MANIFEST.scoring, + num_classes_to_return=3), + stats={'predictions': ArtifactStats(rows=150, columns={ + 'top_score': {'count': 150, 'percentiles': []}}), + 'embeddings': ArtifactStats(rows=150)}), + RuntimeError('disk full')) + + assert record.params['num_classes_to_return'] == '3' + assert record.metrics == {'prediction_table.rows': 50.0, + 'predictions.rows': 150.0, + 'predictions.top_score.count': 150.0, + 'embeddings.rows': 150.0} + assert record.inputs == {} + assert record.tags['kumo.error'] == 'RuntimeError: disk full' + + +def test_adapter_records_the_same_prediction_run(adapter): + track_prediction(PREDICTION_MANIFEST, adapter.make()) + + expected = {PREDICTION_RUN: prediction_record(PREDICTION_MANIFEST)} + assert gaps(expected, adapter.observe(), adapter) == [] + + +def test_prediction_run_without_a_manifest_or_hook_is_skipped(caplog): + recorder = _Recorder() + recorder.log_prediction_run = recorder.calls.append + track_prediction(None, recorder) + track_prediction(PREDICTION_MANIFEST, _Recorder()) + + assert recorder.calls == [] + assert 'not tracked' in caplog.text