Record fits in experiment tracking and publish models to registries

.github/workflows/testing.yml

@@ -720,6 +720,41 @@ jobs:
           cd packages/kumo-sdk && python -m pytest -rs \
             -W "ignore:You have an incompatible version of 'pyarrow' installed:UserWarning"
 
+  mlflow-floor:
+    # kumo-engine's `mlflow` extra allows the MLflow a Databricks Runtime
+    # ships; run the MLflow suites against the extra's floor. The tests left
+    # out read logs, which MLflow 3.5's SQLite migrations reconfigure.
+    needs: [diskgraph-wheel]
+    if: ${{ !cancelled() }}
+    runs-on: ubuntu-24.04
+    permissions:
+      contents: read
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6
+        with:
+          ref: ${{ github.sha }}
+
+      - name: Set up packages
+        uses: ./.github/actions/setup
+        with:
+          github_token: ${{ secrets.KUMO_GITHUB_BOT_TOKEN }}
+          diskgraph-wheel-artifact: diskgraph-wheel
+
+      - name: Run the MLflow suites at the mlflow extra's floor
+        run: |
+          uv pip install --no-deps -e packages/kumo-engine
+          # mlflow 3.5.0 pins a pyarrow with no wheel for this Python; the
+          # extra installs mlflow-skinny, which holds the client and stores.
+          uv pip uninstall mlflow mlflow-skinny mlflow-tracing
+          uv pip install "mlflow-skinny==3.5.0"
+          python -c "import mlflow; assert mlflow.__version__ == '3.5.0'"
+          cd packages/kumo-engine && python -m pytest -rs \
+            test/test_mlflow_target.py test/test_mlflow_model.py \
+            test/test_mlflow_tracking.py -k "not (reuses_its_version or \
+            late_one_is_dropped or names_its_cause_and_can_raise or \
+            ends_runs_failed_and_is_not_raised)"
+
   durations-aggregate:
     runs-on: ubuntu-24.04
     needs: [pytest]

examples/mlops/smoke.py

@@ -0,0 +1,133 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Manual live check, not run in CI: fit a small model on synthetic data on
+this machine, record it in Databricks, SageMaker or Snowflake with an
+explicit identity, register it, and load it back. A few minutes on a CPU.
+Install kumo-engine[mlflow,databricks], kumo-engine[mlflow,sagemaker] or
+kumo-engine[snowflake].
+
+    python smoke.py databricks --host https://<workspace> --client-id <id> \
+        --catalog main --schema ml --warehouse <id> --experiment /Users/<you>/kumo-smoke
+    python smoke.py sagemaker --arn <tracking server ARN> --role-arn <role ARN>
+    python smoke.py snowflake --database ML --schema RUNS --warehouse <wh> \
+        --account <account> --user <user> --private-key-file key.p8
+
+Secrets come from DATABRICKS_CLIENT_SECRET (Databricks) or the key file
+(Snowflake); --role-arn may be left out to use the ambient AWS identity,
+with no MLFLOW_TRACKING_URI needed.
+"""
+
+import argparse
+import os
+import tempfile
+import time
+from pathlib import Path
+from typing import Optional, Sequence
+
+import numpy as np
+import pandas as pd
+
+import kumo.engine as kumo
+from kumo.engine import ModelRegistryLocation
+from kumo.engine.connector import (
+    DatabricksConnector,
+    DatabricksIdentity,
+    S3Identity,
+    SnowflakeConnector,
+    SnowflakeIdentity,
+)
+
+QUERY = 'PREDICT COUNT(orders.*, 0, 7, days) > 0 FOR EACH users.USER_ID'
+
+
+def data(root: Path) -> kumo.Graph:
+    rng = np.random.default_rng(0)
+    root.mkdir(parents=True)
+    pd.DataFrame({'USER_ID': np.arange(200)}).to_parquet(root / 'users.parquet')
+    pd.DataFrame({
+        'ORDER_ID': np.arange(6000),
+        'USER_ID': rng.integers(0, 200, 6000),
+        'ORDER_TIME': pd.Timestamp('2024-01-01') +
+        pd.to_timedelta(rng.integers(0, 120, 6000), unit='D'),
+    }).to_parquet(root / 'orders.parquet')
+    files = kumo.FileConnector(root_dir=str(root))
+    orders = kumo.Table.from_source_table(
+        files['orders.parquet'], primary_key='ORDER_ID',
+        time_column='ORDER_TIME').infer_metadata()
+    return kumo.Graph(
+        tables={
+            'users': kumo.Table.from_source_table(
+                files['users.parquet'], primary_key='USER_ID'),
+            'orders': orders,
+        }, edges=[kumo.Edge(src_table='orders', fkey='USER_ID',
+                            dst_table='users')])
+
+
+def tracking(args: argparse.Namespace, suffix: str):
+    from kumo.engine.integrations.mlflow import MlflowTracking, SageMakerMLflow
+
+    if args.platform == 'databricks':
+        workspace = DatabricksConnector(
+            catalog=args.catalog, schema_name=args.schema,
+            warehouse_id=args.warehouse, identity=DatabricksIdentity(
+                host=args.host, client_id=args.client_id,
+                client_secret=os.environ['DATABRICKS_CLIENT_SECRET']))
+        return (MlflowTracking(target=workspace,
+                               experiment_name=args.experiment),
+                f'{args.catalog}.{args.schema}.kumo_smoke_{suffix}')
+    if args.platform == 'sagemaker':
+        identity = S3Identity(role_arn=args.role_arn) if args.role_arn else None
+        return (MlflowTracking(
+            target=SageMakerMLflow(args.arn, identity=identity),
+            experiment_name='kumo-smoke'), f'kumo-smoke-{suffix}')
+    from kumo.engine.integrations.snowflake import SnowflakeTracking
+
+    connector = SnowflakeConnector(
+        warehouse=args.warehouse, database=args.database,
+        schema_name=args.schema, identity=SnowflakeIdentity(
+            account=args.account, username=args.user,
+            private_key=Path(args.private_key_file).read_text()))
+    return (SnowflakeTracking(connector, experiment_name='KUMO_SMOKE'),
+            f'KUMO_SMOKE_{suffix}')
+
+
+def main(argv: Optional[Sequence[str]] = None) -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument('platform',
+                        choices=('databricks', 'sagemaker', 'snowflake'))
+    for flag in ('--host', '--client-id', '--catalog', '--schema',
+                 '--warehouse', '--experiment', '--arn', '--role-arn',
+                 '--database', '--account', '--user', '--private-key-file'):
+        parser.add_argument(flag)
+    parser.add_argument('--artifact-root',
+                        help='defaults to a local temporary directory')
+    args = parser.parse_args(argv)
+
+    started = time.monotonic()
+    work = Path(tempfile.mkdtemp(prefix='kumo-smoke-'))
+    kumo.configure(scratch_dir=work / 'scratch')
+    graph = data(work / 'data')
+    tracker, model_name = tracking(args, time.strftime('%Y%m%d%H%M%S'))
+    location = ModelRegistryLocation(model_name, alias='candidate')
+    cfg = kumo.MLOps(args.artifact_root or str(work / 'root'),
+                     tracking=tracker, registry=location, on_error='raise')
+    print(f'checked the target in {time.monotonic() - started:.0f} s')
+    with kumo.mlops(cfg):
+        snapshot = graph.snapshot()
+        query = kumo.PredictiveQuery(graph=graph, query=QUERY)
+        table = query.generate_training_table(snapshot)
+        plan = query.suggest_model_plan(snapshot)
+        plan.training_job.num_experiments = 1
+        plan.optimization.max_epochs = 1
+        result = kumo.Trainer(plan).fit(table, snapshot=snapshot)
+    print(f'fit run {result.tracking}, bundle {result.id}, '
+          f'registered {result.registered_model}')
+    loaded = kumo.Trainer.load_from_connector(cfg, location=location)
+    assert loaded.result.id == result.id, (loaded.result.id, result.id)
+    print('loaded the registered model back')
+    print(f'done in {time.monotonic() - started:.0f} s')
+
+
+if __name__ == '__main__':
+    main()

packages/kumo-api/kumoapi/plugin.py

@@ -21,10 +21,12 @@
 :class:`~kumoai_studio.backend.LocalBackend`.
 """
 
+import re
 from dataclasses import dataclass, field
 from typing import (
     Any,
     Callable,
+    Dict,
     List,
     Mapping,
     Optional,
@@ -42,10 +44,14 @@
 from kumoapi.typing import Dtype
 
 __all__ = [
+    'BUNDLE_ID_KEY',
+    'ENGINE_MARKER',
+    'REGISTRY_KEY_PREFIX',
     'ArtifactLocation',
     'ArtifactSink',
     'FileStoreSink',
     'IngestResult',
+    'ModelAliasingSink',
     'ModelFetchingSink',
     'ModelListingSink',
     'ModelPublishingSink',
@@ -56,8 +62,20 @@
     'SourceBackend',
     'SourceConnector',
     'SourceCredentials',
+    'check_model_alias',
+    'registry_lineage',
 ]
 
+#: Prefix of the keys an engine export records in registry metadata.
+REGISTRY_KEY_PREFIX = 'kumo_'
+#: Registry metadata key marking a version as a Kumo engine bundle.
+ENGINE_MARKER = 'kumo_engine'
+#: Registry metadata key holding the version's bundle id; the bundle's own
+#: provenance.json holds its lineage.
+BUNDLE_ID_KEY = 'kumo_bundle_id'
+_ALIAS = re.compile(r'[A-Za-z_][A-Za-z0-9_]{0,254}')
+_RESERVED_ALIAS = re.compile(r'DEFAULT|FIRST|LAST|LATEST|V[0-9]+', re.I)
+
 
 @dataclass(frozen=True)
 class SourceCredentials:
@@ -149,13 +167,16 @@ class ModelRegistryLocation:
     Args:
         model_name: Registry-local model name.
         version: Exact version to write or read. ``None`` lets the registry
-            assign the next one on export, and selects the latest on import.
-            An assigned version comes back on the export result.
+            assign the next one on export; on import it selects Snowflake's
+            default or Databricks' latest Kumo version.
         comment: Optional description recorded on the version.
+        alias: Set on the new version after export; on import, the version
+            it points at. Exclusive with :obj:`version`.
     """
     model_name: str
     version: Optional[str] = None
     comment: Optional[str] = None
+    alias: Optional[str] = None
 
     def __post_init__(self) -> None:
         if (not isinstance(self.model_name, str) or not self.model_name
@@ -166,6 +187,10 @@ def __post_init__(self) -> None:
             raise ValueError(
                 'version must be a non-empty version, or None to let the '
                 'registry choose')
+        if self.alias is not None and (not self.alias
+                                       or self.version is not None):
+            raise ValueError('alias must be a non-empty alias, and is '
+                             'mutually exclusive with version')
 
 
 @dataclass(frozen=True)
@@ -705,7 +730,8 @@ def fetch_model(
             local_dir: Directory to place the bundle in. Created if absent.
             path: Source for a path-addressed target.
             location: A registry location. A version of ``None`` selects the
-                latest, which is resolved here and reported back.
+                registry's choice (see :class:`ModelRegistryLocation`), which
+                is resolved here and reported back.
             project_id: Authorizing project, where one applies.
 
         Returns:
@@ -715,6 +741,44 @@ def fetch_model(
         ...
 
 
+@runtime_checkable
+class ModelAliasingSink(ArtifactSink, Protocol):
+    r"""A registry sink that sets aliases and fetches by
+    :attr:`ModelRegistryLocation.alias`."""
+    def set_model_alias(
+        self,
+        connector: SourceConnector,
+        location: ModelRegistryLocation,
+        alias: str,
+    ) -> None:
+        r"""Points :obj:`alias` at ``location.version``, moving it if held."""
+        ...
+
+
+def check_model_alias(alias: str) -> None:
+    r"""Raises :class:`ValueError` unless every model registry can set
+    :obj:`alias`: a letter or underscore, then letters, digits or underscores
+    (at most 255), and not ``DEFAULT``, ``FIRST``, ``LAST``, ``LATEST`` or
+    ``v<digits>`` in any case."""
+    if (not isinstance(alias, str) or not _ALIAS.fullmatch(alias)
+            or _RESERVED_ALIAS.fullmatch(alias)):
+        raise ValueError(
+            f'{alias!r} cannot be set as a model alias: use letters, digits '
+            f'and underscores, starting with a letter or underscore, and '
+            f'avoid the reserved DEFAULT, FIRST, LAST, LATEST and v<digits> '
+            f'(for DEFAULT use ALTER MODEL ... SET DEFAULT_VERSION).')
+
+
+def registry_lineage(recorded: Mapping[str, Any]) -> Dict[str, str]:
+    r"""The ``kumo_*`` entries of a version's registry metadata, such as
+    :data:`BUNDLE_ID_KEY`, as strings and without :data:`ENGINE_MARKER`."""
+    return {
+        key: str(value)
+        for key, value in recorded.items()
+        if key.startswith(REGISTRY_KEY_PREFIX) and key != ENGINE_MARKER
+    }
+
+
 @runtime_checkable
 class ModelListingSink(ArtifactSink, Protocol):
     r"""A sink that can enumerate the models published through it.

packages/kumo-api/test/test_plugin_registry.py

@@ -0,0 +1,33 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+import pytest
+
+from kumoapi.plugin import check_model_alias, registry_lineage
+
+
+@pytest.mark.parametrize('alias', [
+    '', ' ', 'default', 'LAST', 'Latest', 'v3', 'my-alias', '2nd', 'a' * 256,
+    None
+])
+def test_check_model_alias_refuses_what_a_registry_reserves(alias):
+    with pytest.raises(ValueError):
+        check_model_alias(alias)
+
+
+def test_check_model_alias_accepts_a_plain_name():
+    check_model_alias('champion')
+    check_model_alias('_v2_prod')
+
+
+def test_registry_lineage_keeps_kumo_keys_as_strings():
+    recorded = {
+        'kumo_engine': '1',
+        'kumo_bundle_id': 'b',
+        'auc': 0.9,
+        'kumo_seconds': 3
+    }
+    assert registry_lineage(recorded) == {
+        'kumo_bundle_id': 'b',
+        'kumo_seconds': '3'
+    }