Add model registry aliases and a bundle id pointer on registered versions

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'
+    }

packages/kumo-databricks/src/kumo_databricks/registry.py

@@ -10,7 +10,6 @@
 import os
 import shutil
 import tempfile
-from dataclasses import replace
 from pathlib import Path
 from typing import Any, Iterator, Mapping, Optional, cast
 from urllib.parse import quote
@@ -20,12 +19,15 @@
 from kumo_databricks.session import _DatabricksSessions, _target_config
 from kumo_databricks.sql import normalize_table_coordinate
 from kumoapi.plugin import (
+    BUNDLE_ID_KEY,
+    ENGINE_MARKER,
     ModelRegistryLocation,
     PublishedModel,
     SourceConnector,
+    check_model_alias,
+    registry_lineage,
 )
 
-_ENGINE_MARKER = 'kumo_engine'
 _MODELS = '/api/2.1/unity-catalog/models'
 _LOG = logging.getLogger(__name__)
 
@@ -56,6 +58,17 @@ def _tags(version: Mapping[str, Any]) -> dict[str, str]:
     return {tag['key']: tag['value'] for tag in version.get('tags', [])}
 
 
+def _alias_path(name: str, alias: str) -> str:
+    return f'{_MODELS}/{quote(name, safe="")}/aliases/{quote(alias, safe="")}'
+
+
+def _version_number(location: ModelRegistryLocation) -> Optional[int]:
+    version = location.version
+    if version is not None and (not version.isdecimal() or int(version) < 1):
+        raise ValueError('Databricks model versions are positive integers.')
+    return None if version is None else int(version)
+
+
 def _published(
     name: str,
     version: Mapping[str, Any],
@@ -133,7 +146,9 @@ def publish_bundle(
     verify_model_bundle(root)
     _validate_catalog(dict(metadata=dict(metadata), metrics=dict(metrics)))
     catalog_name, schema, model = name.split('.')
-    tags = {_ENGINE_MARKER: '1'}
+    tags = {ENGINE_MARKER: '1'}
+    if metadata.get(BUNDLE_ID_KEY):
+        tags[BUNDLE_ID_KEY] = str(metadata[BUNDLE_ID_KEY])
     try:
         client.registered_models.create(catalog_name=catalog_name,
                                         schema_name=schema, name=model)
@@ -200,7 +215,7 @@ def _bundle_metadata(bundle: Path) -> tuple[dict[str, str], dict[str, float]]:
 def _download(client: WorkspaceClient, remote: str, destination: Path) -> None:
     for entry in client.files.list_directory_contents(remote):
         # A server-provided path must stay an immediate child of this directory.
-        path = entry.path or ''
+        path = (entry.path or '').rstrip('/')
         prefix = remote + '/'
         child = path[len(prefix):] if path.startswith(prefix) else ''
         if not child or '/' in child or child in ('.', '..') or '\\' in child:
@@ -219,21 +234,26 @@ def _download(client: WorkspaceClient, remote: str, destination: Path) -> None:
 
 
 def fetch_bundle(client: WorkspaceClient, name: str, local_dir: Path, *,
-                 version: Optional[str] = None) -> PublishedModel:
+                 version: Optional[str] = None,
+                 alias: Optional[str] = None) -> PublishedModel:
     from kumo_core.artifacts import verify_model_bundle
 
-    if version is None:
+    if alias is not None:
+        aliased = cast(dict[str, Any],
+                       client.api_client.do('GET', _alias_path(name, alias)))
+        model_version = _get_version(client, name, aliased['version'])
+    elif version is None:
         versions = [
             v for v in _versions(client, name) if v.get('status') == 'READY'
-            and _tags(v).get(_ENGINE_MARKER) == '1'
+            and _tags(v).get(ENGINE_MARKER) == '1'
         ]
         if not versions:
             raise ValueError('No ready Kumo model version in this registry.')
         model_version = max(versions, key=lambda v: int(v['version']))
     else:
         model_version = _get_version(client, name, version)
     if model_version.get('status') != 'READY' or _tags(model_version).get(
-            _ENGINE_MARKER) != '1':
+            ENGINE_MARKER) != '1':
         raise ValueError('This version is not a ready Kumo model bundle.')
     destination = local_dir
     destination.parent.mkdir(parents=True, exist_ok=True)
@@ -251,9 +271,12 @@ def fetch_bundle(client: WorkspaceClient, name: str, local_dir: Path, *,
                 'The registered model must contain one Kumo bundle.')
         bundle = bundles[0]
         verify_model_bundle(bundle)
-        published = _published(name, model_version)
         metadata, metrics = _bundle_metadata(bundle)
-        published = replace(published, metadata=metadata, metrics=metrics)
+        published = _published(
+            name, model_version, metadata={
+                **metadata,
+                **registry_lineage(_tags(model_version))
+            }, metrics=metrics)
         backup = Path(tmp) / 'previous'
         if destination.exists():
             os.replace(destination, backup)
@@ -277,8 +300,10 @@ def list_models(client: WorkspaceClient, catalog: str,
             continue
         for version in _versions(client, name):
             if version.get('status') == 'READY' and _tags(version).get(
-                    _ENGINE_MARKER) == '1':
-                found.append(_published(name, version))
+                    ENGINE_MARKER) == '1':
+                found.append(
+                    _published(name, version,
+                               metadata=registry_lineage(_tags(version))))
     return sorted(found, key=lambda item: (item.name, int(item.version)))
 
 
@@ -321,15 +346,26 @@ def fetch_model(
         location: Optional[ModelRegistryLocation] = None,
         project_id: Optional[str] = None,
     ) -> PublishedModel:
-        """Download the named version, or the latest ready Kumo version."""
+        """Download the named or aliased version, else the latest Kumo one."""
         del project_id
         name = _address(connector, path, location)
         assert location is not None
-        if location.version is not None and (not location.version.isdecimal()
-                                             or int(location.version) < 1):
-            raise ValueError('Databricks model versions are positive integers.')
+        _version_number(location)
         return fetch_bundle(self._client(connector), name,
-                            Path(local_dir).resolve(), version=location.version)
+                            Path(local_dir).resolve(), version=location.version,
+                            alias=getattr(location, 'alias', None))
+
+    def set_model_alias(self, connector: SourceConnector,
+                        location: ModelRegistryLocation, alias: str) -> None:
+        """Point an alias at a version; UC moves it from any other."""
+        name = _address(connector, None, location)
+        number = _version_number(location)
+        if number is None:
+            raise ValueError('Setting an alias needs a ModelRegistryLocation '
+                             'naming a version.')
+        check_model_alias(alias)
+        self._client(connector).api_client.do('PUT', _alias_path(name, alias),
+                                              body=dict(version_num=number))
 
     def list_models(self, connector: SourceConnector, *,
                     project_id: Optional[str] = None) -> list[PublishedModel]: