Record fits in experiment tracking and publish models to registries

.github/workflows/testing.yml

@@ -723,6 +723,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,118 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Manual live check per platform; Databricks reads DATABRICKS_CLIENT_SECRET."""
+
+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,13 @@
 from kumoapi.typing import Dtype
 
 __all__ = [
+    'BUNDLE_ID_KEY',
+    'ENGINE_MARKER',
     'ArtifactLocation',
     'ArtifactSink',
     'FileStoreSink',
     'IngestResult',
+    'ModelAliasingSink',
     'ModelFetchingSink',
     'ModelListingSink',
     'ModelPublishingSink',
@@ -56,8 +61,16 @@
     'SourceBackend',
     'SourceConnector',
     'SourceCredentials',
+    'check_model_alias',
+    'registry_lineage',
 ]
 
+_REGISTRY_KEY_PREFIX = 'kumo_'
+ENGINE_MARKER = 'kumo_engine'
+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 +162,15 @@ 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 on export; read on import.
     """
     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 +181,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)
@@ -702,7 +721,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:
@@ -712,6 +732,38 @@ def fetch_model(
         ...
 
 
+@runtime_checkable
+class ModelAliasingSink(ArtifactSink, Protocol):
+    r"""A registry sink that sets and fetches by 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 registry accepts ``alias``."""
+    if (not isinstance(alias, str) or not _ALIAS.fullmatch(alias)
+            or _RESERVED_ALIAS.fullmatch(alias)):
+        raise ValueError(
+            f'{alias!r} cannot be a model alias: use letters, digits and '
+            f'underscores, not starting with a digit, other than the reserved '
+            f'DEFAULT, FIRST, LAST, LATEST and v<digits>.')
+
+
+def registry_lineage(recorded: Mapping[str, Any]) -> Dict[str, str]:
+    r"""A version's ``kumo_*`` registry metadata, without the 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'
+    }

packages/kumo-core/src/kumo_core/artifact_export.py

@@ -304,10 +304,14 @@ def prepared_artifact_file(path, extra_columns=None, *, parquet_options=None,
         writer = None
         try:
             for batch in local_parquet_batches(path):
-                frame = batch.to_pandas()
+                # Pandas turns all-null bool columns into Arrow nulls.
+                table = pa.Table.from_batches([batch])
                 for name, value in (extra_columns or {}).items():
-                    frame[name] = value
-                table = pa.Table.from_pandas(frame, preserve_index=False)
+                    column = pa.repeat(pa.scalar(value), table.num_rows)
+                    table = (table.set_column(
+                        table.column_names.index(name), name, column)
+                             if name in table.column_names else
+                             table.append_column(name, column))
                 if writer is None:
                     writer = pq.ParquetWriter(staged, table.schema,
                                               **(parquet_options or {}))

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
+
 PROVENANCE_FILE = '_provenance.json'
 
 
@@ -378,6 +382,26 @@ def add_bundle_file(bundle_dir: Path, role: str, relative_path: str) -> None:
     _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'
@@ -475,6 +499,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.