Log, register and serve fitted models through MLflow pyfunc

packages/kumo-core/src/kumo_core/train/formatting.py

@@ -31,6 +31,15 @@ class LinkPredictionTable(Protocol):
 
 
 _ANCHOR_TIME_COL = "ANCHOR_TIMESTAMP"
+_CLASS_TASKS = frozenset({
+    TaskType.MULTICLASS_CLASSIFICATION,
+    TaskType.MULTILABEL_CLASSIFICATION,
+})
+_LONG_FORMAT_TASKS = _CLASS_TASKS | {
+    TaskType.MULTILABEL_RANKING,
+    TaskType.TEMPORAL_LINK_PREDICTION,
+    TaskType.STATIC_LINK_PREDICTION,
+}
 
 
 def _get_col_names(
@@ -67,6 +76,23 @@ def _get_col_names(
     return result
 
 
+def prediction_output_columns(task_type: TaskType) -> list[str]:
+    r"""The prediction columns a scored row carries after its entity column,
+    for a target named ``TARGET`` and a time column named ``TIMESTAMP``."""
+    if task_type in _CLASS_TASKS:
+        return ['PREDICTED', 'CLASS', 'SCORE']
+    if task_type.is_ranking:
+        return ['CLASS', 'SCORE']
+    if task_type == TaskType.FORECASTING:
+        return ['TIMESTAMP', 'TARGET_PRED']
+    if task_type == TaskType.BINARY_CLASSIFICATION:
+        return [
+            'TARGET_PRED',
+            *_get_col_names(pd.Index([]), task_type, is_wide_format=True)
+        ]
+    return ['TARGET_PRED']
+
+
 def _format_prediction_sanity_checks(
     target_encoder: TargetEncoder,
     task_type: TaskType,
@@ -284,13 +310,7 @@ def format_prediction(
             'metadata. Rename the time column before forecasting.')
 
     format: Literal['wide', 'long']
-    if task_type in {
-            TaskType.MULTICLASS_CLASSIFICATION,
-            TaskType.MULTILABEL_CLASSIFICATION,
-            TaskType.MULTILABEL_RANKING,
-            TaskType.TEMPORAL_LINK_PREDICTION,
-            TaskType.STATIC_LINK_PREDICTION,
-    }:
+    if task_type in _LONG_FORMAT_TASKS:
         format = 'long'
     else:
         assert top_k is None, f"`top_k` is not supported for {task_type}."

packages/kumo-core/tests/test_prediction_output_columns.py

@@ -0,0 +1,23 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+import pytest
+from kumo_core.train.formatting import prediction_output_columns
+from kumo_core.types import TaskType
+
+
+@pytest.mark.parametrize('task_type, columns', [
+    (TaskType.BINARY_CLASSIFICATION,
+     ['TARGET_PRED', 'False_PROB', 'True_PROB']),
+    (TaskType.REGRESSION, ['TARGET_PRED']),
+    (TaskType.FORECASTING, ['TIMESTAMP', 'TARGET_PRED']),
+    (TaskType.MULTICLASS_CLASSIFICATION, ['PREDICTED', 'CLASS', 'SCORE']),
+    (TaskType.MULTILABEL_CLASSIFICATION, ['PREDICTED', 'CLASS', 'SCORE']),
+    (TaskType.MULTILABEL_RANKING, ['CLASS', 'SCORE']),
+    (TaskType.TEMPORAL_LINK_PREDICTION, ['CLASS', 'SCORE']),
+    (TaskType.STATIC_LINK_PREDICTION, ['CLASS', 'SCORE']),
+])
+def test_prediction_output_columns(task_type, columns) -> None:
+    assert prediction_output_columns(task_type) == columns

packages/kumo-engine/docs/mlops.md

@@ -203,9 +203,50 @@ 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`, and the best-effort hooks `log_model`,
+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
 (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
 and param names.
+
+### MLflow models
+
+With `registry=ModelRegistryLocation(name, alias=...)` and an `MlflowTracking` tracker, every fit in the block logs its model on the fit run as an MLflow pyfunc model and registers it as a new version of `name` (Unity Catalog with `MLFLOW_REGISTRY_URI=databricks-uc`, SageMaker managed MLflow, or any MLflow registry). One alias is set on the new version. A bare `ModelRegistryLocation` without such a tracker, a fixed `version`, or an invalid alias raises when `kumo.MLOps` is built.
+
+```python
+from kumo.engine import ModelRegistryLocation
+from kumo.engine.integrations.mlflow import MlflowTracking, load_model
+
+cfg = kumo.MLOps("s3://bucket/kumo",
+                 tracking=MlflowTracking(experiment_name="/Shared/users_spend"),
+                 registry=ModelRegistryLocation("main.ml.users_spend", alias="champion"))
+with kumo.mlops(cfg):
+    result = trainer.fit(training_table, snapshot=snapshot)
+result.registered_model   # PublishedModel(name='main.ml.users_spend', version='4', uri='models:/main.ml.users_spend/4')
+
+trainer = load_model("models:/main.ml.users_spend@champion")   # the Trainer back
+model = mlflow.pyfunc.load_model("models:/main.ml.users_spend@champion")
+scores = model.predict(entities, params={"snapshot": "s3://bucket/kumo/snapshots/<graph key>/<version>"})
+```
+
+- `fit` waits for the upload and registration before it returns, for up to an hour, with a 600 s MLflow HTTP timeout during the upload (unless `MLFLOW_HTTP_REQUEST_TIMEOUT` is set). If they fail, the tracking thread stops or the hour passes, `fit` warns and still returns its result with `registered_model=None`, the run ends `FAILED`, and a model that finishes uploading later is not registered.
+- `MlflowTracking(log_model=True)` logs the model without registering it. `log_model(trainer, snapshot=...)` logs a fitted or loaded model in any run, with the bundle's source tables as run inputs.
+- The model's metadata is `kumo_bundle_id`; its `provenance` artifact is the bundle's `_provenance.json`. Its pip requirement is `kumo-engine[s3,gcs]` of the same version.
+- The pyfunc model loads a snapshot on its first `predict` with that URI and keeps it for the life of the process: each serving replica, and each Spark executor of `mlflow.pyfunc.spark_udf`, downloads the snapshot once and keeps every snapshot it has loaded (no eviction), so its disk must hold them. With `serving_snapshot=<snapshot URI>` (on `log_model` or `MlflowTracking`), each replica downloads that snapshot when it starts, and `predict` uses it when no `snapshot` param is given. `params={"population": True}` scores every entity and returns the whole frame in memory; for large populations use `Trainer.predict` with a `PredictionOutputConfig`. Each `predict` call works in its own scratch directory and removes it before returning.
+- The `mlflow` extra needs an MLflow 3.13 or later client: DiskGraph pins pyarrow 24, which older MLflow releases exclude. SageMaker managed MLflow servers run MLflow 3.0; a 3.17 client against a 3.0 server handles logging, registration, aliases and `load_model`.
+- On Databricks, the parent folder of an MLflow experiment path must exist. Registering to Unity Catalog from outside Databricks also needs `boto3`, which the `mlflow` extra does not install.
+
+#### Which registry
+
+Pick one per registered model, since the two store different formats. Both record only the `kumo_bundle_id` pointer; lineage stays in the bundle's `_provenance.json`. Both accept the same aliases (`kumoapi.plugin.check_model_alias`).
+
+| | Connector registry | MLflow model |
+|---|---|---|
+| Register | `kumo.MLOps(registry=ModelOutputConfig(connector, location=ModelRegistryLocation(name, alias=...)))` | `kumo.MLOps(tracking=MlflowTracking(), registry=ModelRegistryLocation(name, alias=...))` |
+| Load | `Trainer.load_from_connector(connector, location=...)` | `mlflow.pyfunc.load_model(uri)` or `load_model(uri)` |
+| Promote | `set_model_alias(connector, location, alias)` | MLflow's `set_registered_model_alias` |
+| Databricks | `DatabricksConnector` (Unity Catalog) | Unity Catalog through MLflow |
+| SageMaker | not available | SageMaker managed MLflow (`MLFLOW_TRACKING_URI` = tracking server ARN) |
+| Snowflake | `SnowflakeConnector` (Snowflake Model Registry) | not available |
+| Plain VM | not available (store under the artifact root instead) | any MLflow tracking server |

packages/kumo-engine/kumo/engine/_mlops.py

@@ -8,6 +8,8 @@
 from dataclasses import dataclass, field
 from typing import Any, FrozenSet, Iterable, Iterator, Optional, Tuple, Union
 
+from kumoapi.plugin import ModelRegistryLocation, check_model_alias
+
 from kumo.engine._manifest import KINDS
 
 
@@ -79,6 +81,17 @@ def __post_init__(self) -> None:
         if unknown:
             raise ValueError(f'Unknown export kinds {unknown}; choose from '
                              f'{list(KINDS)}.')
+        if isinstance(self.registry, ModelRegistryLocation):
+            if not callable(getattr(self.tracking, 'log_model', None)):
+                raise ValueError(
+                    'registry=ModelRegistryLocation(...) registers through '
+                    'the MLflow tracker: pass tracking=MlflowTracking(...), '
+                    'or registry=ModelOutputConfig(connector, location=...).')
+            if self.registry.version is not None:
+                raise ValueError('MLflow assigns registered model versions; '
+                                 'leave ModelRegistryLocation.version unset.')
+            if self.registry.alias is not None:
+                check_model_alias(self.registry.alias)
         object.__setattr__(self, 'root', root)
         object.__setattr__(self, 'exports', kinds)
 

packages/kumo-engine/kumo/engine/_scratch.py

@@ -35,11 +35,15 @@
 
 import hashlib
 import uuid
+from contextlib import contextmanager
+from contextvars import ContextVar
 from pathlib import Path
-from typing import Union
+from typing import Iterator, Optional, Union
 
 _DEFAULT_ROOT = Path.home() / ".cache" / "kumo-engine"
 _root = _DEFAULT_ROOT
+_scoped_root: ContextVar[Optional[Path]] = ContextVar('kumo_scratch_root',
+                                                      default=None)
 
 
 def set_scratch_dir(path: Union[str, Path]) -> None:
@@ -65,7 +69,18 @@ def set_scratch_dir(path: Union[str, Path]) -> None:
 
 
 def get_scratch_dir() -> Path:
-    return _root
+    return _scoped_root.get() or _root
+
+
+@contextmanager
+def scoped_scratch_dir(path: Path) -> Iterator[None]:
+    r"""Nests this thread's (context's) new scratch directories under
+    ``path`` instead of the configured root."""
+    token = _scoped_root.set(path)
+    try:
+        yield
+    finally:
+        _scoped_root.reset(token)
 
 
 def _spool_dir() -> str:

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

@@ -32,6 +32,9 @@
 from kumo.engine._events import Event, _finite_floats
 
 if TYPE_CHECKING:
+    from kumoapi.plugin import ModelRegistryLocation, PublishedModel
+
+    from kumo.engine.graph import GraphSnapshot
     from kumo.engine.trainer.trainer import TrainingResult
 
 logger = logging.getLogger(__name__)
@@ -42,6 +45,7 @@
 # never stalls training; these bound how long the fit waits for it.
 QUEUE_SIZE = 1000
 FINISH_SECONDS = 30.0
+MODEL_SECONDS = 3600.0
 
 
 @dataclass(frozen=True)
@@ -219,8 +223,12 @@ class Tracking(Protocol):
     the event that updated it, and `finish` the fit run's final record
     (`None` if tracking stopped early, which should end the run as failed).
     Optional members: ``prepare()``, called first on the fit's own thread;
-    ``tags`` (extra fit-run tags); and the best-effort hooks
-    ``log_model(result, snapshot)``, ``log_prediction_run(record)`` and
+    ``tags`` (extra fit-run tags); ``log_model(result, snapshot, registry)``,
+    called after a fit that succeeded and before `finish`, which the fit waits
+    for while the tracking thread is alive, up to an hour (an error, a stopped
+    thread or the time limit ends the run as failed), and which returns the
+    version it registered at ``registry``, a ``ModelRegistryLocation``; and the
+    best-effort hooks ``log_prediction_run(record)`` and
     ``log_monitoring(record)``.
     """
     def start(self, record: RunRecord) -> None:
@@ -286,6 +294,56 @@ def on_event(self, event: Event) -> None:
         self._submit(lambda: self._call(
             lambda tracker: tracker.on_trial(record, event)))
 
+    def log_model(
+        self,
+        result: 'TrainingResult',
+        snapshot: 'GraphSnapshot',
+        registry: Optional['ModelRegistryLocation'] = None,
+    ) -> Tuple[Optional['PublishedModel'], Optional[BaseException]]:
+        r"""Runs the tracker's ``log_model`` on the tracking thread and waits
+        for it; returns the registered version and the error it raised."""
+        if (self._worker is None
+                or not hasattr(self._tracker, 'log_model')):
+            return None, None
+        done = threading.Event()
+        failed: List[BaseException] = []
+        registered: List[Any] = []
+
+        def log() -> None:
+            try:
+                tracker = self._tracker
+                if tracker is not None:
+                    registered.append(
+                        getattr(tracker, 'log_model')(result, snapshot,
+                                                      registry))
+            except BaseException as error:
+                failed.append(error)
+            finally:
+                done.set()
+
+        worker = self._worker
+        deadline = time.monotonic() + MODEL_SECONDS
+        try:
+            self._jobs.put(log, timeout=MODEL_SECONDS)
+        except queue.Full:
+            pass
+        while not done.wait(1.0):
+            if worker.is_alive() and time.monotonic() < deadline:
+                continue
+            failed.append(
+                TimeoutError(f'it did not finish within {MODEL_SECONDS:.0f}s')
+                if worker.is_alive() else
+                RuntimeError('the tracking thread stopped'))
+            tracker, self._tracker = self._tracker, None
+            if tracker is not None:
+                _abandon(tracker)
+            break
+        if failed:
+            warnings.warn(f'Could not log the model: {failed[0]}',
+                          UserWarning, stacklevel=3)
+            return None, failed[0]
+        return (registered or [None])[0], None
+
     def finish(self, result: Optional['TrainingResult'],
                error: Optional[BaseException]) -> None:
         worker = self._worker