From 0e6ea1cb5a05e49c877a0040e8706633e76f5438 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 22:14:31 +0000 Subject: [PATCH 001/278] Read the holdout of the published model, not the last-evaluated one fit() read test_df.parquet from best_trainer.log_dir, which is whichever refit stage ran last, not necessarily the model the bundle publishes. With refit_full the published model has seen the test split, so its holdout is meaningless. Read the holdout from the published model's log dir (Trainer.published_fit_mode, also used to pick the prediction model) and return none for refit_full. get_log_dir takes an optional fit_mode override. Verified: test_holdout_df (fit / refit_trainval with explain / refit_full), test_training_explanations, test_semantic_carry_forward: 35 passed. Co-Authored-By: Claude Opus 5.5 --- .../kumo-core/src/kumo_core/train/trainer.py | 33 +++++++++++-------- .../kumo/engine/trainer/trainer.py | 26 +++++++++++---- packages/kumo-engine/test/test_holdout_df.py | 28 ++++++++++++++-- .../test/test_semantic_carry_forward.py | 15 ++++++--- .../test/test_training_explanations.py | 3 +- 5 files changed, 76 insertions(+), 29 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/train/trainer.py b/packages/kumo-core/src/kumo_core/train/trainer.py index b18f552e8..efe6968f2 100644 --- a/packages/kumo-core/src/kumo_core/train/trainer.py +++ b/packages/kumo-core/src/kumo_core/train/trainer.py @@ -1507,6 +1507,7 @@ def get_log_dir( self, local: bool = False, original: bool = False, + fit_mode: FitMode | None = None, ) -> str: r"""Switch the log directory based on `self.fit_mode`. `version_xx/`: fit mode. @@ -1516,16 +1517,17 @@ def get_log_dir( If `original=True`, then we get the original (non-refitting) log directory. This way, we can update `config.json` in the original - directory""" + directory. `fit_mode` overrides `self.fit_mode`.""" + fit_mode = self.fit_mode if fit_mode is None else fit_mode path = self.root_dir if not local else self.local_root_dir # In case the `root_dir` contains a trainer config, it is a log dir # for a previous model version. if fsutil.exists(osp.join(path, 'config.json')): - if not original and self.fit_mode.is_refit: + if not original and fit_mode.is_refit: if path[-1] == '/': path = path[:-1] - path = f"{path}_{self.fit_mode.value}" + path = f"{path}_{fit_mode.value}" fsutil.makedirs(path, exist_ok=True) return path else: @@ -1535,10 +1537,10 @@ def get_log_dir( path = osp.join(path, self.name) if self.track_version: path = osp.join(path, self.version_dir_name) - if not original and self.fit_mode.is_refit: + if not original and fit_mode.is_refit: if path[-1] == '/': path = path[:-1] - path = f"{path}_{self.fit_mode.value}" + path = f"{path}_{fit_mode.value}" fsutil.makedirs(path, exist_ok=True) return path @@ -2531,16 +2533,21 @@ def _preprocess_top_k(self, top_k: int | None): # Batch prediction ######################################################## + def published_fit_mode(self) -> FitMode | None: + r"""The fit mode predictions use by default: the model trained on the + most data, or `None` before fitting.""" + if self.refit_cfg.refitted_full: + return FitMode.refit_full + if self.refit_cfg.refitted_trainval: + return FitMode.refit_trainval + if self.refit_cfg.fitted: + return FitMode.fit + return None + def _set_fit_mode_to_use(self, fit_mode_to_use: FitMode | None = None): - # Infer `fit_mode_to_use` - Use model that's trained on most data: if fit_mode_to_use is None: - if self.refit_cfg.refitted_full: - fit_mode_to_use = FitMode.refit_full - elif self.refit_cfg.refitted_trainval: - fit_mode_to_use = FitMode.refit_trainval - elif self.refit_cfg.fitted: - fit_mode_to_use = FitMode.fit - else: + fit_mode_to_use = self.published_fit_mode() + if fit_mode_to_use is None: raise RuntimeError("Trainer is not fitted yet.") else: self._check_readiness_for_prediction(fit_mode_to_use) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 1b525b010..28c3156a3 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -77,11 +77,12 @@ class TrainingResult: `published_model` records the destination's assigned name, version and URI when `fit()` exports the bundle. It is `None` without an export. - `holdout_path` is the holdout Parquet file written by `fit()`, or `None` - when the task produced none (also for results of `Trainer.load`). It is - not part of the published model bundle. The file is capped at 1 GiB in - memory; `holdout_truncated` is true when rows were dropped to fit, and - `None` when the size record could not be read. + `holdout_path` is the published model's holdout Parquet file written by + `fit()`, or `None` when it has none (no holdout task, a model refit on the + full data, or results of `Trainer.load`). It is not part of the published + model bundle. The file is capped at 1 GiB in memory; `holdout_truncated` is + true when rows were dropped to fit, and `None` when the size record could + not be read. """ metrics: Dict[str, float] @@ -187,6 +188,17 @@ def _find_holdout(log_dir: str) -> tuple[Optional[Path], Optional[bool]]: return path, truncated +def _published_holdout(trainer: Any) -> tuple[Optional[Path], Optional[bool]]: + from kumo_core.training.config import FitMode + + published = trainer.published_fit_mode() + if published is FitMode.refit_full: + logger.info('The published model is refit on the full data, including ' + 'the test split, so it has no holdout predictions.') + return None, False + return _find_holdout(trainer.get_log_dir(fit_mode=published)) + + def _collect_garbage(stage: str) -> None: # Lightning links each trainer to its data module and model in reference # cycles; free a finished run now rather than at a later collection. @@ -494,8 +506,8 @@ def fit( 'failed for run %s. The trained model remains available.', run_id) - holdout_path, holdout_truncated = _find_holdout( - result.best_trainer.log_dir) + holdout_path, holdout_truncated = _published_holdout( + result.best_trainer) training_result = TrainingResult( metrics=dict(result.metrics), artifact_paths=dict(result.artifact_paths), diff --git a/packages/kumo-engine/test/test_holdout_df.py b/packages/kumo-engine/test/test_holdout_df.py index 8c03af648..70bfe82c6 100644 --- a/packages/kumo-engine/test/test_holdout_df.py +++ b/packages/kumo-engine/test/test_holdout_df.py @@ -2,6 +2,8 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import re + import numpy as np import pandas as pd import pytest @@ -65,13 +67,16 @@ def graph(tmp_path_factory): QUERY = 'PREDICT COUNT(events.*, 0, 7, days) > 0 FOR EACH ents.ID' -def _fit(graph, query): +def _fit(graph, query, refit_trainval=True, refit_full=False, explain=False): graph, snapshot = graph pquery = kumo.PredictiveQuery(graph=graph, query=query) plan = ModelPlan() - plan.training_job.disable_explain = True + plan.training_job.disable_explain = not explain + plan.training_job.refit_trainval = refit_trainval + plan.training_job.refit_full = refit_full plan.optimization.max_epochs = 1 - plan.optimization.max_steps_per_epoch = 2 + # refit_full splits each epoch's batches across 3 stages. + plan.optimization.max_steps_per_epoch = 3 if refit_full else 2 trainer = kumo.Trainer(model_plan=plan) return trainer.fit(pquery.generate_training_table(snapshot), snapshot=snapshot) @@ -90,6 +95,23 @@ def test_holdout_df_after_fit(graph): assert not (result.bundle_dir / 'test_df.parquet').exists() +@pytest.mark.parametrize('refit_trainval,refit_full,explain,suffix', [ + (False, False, False, ''), + (True, False, True, '_refit_trainval'), + (True, True, False, None), +]) +def test_holdout_is_the_published_models(graph, refit_trainval, refit_full, + explain, suffix): + result = _fit(graph, QUERY, refit_trainval, refit_full, explain) + + if suffix is None: + assert result.holdout_path is None + else: + assert re.fullmatch(rf'version_\d+{suffix}', + result.holdout_path.parent.name) + assert len(result.holdout_df()) > 0 + + def test_no_holdout(graph): result = _fit( graph, 'PREDICT LIST_DISTINCT(events.ITEM_ID, 0, 7, days) ' diff --git a/packages/kumo-engine/test/test_semantic_carry_forward.py b/packages/kumo-engine/test/test_semantic_carry_forward.py index 15a4e2330..a1533b131 100644 --- a/packages/kumo-engine/test/test_semantic_carry_forward.py +++ b/packages/kumo-engine/test/test_semantic_carry_forward.py @@ -44,6 +44,11 @@ def _stub_generated_task_dataset_builder(monkeypatch): ) +def _best_trainer(log_dir): + return SimpleNamespace(get_log_dir=lambda fit_mode: str(log_dir), + published_fit_mode=lambda: None) + + def _write_prediction_shard(options, frame): from pathlib import Path @@ -406,7 +411,7 @@ def fake_materialize(request, task_dataset, generated, **kwargs): metrics={}, artifact_paths={}, explain_disabled=True, - best_trainer=SimpleNamespace(log_dir=str(tmp_path)), + best_trainer=_best_trainer(tmp_path), )) monkeypatch.setattr('kumo.engine.trainer.trainer.write_training_spec', lambda *args, **kwargs: None) @@ -536,7 +541,7 @@ def fake_materialize(request, task_dataset, generated, **kwargs): metrics={}, artifact_paths={}, explain_disabled=True, - best_trainer=SimpleNamespace(log_dir=str(tmp_path)), + best_trainer=_best_trainer(tmp_path), )) monkeypatch.setattr('kumo.engine.trainer.trainer.write_training_spec', lambda *args, **kwargs: None) @@ -612,7 +617,7 @@ def fake_materialize(request, task_dataset, generated, **kwargs): metrics={}, artifact_paths={}, explain_disabled=True, - best_trainer=SimpleNamespace(log_dir=str(tmp_path)), + best_trainer=_best_trainer(tmp_path), )) monkeypatch.setattr('kumo.engine.trainer.trainer.write_training_spec', lambda *args, **kwargs: None) @@ -791,7 +796,7 @@ def fake_materialize(request, task_dataset, generated, **kwargs): metrics={}, artifact_paths={}, explain_disabled=True, - best_trainer=SimpleNamespace(log_dir=str(tmp_path)), + best_trainer=_best_trainer(tmp_path), )) monkeypatch.setattr('kumo.engine.trainer.trainer.write_training_spec', lambda *args, **kwargs: None) @@ -872,7 +877,7 @@ def fake_train(training_input, **kwargs): return SimpleNamespace( metrics={'run_name': run_name}, artifact_paths={}, explain_disabled=True, - best_trainer=SimpleNamespace(log_dir=str(tmp_path))) + best_trainer=_best_trainer(tmp_path)) monkeypatch.setattr( 'kumo.engine.trainer.trainer.build_generated_training_task_dataset', diff --git a/packages/kumo-engine/test/test_training_explanations.py b/packages/kumo-engine/test/test_training_explanations.py index 5e8ed9df1..399dd1eed 100644 --- a/packages/kumo-engine/test/test_training_explanations.py +++ b/packages/kumo-engine/test/test_training_explanations.py @@ -51,7 +51,8 @@ def explain(**kwargs): best_trainer = SimpleNamespace( explain=Mock(side_effect=explain), subgraph_collector_dir=lambda: str(tmp_path / 'subgraph_collector'), - log_dir=str(tmp_path), + get_log_dir=lambda fit_mode: str(tmp_path), + published_fit_mode=lambda: None, ) def aggregate_features(*args): From fc4d709c0987d0713bf276cae98129b84af70b51 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 27 Oct 2026 07:30:11 +0000 Subject: [PATCH 002/278] Let any sink hold files: FileStoreSink protocol and shared writers Snapshot exports, model bundles and export manifests could only go to object stores: the engine checked for ObjectStoreSinkBackend, and output configs accepted output_paths only for S3, GCS, ADLS and local connectors. A store that is not an object store (a Databricks volume, a Snowflake stage) had no way in, even with a pyarrow filesystem to offer. kumoapi.plugin.FileStoreSink is an optional sink capability: holds_files(connector, path) says from the path alone whether it is a file location rather than a table, and filesystem(connector, path) returns the pyarrow filesystem. ObjectStoreSinkBackend implements it. kumo.engine.connector.is_file_destination/file_store are the engine's one file-vs-table gate; PredictionOutputConfig(output_paths=) and the TrainingTableOutputConfig directory slash use it, so object stores behave as before. kumo_core gains FilesystemSink (the object-store file writers over any filesystem()), upload_bundle/download_bundle (the manifest file last), open_sized (tells a chunked filesystem the size) and PROVENANCE_FILE. The chunked filesystem streams a sized file of up to one part as one plain object, writes frames through a local spool so they stay plain too, spools under the engine scratch directory read at each write, and reads parts into the caller's buffer with one copy. Verified: test_chunked_fs.py (26), test_filestore_sink.py (protocol, output configs, bundles and Parquet over a filesystem-only sink), file connector, completion, preflight, registry and kumo-core export suites. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-api/kumoapi/json_serde.py | 4 +- packages/kumo-api/kumoapi/plugin.py | 18 + packages/kumo-core/src/kumo_core/artifacts.py | 4 + .../kumo-core/src/kumo_core/chunked_fs.py | 441 ++++++++++++++++++ .../src/kumo_core/object_store_sink.py | 164 +++++-- packages/kumo-core/tests/test_chunked_fs.py | 320 +++++++++++++ packages/kumo-engine/kumo/engine/_scratch.py | 13 + .../kumo/engine/connector/__init__.py | 27 ++ .../kumo/engine/trainer/output_config.py | 16 +- .../kumo-engine/test/test_filestore_sink.py | 155 ++++++ .../test/test_prediction_completion.py | 1 + 11 files changed, 1118 insertions(+), 45 deletions(-) create mode 100644 packages/kumo-core/src/kumo_core/chunked_fs.py create mode 100644 packages/kumo-core/tests/test_chunked_fs.py create mode 100644 packages/kumo-engine/test/test_filestore_sink.py diff --git a/packages/kumo-api/kumoapi/json_serde.py b/packages/kumo-api/kumoapi/json_serde.py index 8f3356019..0a62c949f 100644 --- a/packages/kumo-api/kumoapi/json_serde.py +++ b/packages/kumo-api/kumoapi/json_serde.py @@ -9,7 +9,7 @@ from decimal import Decimal from enum import Enum from pathlib import PurePath -from typing import Any, Dict, Type, TypeVar, Union +from typing import Any, Dict, Type, TypeVar, Union, cast from uuid import UUID from pydantic import BaseModel, SecretStr @@ -139,5 +139,5 @@ def from_json(obj: Any, cls: Union[Type[T], object]) -> T: if isinstance(obj, str): obj = json.loads(obj) from pydantic import TypeAdapter - adapter = TypeAdapter(cls) + adapter = TypeAdapter(cast(Any, cls)) return adapter.validate_python(obj) diff --git a/packages/kumo-api/kumoapi/plugin.py b/packages/kumo-api/kumoapi/plugin.py index 2bb9dd998..b649e302b 100644 --- a/packages/kumo-api/kumoapi/plugin.py +++ b/packages/kumo-api/kumoapi/plugin.py @@ -29,6 +29,7 @@ Mapping, Optional, Protocol, + Tuple, TypeVar, runtime_checkable, ) @@ -43,6 +44,7 @@ __all__ = [ 'ArtifactLocation', 'ArtifactSink', + 'FileStoreSink', 'IngestResult', 'ModelFetchingSink', 'ModelListingSink', @@ -604,6 +606,22 @@ def write_files( ... +@runtime_checkable +class FileStoreSink(Protocol): + r"""A sink that also holds files: snapshot exports, model bundles and + export manifests go through its filesystem.""" + def holds_files(self, connector: SourceConnector, path: str) -> bool: + r"""Whether ``path`` is a file location of this sink rather than a + table name; decided from the path alone, without I/O.""" + ... + + def filesystem(self, connector: SourceConnector, + path: str) -> Tuple[Any, str]: + r"""A ``pyarrow.fs.FileSystem`` and ``path`` resolved on it, for a + path :meth:`holds_files` accepts.""" + ... + + @runtime_checkable class ModelPublishingSink(ArtifactSink, Protocol): r"""A sink that can also receive a trained model bundle. diff --git a/packages/kumo-core/src/kumo_core/artifacts.py b/packages/kumo-core/src/kumo_core/artifacts.py index b6e7a49c4..7a4226974 100644 --- a/packages/kumo-core/src/kumo_core/artifacts.py +++ b/packages/kumo-core/src/kumo_core/artifacts.py @@ -19,6 +19,10 @@ from .io.json import load_json, save_json from .manifest import RunManifest, RunPaths +#: Every engine artifact's manifest file; underscore-prefixed so Parquet +#: readers skip it. +PROVENANCE_FILE = '_provenance.json' + @runtime_checkable class ArtifactStore(Protocol): diff --git a/packages/kumo-core/src/kumo_core/chunked_fs.py b/packages/kumo-core/src/kumo_core/chunked_fs.py new file mode 100644 index 000000000..8972badd8 --- /dev/null +++ b/packages/kumo-core/src/kumo_core/chunked_fs.py @@ -0,0 +1,441 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""A pyarrow filesystem over stores that only put, list, read and remove. + +A file larger than :attr:`ChunkedFileHandler.part_bytes` is stored as +``.chunks/00000 ... NNNNN`` followed by an empty ``_SUCCESS``; anything +smaller is one plain object. Readers see one logical file either way. A write +removes the other form; if a crash leaves both, the plain object wins. A writer +spools one part at a time on local disk, under :func:`spool_root`. A handler +that can stream an upload implements ``_stream``: a writer told the size +(:data:`SIZE_KEY` in the stream metadata) streams a file of up to +``part_bytes`` as one plain object and a larger one as parts, without local +disk; without a size, up to ``plain_bytes`` is spooled and stays plain. +""" + +import io +import os +import posixpath +import shutil +import tempfile +from typing import ( + BinaryIO, + Callable, + Dict, + List, + Optional, + Protocol, + Tuple, + Union, +) + +import pyarrow as pa +from pyarrow.fs import FileInfo, FileSystemHandler, FileType + +PART_BYTES = 256 << 20 +CHUNKS = '.chunks' +DONE = '_SUCCESS' +#: Stream metadata key carrying the size a written file will have. +SIZE_KEY = 'kumo-size' +_spool_dir: Union[None, str, Callable[[], str]] = None + + +def set_spool_dir(path: Union[None, str, Callable[[], str]]) -> None: + r"""Where writers spool parts unless a handler sets ``spool_dir``: a + path, a callable read at each write, or ``None`` (tempfile's).""" + global _spool_dir + _spool_dir = path + + +class Part(Protocol): + def read_at(self, offset: int, length: int) -> bytes: + ... + + def close(self) -> None: + ... + + +def spool_root(handler: 'ChunkedFileHandler') -> Optional[str]: + r"""Where ``handler`` stages local copies: its ``spool_dir``, else the + :func:`set_spool_dir` directory (created), else None (tempfile's).""" + root = handler.spool_dir or (_spool_dir() + if callable(_spool_dir) else _spool_dir) + if root is not None: + os.makedirs(root, exist_ok=True) + return root + + +class ChunkedFileHandler(FileSystemHandler): + r"""Subclasses supply the five store primitives below, and optionally + ``_stream``.""" + part_bytes = PART_BYTES + spool_dir: Optional[str] = None + #: Largest plain file when ``_stream`` is implemented (<= part_bytes). + plain_bytes = PART_BYTES + + def _stream(self, path: str) -> BinaryIO: + r"""Override to stream parts: a writable whose ``close`` commits the + object at ``path``.""" + raise NotImplementedError + + def _size(self, path: str) -> Optional[int]: + r"""Size of exactly the object at ``path``, or ``None``.""" + raise NotImplementedError + + def _list(self, directory: str) -> Dict[str, int]: + r"""Every object below ``directory/``, by full path.""" + raise NotImplementedError + + def _put(self, local: str, path: str) -> None: + r"""Uploads ``local``, whose basename matches ``path``'s.""" + raise NotImplementedError + + def _part(self, path: str, size: int) -> Part: + raise NotImplementedError + + def _remove(self, path: str, *, tree: bool) -> None: + r"""Removes the object at ``path``, or everything below it.""" + raise NotImplementedError + + def normalize_path(self, path): + return path.rstrip('/') + + def _layout(self, path: str) -> Optional[List[Tuple[str, int]]]: + cache = self.__dict__.setdefault('_layouts', {}) + layout = cache.get(path) + if layout is not None: + return layout + size = self._size(path) + layout = ([(path, size)] if size is not None else _logical_parts( + self._list(path + CHUNKS), path + CHUNKS)) + if layout is not None: + cache[path] = layout + return layout + + def _forget(self, path: str) -> None: + r"""Drops cached layouts at or below ``path``; writes and deletes + through this handler call it.""" + cache = self.__dict__.get('_layouts', {}) + path = path.rstrip('/') + for name in list(cache): + if name == path or name.startswith(path + '/'): + cache.pop(name, None) + + def get_file_info(self, paths): + infos = [] + for path in paths: + layout = self._layout(path) + if layout is not None: + infos.append( + FileInfo(path, FileType.File, + size=sum(size for _, size in layout))) + elif self._list(path): + infos.append(FileInfo(path, FileType.Directory)) + else: + infos.append(FileInfo(path, FileType.NotFound)) + return infos + + def get_file_info_selector(self, selector): + base = selector.base_dir.rstrip('/') + files = _logical_files(self._list(base)) + if not files and not selector.allow_not_found: + raise FileNotFoundError(f'{base} does not exist.') + infos: Dict[str, FileInfo] = {} + for path, size in files.items(): + parts = path[len(base) + 1:].split('/') + depths = len(parts) if selector.recursive else min(len(parts), 2) + for depth in range(1, depths): + directory = f"{base}/{'/'.join(parts[:depth])}" + infos[directory] = FileInfo(directory, FileType.Directory) + if selector.recursive or len(parts) == 1: + infos[path] = FileInfo(path, FileType.File, size=size) + return [infos[path] for path in sorted(infos)] + + def create_dir(self, path, recursive): + pass + + def delete_dir(self, path): + path = path.rstrip('/') + self._forget(path) + self._remove(path, tree=True) + + def delete_dir_contents(self, path, missing_dir_ok=False): + path = path.rstrip('/') + self._forget(path) + if not missing_dir_ok and not self._list(path): + raise FileNotFoundError(f'{path} does not exist.') + self._remove(path, tree=True) + + def delete_root_dir_contents(self): + raise NotImplementedError('Refusing to delete the whole store.') + + def delete_file(self, path): + self._forget(path) + plain = self._size(path) is not None + chunked = bool(self._list(path + CHUNKS)) + if not plain and not chunked: + raise FileNotFoundError(f'{path} does not exist.') + if plain: + self._remove(path, tree=False) + if chunked: + self._remove(path + CHUNKS, tree=True) + + def move(self, src, dest): + raise NotImplementedError(f'{self.get_type_name()} cannot rename.') + + def copy_file(self, src, dest): + raise NotImplementedError(f'{self.get_type_name()} cannot copy.') + + def open_append_stream(self, path, metadata): + raise NotImplementedError(f'{self.get_type_name()} cannot append.') + + def open_input_stream(self, path): + return self.open_input_file(path) + + def open_input_file(self, path): + layout = self._layout(path) + if layout is None: + raise FileNotFoundError(f'{path} does not exist.') + return pa.PythonFile(_Reader(self, layout), mode='r') + + def open_output_stream(self, path, metadata): + self._forget(path) + size = (metadata or {}).get(SIZE_KEY) + return pa.PythonFile( + _Writer(self, path, None if size is None else int(size)), + mode='w') + + +def _logical_parts(objects: Dict[str, int], + directory: str) -> Optional[List[Tuple[str, int]]]: + if f'{directory}/{DONE}' not in objects: + return None + return sorted((path, size) for path, size in objects.items() + if posixpath.basename(path) != DONE) + + +def _logical_files(objects: Dict[str, int]) -> Dict[str, int]: + r"""Collapses committed ``.chunks`` directories into their files.""" + files: Dict[str, int] = {} + chunked: Dict[str, Dict[str, int]] = {} + for path, size in objects.items(): + directory = posixpath.dirname(path) + if directory.endswith(CHUNKS): + chunked.setdefault(directory, {})[path] = size + else: + files[path] = size + for directory, parts in chunked.items(): + name = directory[:-len(CHUNKS)] + layout = _logical_parts(parts, directory) + if layout is not None and name not in files: + files[name] = sum(size for _, size in layout) + return files + + +class _Reader(io.RawIOBase): + def __init__(self, handler: ChunkedFileHandler, + layout: List[Tuple[str, int]]) -> None: + self._handler = handler + self._layout: List[Tuple[str, int, int]] = [] + start = 0 + for path, size in layout: + self._layout.append((path, start, size)) + start += size + self._size = start + self._open: Dict[str, Part] = {} + self._position = 0 + + def readable(self): + return True + + def seekable(self): + return True + + def tell(self): + return self._position + + def seek(self, offset, whence=io.SEEK_SET): + base = { + io.SEEK_SET: 0, + io.SEEK_CUR: self._position, + io.SEEK_END: self._size + }[whence] + self._position = max(base + offset, 0) + return self._position + + def _pieces(self, end: int): + for path, start, length in self._layout: + if start + length <= self._position or start >= end: + continue + offset = max(self._position, start) - start + want = min(end, start + length) - start - offset + part = self._open.get(path) + if part is None: + part = self._open[path] = self._handler._part(path, length) + chunk = part.read_at(offset, want) + if len(chunk) != want: + raise OSError(f'{path} returned {len(chunk)} of {want} bytes') + yield chunk + + def read(self, size=-1): + end = self._size if size is None or size < 0 else min( + self._position + size, self._size) + if end <= self._position: + return b'' + pieces = self._pieces(end) + first = next(pieces) + if len(first) == end - self._position: + self._position = end + return first + out = bytearray(end - self._position) + out[:len(first)] = first + filled = len(first) + for chunk in pieces: + out[filled:filled + len(chunk)] = chunk + filled += len(chunk) + self._position = end + return out + + def readall(self): + return self.read() + + def readinto(self, buffer): + view = memoryview(buffer).cast('B') + end = min(self._position + len(view), self._size) + filled = 0 + for chunk in self._pieces(end): + view[filled:filled + len(chunk)] = chunk + filled += len(chunk) + self._position += filled + return filled + + def close(self): + if not self.closed: + for part in self._open.values(): + part.close() + super().close() + + +class _Writer(io.RawIOBase): + def __init__(self, handler: ChunkedFileHandler, path: str, + size: Optional[int] = None) -> None: + self._handler = handler + self._path = path + self._dir = tempfile.mkdtemp(prefix='kumo-chunked-', + dir=spool_root(handler)) + self._spool = open(os.path.join(self._dir, 'spool'), 'wb') + self._filled = 0 + self._parts = 0 + self._written = 0 + self._part: Optional[BinaryIO] = None + self._streams = type(handler)._stream is not ChunkedFileHandler._stream + self._limit = (min(handler.plain_bytes, handler.part_bytes) + if self._streams else handler.part_bytes) + self._plain: Optional[BinaryIO] = None + if self._streams and size is not None: + if size <= handler.part_bytes: + self._plain = handler._stream(path) + else: + self._clear() + self._part = self._open_part() + + def writable(self): + return True + + def tell(self): + return self._written + + def write(self, data): + view = memoryview(data).cast('B') + total = len(view) + if self._plain is not None: + self._plain.write(view) + self._written += total + return total + while view: + if self._part is not None: + if self._filled == self._handler.part_bytes: + self._next_stream() + take = min(len(view), self._handler.part_bytes - self._filled) + self._part.write(view[:take]) + else: + if self._filled == self._limit: + if self._streams: + self._start_stream() + continue + self._flush_part() + take = min(len(view), self._limit - self._filled) + self._spool.write(view[:take]) + self._filled += take + view = view[take:] + self._written += total + return total + + def _upload(self, target: str) -> None: + self._spool.close() + local = os.path.join(self._dir, posixpath.basename(target)) + os.replace(os.path.join(self._dir, 'spool'), local) + try: + self._handler._put(local, target) + finally: + os.remove(local) + self._spool = open(os.path.join(self._dir, 'spool'), 'wb') + self._filled = 0 + + def _clear(self) -> None: + self._handler._remove(self._path + CHUNKS, tree=True) + if self._handler._size(self._path) is not None: + self._handler._remove(self._path, tree=False) + + def _flush_part(self) -> None: + if not self._parts: + self._clear() + self._upload(f'{self._path}{CHUNKS}/{self._parts:05d}') + self._parts += 1 + + def _open_part(self) -> BinaryIO: + return self._handler._stream( + f'{self._path}{CHUNKS}/{self._parts:05d}') + + def _start_stream(self) -> None: + self._clear() + self._spool.close() + spooled = os.path.join(self._dir, 'spool') + self._part = self._open_part() + with open(spooled, 'rb') as reader: + shutil.copyfileobj(reader, self._part, 8 << 20) + self._spool = open(spooled, 'wb') + + def _next_stream(self) -> None: + assert self._part is not None + self._part.close() + self._parts += 1 + self._part = self._open_part() + self._filled = 0 + + def close(self): + if self.closed: + return + try: + if self._plain is not None: + self._plain.close() + self._plain = None + if self._handler._list(self._path + CHUNKS): + self._handler._remove(self._path + CHUNKS, tree=True) + elif self._part is not None: + self._part.close() + self._part = None + self._upload(f'{self._path}{CHUNKS}/{DONE}') + elif self._parts: + self._flush_part() + self._upload(f'{self._path}{CHUNKS}/{DONE}') + else: + self._upload(self._path) + if self._handler._list(self._path + CHUNKS): + self._handler._remove(self._path + CHUNKS, tree=True) + finally: + self._handler._forget(self._path) + self._spool.close() + shutil.rmtree(self._dir, ignore_errors=True) + super().close() diff --git a/packages/kumo-core/src/kumo_core/object_store_sink.py b/packages/kumo-core/src/kumo_core/object_store_sink.py index ec34f89f9..3951ea5b3 100644 --- a/packages/kumo-core/src/kumo_core/object_store_sink.py +++ b/packages/kumo-core/src/kumo_core/object_store_sink.py @@ -2,11 +2,13 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import os import posixpath import shutil from abc import ABC, abstractmethod +from dataclasses import replace from pathlib import Path -from typing import Mapping, Optional, Tuple +from typing import Any, Mapping, Optional, Tuple import pandas as pd from kumoapi.jobs import WriteMode @@ -21,6 +23,69 @@ _OBJECT_STORES = frozenset({'s3', 'gcs', 'abfs'}) +def _chunked(filesystem: FileSystem) -> bool: + from kumo_core.chunked_fs import ChunkedFileHandler + from pyarrow.fs import PyFileSystem + + return isinstance(filesystem, PyFileSystem) and isinstance( + filesystem.handler, ChunkedFileHandler) + + +def open_sized(filesystem: FileSystem, path: str, size: int) -> Any: + """``open_output_stream(path)``, telling a chunked filesystem the size + the file will have.""" + if _chunked(filesystem): + from kumo_core.chunked_fs import SIZE_KEY + + return filesystem.open_output_stream( + path, metadata={SIZE_KEY: str(size)}) + return filesystem.open_output_stream(path) + + +def upload_bundle(filesystem: FileSystem, destination: str, + bundle_dir: str) -> None: + """Copy every file of ``bundle_dir`` under ``destination``, the + provenance file last.""" + from kumo_core.artifacts import PROVENANCE_FILE + + root = Path(bundle_dir) + for source in sorted((item for item in root.rglob('*') if item.is_file()), + key=lambda item: + (item.name == PROVENANCE_FILE, item)): + target = f'{destination}/{source.relative_to(root).as_posix()}' + _prepared(filesystem, posixpath.dirname(target)) + with open(source, 'rb') as reader, open_sized( + filesystem, target, source.stat().st_size) as writer: + shutil.copyfileobj(reader, writer, _CHUNK) + + +def download_bundle(filesystem: FileSystem, source: str, + local_dir: str) -> bool: + """Replace ``local_dir`` with the files under ``source``; False if there + are none.""" + from kumo_core.sources.layout import confined + from pyarrow.fs import FileSelector, FileType + + destination = Path(local_dir) + if destination.exists(): + shutil.rmtree(destination) + destination.mkdir(parents=True) + found = False + selector = FileSelector(source, recursive=True, allow_not_found=True) + for info in filesystem.get_file_info(selector): + if info.type is not FileType.File: + continue + found = True + # Remote keys are untrusted; each relative path must stay inside + # the local bundle, even when it contains parent components. + target = destination / confined(info.path[len(source) + 1:]) + target.parent.mkdir(parents=True, exist_ok=True) + with filesystem.open_input_stream(info.path) as reader, \ + open(target, 'wb') as writer: + shutil.copyfileobj(reader, writer, _CHUNK) + return found + + class ObjectStoreSinkBackend(ABC): scheme: str provider: str @@ -31,6 +96,16 @@ def _destination(self, connector: SourceConnector, """Return the authenticated filesystem and validated bucket/key path.""" raise NotImplementedError + def holds_files(self, connector: SourceConnector, path: str) -> bool: + """Every path of an object store is a file location.""" + return True + + def filesystem(self, connector: SourceConnector, + path: str) -> Tuple[FileSystem, str]: + """The authenticated filesystem and resolved ``path``.""" + filesystem, resolved = self._destination(connector, path) + return filesystem, resolved.rstrip('/') + def write_table( self, connector: SourceConnector, @@ -98,14 +173,7 @@ def publish_model( filesystem, destination = self._destination(connector, path) destination = destination.rstrip('/') - root = Path(bundle_dir) - for source in sorted(item for item in root.rglob('*') - if item.is_file()): - target = f'{destination}/{source.relative_to(root).as_posix()}' - _prepared(filesystem, posixpath.dirname(target)) - with open(source, 'rb') as reader, \ - filesystem.open_output_stream(target) as writer: - shutil.copyfileobj(reader, writer, _CHUNK) + upload_bundle(filesystem, destination, bundle_dir) return PublishedModel( name=posixpath.basename(destination), version=destination, @@ -124,9 +192,6 @@ def fetch_model( ) -> PublishedModel: """Replace the local directory with the published bundle.""" del project_id - from kumo_core.sources.layout import confined - from pyarrow.fs import FileSelector, FileType - if location is not None: raise NotImplementedError( f"The local backend does not import a model by name from " @@ -141,25 +206,7 @@ def fetch_model( filesystem, source = self._destination(connector, path) source = source.rstrip('/') - destination = Path(local_dir) - if destination.exists(): - shutil.rmtree(destination) - destination.mkdir(parents=True) - - found = False - selector = FileSelector(source, recursive=True, allow_not_found=True) - for info in filesystem.get_file_info(selector): - if info.type is not FileType.File: - continue - found = True - # Remote keys are untrusted; each relative path must stay inside - # the local bundle, even when it contains parent components. - target = destination / confined(info.path[len(source) + 1:]) - target.parent.mkdir(parents=True, exist_ok=True) - with filesystem.open_input_stream(info.path) as reader, \ - open(target, 'wb') as writer: - shutil.copyfileobj(reader, writer, _CHUNK) - if not found: + if not download_bundle(filesystem, source, local_dir): raise ValueError( f"Nothing is published at {self.scheme}://{source}. Check the " f"path the model was exported to.") @@ -232,8 +279,21 @@ def _write_one(filesystem: FileSystem, frame: pd.DataFrame, path: str, with filesystem.open_output_stream(path) as writer: csv.write_csv(table, writer) return - pq.write_table(pa.Table.from_pandas(frame, preserve_index=False), path, - filesystem=filesystem) + table = pa.Table.from_pandas(frame, preserve_index=False) + if not _chunked(filesystem): + pq.write_table(table, path, filesystem=filesystem) + return + import tempfile + + from kumo_core.chunked_fs import spool_root + # Spool so the size is known and a file up to one part stays plain. + with tempfile.TemporaryDirectory( + dir=spool_root(filesystem.handler)) as spool: + local = os.path.join(spool, 'frame.parquet') + pq.write_table(table, local) + with open(local, 'rb') as reader, open_sized( + filesystem, path, os.path.getsize(local)) as writer: + shutil.copyfileobj(reader, writer, _CHUNK) def _prepared(filesystem: FileSystem, directory: str) -> None: @@ -309,9 +369,9 @@ def copy_file(self, source, destination, overwrite=True): # CSV conversion already applied this destination's metadata. columns = self.extra_columns if source.endswith('.parquet') else None with prepared_artifact_file(source, columns) as prepared: - with open(prepared, - 'rb') as reader, self.filesystem.open_output_stream( - destination) as writer: + with open(prepared, 'rb') as reader, open_sized( + self.filesystem, destination, + os.path.getsize(prepared)) as writer: shutil.copyfileobj(reader, writer, _CHUNK) def copy_files_threaded(self, pairs, overwrite, max_workers): @@ -329,3 +389,35 @@ def convert_csv(self, source, destination): ) with prepared_artifact_file(source, self.extra_columns) as prepared: parquet_to_csv(prepared, destination) + + +class FilesystemSink(ObjectStoreSinkBackend): + """The object-store file writers over any sink's ``filesystem()``, for + the file paths of a sink that also writes tables: paths in, bare + resolved paths out.""" + scheme = 'filesystem' + + def __init__(self, sink: Any) -> None: + self._sink = sink + self.provider = type(sink).__name__ + + def _destination(self, connector, coordinate): + filesystem, path = self._sink.filesystem(connector, coordinate) + return filesystem, path + ('/' if coordinate.endswith('/') else '') + + def _bare(self, uri: str) -> str: + return uri.removeprefix(f'{self.scheme}://') + + def write_table(self, *args, **kwargs): + return self._bare(super().write_table(*args, **kwargs)) + + def write_files(self, *args, **kwargs): + return self._bare(super().write_files(*args, **kwargs)) + + def publish_model(self, *args, **kwargs): + published = super().publish_model(*args, **kwargs) + return replace(published, uri=self._bare(published.uri)) + + def fetch_model(self, *args, **kwargs): + fetched = super().fetch_model(*args, **kwargs) + return replace(fetched, uri=self._bare(fetched.uri)) diff --git a/packages/kumo-core/tests/test_chunked_fs.py b/packages/kumo-core/tests/test_chunked_fs.py new file mode 100644 index 000000000..f93b19e60 --- /dev/null +++ b/packages/kumo-core/tests/test_chunked_fs.py @@ -0,0 +1,320 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import io +import os + +import pytest +from kumo_core import chunked_fs +from kumo_core.chunked_fs import ChunkedFileHandler +from pyarrow.fs import FileSelector, FileType, PyFileSystem + +PART = 16 + + +class _Bytes: + def __init__(self, data: bytes) -> None: + self.data = data + + def read_at(self, offset, length): + return self.data[offset:offset + length] + + def close(self): + pass + + +class MemoryHandler(ChunkedFileHandler): + part_bytes = PART + + def __init__(self): + self.objects = {} + self.spooled = [] + + def get_type_name(self): + return 'memory' + + def _size(self, path): + data = self.objects.get(path) + return None if data is None else len(data) + + def _list(self, directory): + return { + path: len(data) + for path, data in self.objects.items() + if path.startswith(directory + '/') + } + + def _put(self, local, path): + assert os.path.basename(local) == os.path.basename(path) + self.spooled.append(os.path.dirname(os.path.dirname(local))) + with open(local, 'rb') as reader: + self.objects[path] = reader.read() + + def _part(self, path, size): + return _Bytes(self.objects[path]) + + def _remove(self, path, *, tree): + for name in list(self.objects): + if name.startswith(path + '/') if tree else name == path: + del self.objects[name] + + +@pytest.fixture +def store(): + handler = MemoryHandler() + return handler, PyFileSystem(handler) + + +def _write(fs, path, data): + with fs.open_output_stream(path) as writer: + for start in range(0, len(data), 5): + writer.write(data[start:start + 5]) + + +def _read(fs, path): + with fs.open_input_stream(path) as reader: + return reader.readall() + + +@pytest.mark.parametrize('size', [0, 3, PART]) +def test_files_up_to_one_part_are_one_plain_object(store, size): + handler, fs = store + data = bytes(range(size)) + _write(fs, 'root/a.bin', data) + assert handler.objects == {'root/a.bin': data} + assert fs.get_file_info('root/a.bin').size == size + assert _read(fs, 'root/a.bin') == data + + +@pytest.mark.parametrize('size', [2 * PART + 5, 3 * PART]) +def test_larger_files_are_parts_committed_last_and_read_back_exactly( + store, size): + handler, fs = store + data = os.urandom(size) + _write(fs, 'root/a.bin', data) + assert sorted(handler.objects) == [ + 'root/a.bin.chunks/00000', 'root/a.bin.chunks/00001', + 'root/a.bin.chunks/00002', 'root/a.bin.chunks/_SUCCESS' + ] + info = fs.get_file_info('root/a.bin') + assert (info.type, info.size) == (FileType.File, len(data)) + assert _read(fs, 'root/a.bin') == data + with fs.open_input_file('root/a.bin') as reader: + assert reader.size() == len(data) + assert reader.read_at(10, PART - 4) == data[PART - 4:PART + 6] + assert reader.read_at(0, 2 * PART) == b'' + assert reader.read_at(PART, PART) == data[PART:2 * PART] + assert reader.read_at(100, PART + 1) == data[PART + 1:] + + +def test_an_uncommitted_chunked_write_is_not_a_file(store): + handler, fs = store + handler.objects['root/a.bin.chunks/00000'] = b'x' * PART + assert fs.get_file_info('root/a.bin').type is not FileType.File + with pytest.raises(FileNotFoundError): + fs.open_input_file('root/a.bin') + + +@pytest.mark.parametrize('first,second', [(3, 2 * PART), (2 * PART, 3)]) +def test_a_write_removes_the_other_form(store, first, second): + handler, fs = store + _write(fs, 'root/a.bin', b'1' * first) + _write(fs, 'root/a.bin', b'2' * second) + assert _read(fs, 'root/a.bin') == b'2' * second + plain = 'root/a.bin' in handler.objects + assert plain == (second <= PART) + assert any('.chunks/' in name for name in handler.objects) != plain + + +def test_plain_wins_when_a_crash_left_both_forms(store): + handler, fs = store + _write(fs, 'root/a.bin', b'c' * (2 * PART)) + handler.objects['root/a.bin'] = b'plain' + assert _read(fs, 'root/a.bin') == b'plain' + infos = fs.get_file_info(FileSelector('root', recursive=True)) + assert [(i.path, i.size) for i in infos] == [('root/a.bin', 5)] + + +def test_listing_shows_logical_files_and_directories(store): + _, fs = store + _write(fs, 'root/top.bin', b'1') + _write(fs, 'root/dir/big.bin', b'2' * (PART + 1)) + _write(fs, 'root/dir/sub/small.bin', b'3') + recursive = fs.get_file_info(FileSelector('root', recursive=True)) + assert [(i.path, i.type, i.size) for i in recursive] == [ + ('root/dir', FileType.Directory, None), + ('root/dir/big.bin', FileType.File, PART + 1), + ('root/dir/sub', FileType.Directory, None), + ('root/dir/sub/small.bin', FileType.File, 1), + ('root/top.bin', FileType.File, 1), + ] + shallow = fs.get_file_info(FileSelector('root')) + assert [i.path for i in shallow] == ['root/dir', 'root/top.bin'] + assert fs.get_file_info('root/dir').type is FileType.Directory + assert fs.get_file_info('root/none').type is FileType.NotFound + assert fs.get_file_info(FileSelector('nothing', allow_not_found=True)) == [] + with pytest.raises(FileNotFoundError): + fs.get_file_info(FileSelector('nothing')) + + +def test_deletes(store): + handler, fs = store + _write(fs, 'root/a.bin', b'a' * (PART + 1)) + _write(fs, 'root/b.bin', b'b') + fs.delete_file('root/a.bin') + assert sorted(handler.objects) == ['root/b.bin'] + with pytest.raises(FileNotFoundError): + fs.delete_file('root/a.bin') + fs.delete_dir_contents('root') + assert handler.objects == {} + fs.delete_dir_contents('root', missing_dir_ok=True) + with pytest.raises(FileNotFoundError): + fs.delete_dir_contents('root') + with pytest.raises(NotImplementedError): + fs.move('root/a', 'root/b') + + +def test_parts_spool_under_the_spool_dir_and_are_removed( + store, tmp_path, monkeypatch): + handler, fs = store + monkeypatch.setattr(chunked_fs, '_spool_dir', None) + chunked_fs.set_spool_dir(str(tmp_path / 'scratch')) + _write(fs, 'root/a.bin', b'a' * (2 * PART + 1)) + handler.spool_dir = str(tmp_path / 'own') + _write(fs, 'root/b.bin', b'b') + assert handler.spooled == [str(tmp_path / 'scratch')] * 4 + [ + str(tmp_path / 'own') + ] + assert not os.listdir(tmp_path / 'scratch') + assert not os.listdir(tmp_path / 'own') + + +class _Upload(io.BytesIO): + def __init__(self, handler, path): + super().__init__() + self.handler, self.path = handler, path + + def close(self): + self.handler.objects[self.path] = self.getvalue() + self.handler.streamed.append(self.path) + super().close() + + +class StreamingHandler(MemoryHandler): + plain_bytes = 6 + + def __init__(self): + super().__init__() + self.streamed = [] + + def _stream(self, path): + return _Upload(self, path) + + +@pytest.mark.parametrize('size', [0, 6, 7, PART, 2 * PART + 5, 3 * PART]) +def test_a_streaming_handler_streams_parts_and_spools_only_small_files(size): + handler = StreamingHandler() + fs = PyFileSystem(handler) + _write(fs, 'root/a.bin', b'x' * (2 * PART)) + data = os.urandom(size) + _write(fs, 'root/a.bin', data) + assert _read(fs, 'root/a.bin') == data + if size <= 6: + assert handler.objects == {'root/a.bin': data} + return + parts = -(-size // PART) + assert sorted(handler.objects) == [ + f'root/a.bin.chunks/{i:05d}' for i in range(parts) + ] + ['root/a.bin.chunks/_SUCCESS'] + assert handler.streamed[-parts:] == [ + f'root/a.bin.chunks/{i:05d}' for i in range(parts) + ] + assert len(handler.spooled) == 2 + with fs.open_input_file('root/a.bin') as reader: + assert reader.read_at(4, PART - 2) == data[PART - 2:PART + 2] + + +def test_opening_a_file_again_reuses_its_layout(store): + handler, fs = store + _write(fs, 'root/a.bin', b'a' * (2 * PART)) + probes = [] + size, listing = handler._size, handler._list + handler._size = lambda path: probes.append(path) or size(path) + handler._list = lambda path: probes.append(path) or listing(path) + for _ in range(3): + assert _read(fs, 'root/a.bin') == b'a' * (2 * PART) + assert len(probes) == 2 + _write(fs, 'root/a.bin', b'new') + assert _read(fs, 'root/a.bin') == b'new' + fs.delete_dir_contents('root/') + with pytest.raises(FileNotFoundError): + fs.open_input_file('root/a.bin') + + +def test_layout_cache_survives_concurrent_opens_and_writes(store): + import threading + + handler, fs = store + for i in range(20): + _write(fs, f'root/{i}.bin', b'x' * (PART + 1)) + errors = [] + + def work(i): + try: + for _ in range(50): + _read(fs, f'root/{i}.bin') + _write(fs, f'root/{i}.bin', b'x' * (PART + 1)) + except Exception as error: + errors.append(error) + + threads = [threading.Thread(target=work, args=(i, )) for i in range(20)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + assert errors == [] + + +@pytest.mark.parametrize('size', [7, PART, PART + 1, 3 * PART]) +def test_a_known_size_streams_without_spooling(size): + from kumo_core.object_store_sink import open_sized + + handler = StreamingHandler() + fs = PyFileSystem(handler) + _write(fs, 'root/a.bin', b'x' * (2 * PART)) + handler.spooled.clear() + data = os.urandom(size) + with open_sized(fs, 'root/a.bin', size) as writer: + for start in range(0, size, 5): + writer.write(data[start:start + 5]) + assert _read(fs, 'root/a.bin') == data + plain = 'root/a.bin' in handler.objects + assert plain == (size <= PART) + assert any('.chunks/' in name for name in handler.objects) != plain + # Only the empty _SUCCESS marker of a chunked file is put from disk. + assert len(handler.spooled) == (0 if plain else 1) + + +def test_the_spool_dir_can_be_read_at_each_write(store, tmp_path, + monkeypatch): + handler, fs = store + root = [tmp_path / 'one'] + monkeypatch.setattr(chunked_fs, '_spool_dir', None) + chunked_fs.set_spool_dir(lambda: str(root[0])) + _write(fs, 'root/a.bin', b'a') + root[0] = tmp_path / 'two' + _write(fs, 'root/b.bin', b'b') + assert handler.spooled == [str(tmp_path / 'one'), str(tmp_path / 'two')] + + +def test_reads_spanning_parts_fill_the_callers_buffer(store): + _, fs = store + data = os.urandom(3 * PART + 3) + _write(fs, 'root/a.bin', data) + with fs.open_input_file('root/a.bin') as reader: + buffer = bytearray(2 * PART) + reader.seek(PART - 3) + assert reader.readinto(buffer) == 2 * PART + assert bytes(buffer) == data[PART - 3:3 * PART - 3] + assert reader.read_at(4, 2) == data[2:6] diff --git a/packages/kumo-engine/kumo/engine/_scratch.py b/packages/kumo-engine/kumo/engine/_scratch.py index f3642e9f5..7381e5cb5 100644 --- a/packages/kumo-engine/kumo/engine/_scratch.py +++ b/packages/kumo-engine/kumo/engine/_scratch.py @@ -68,6 +68,19 @@ def get_scratch_dir() -> Path: return _root +def _spool_dir() -> str: + return str(get_scratch_dir() / "spool") + + +def _spool_parts() -> None: + from kumo_core.chunked_fs import set_spool_dir + + set_spool_dir(_spool_dir) + + +_spool_parts() + + def new_run_id() -> str: r"""A fresh, random id for one run -- a :class:`~kumo.engine.trainer.Trainer` instance, a `Trainer.predict` call, a `PredictiveQuery. diff --git a/packages/kumo-engine/kumo/engine/connector/__init__.py b/packages/kumo-engine/kumo/engine/connector/__init__.py index d328cf3ee..64c869506 100644 --- a/packages/kumo-engine/kumo/engine/connector/__init__.py +++ b/packages/kumo-engine/kumo/engine/connector/__init__.py @@ -17,6 +17,8 @@ `_response`) and the `_bind_connector` round-trip. See `kumo.engine.connector.base`. """ +from typing import Any, Optional, Tuple + from kumo.engine.connector.base import Connector from kumo.engine.connector.bigquery import BigQueryConnector from kumo.engine.connector.databricks import DatabricksConnector @@ -34,6 +36,31 @@ from kumo.engine.connector.snowflake import SnowflakeConnector from kumo.engine.connector.source_table import SourceTable, SQLSourceTable + +def _file_sink(connector: Any, path: str) -> Any: + from kumoapi.plugin import FileStoreSink + + try: + sink = connector.sink_backend + except Exception: + return None + return sink if isinstance(sink, FileStoreSink) and sink.holds_files( + connector, path) else None + + +def is_file_destination(connector: Any, path: str) -> bool: + r"""Whether ``path`` is a file location of ``connector``'s sink + (`kumoapi.plugin.FileStoreSink`) rather than a table; no I/O.""" + return _file_sink(connector, path) is not None + + +def file_store(connector: Any, path: str) -> Optional[Tuple[Any, str]]: + r"""``connector``'s ``(pyarrow filesystem, resolved path)`` for a file + destination; None for a table.""" + sink = _file_sink(connector, path) + return None if sink is None else sink.filesystem(connector, path) + + __all__ = [ "Connector", "FileConnector", diff --git a/packages/kumo-engine/kumo/engine/trainer/output_config.py b/packages/kumo-engine/kumo/engine/trainer/output_config.py index 4a325126d..43da33eac 100644 --- a/packages/kumo-engine/kumo/engine/trainer/output_config.py +++ b/packages/kumo-engine/kumo/engine/trainer/output_config.py @@ -11,10 +11,9 @@ from kumoapi.plugin import ModelRegistryLocation, SourceConnector from kumoapi.prediction_outputs import resolve_prediction_output_paths +from kumo.engine.connector import is_file_destination from kumo.engine.connector.base import Connector -_FILE_SOURCES = frozenset({'S3', 'GCS', 'ADLS', 'LOCAL_FILE'}) - @dataclass(frozen=True) class PredictionOutputConfig: @@ -36,12 +35,15 @@ class PredictionOutputConfig: def __post_init__(self) -> None: if self.output_paths is not None: - if self.connector.source_type not in _FILE_SOURCES: - raise ValueError( - 'output_paths requires a file or object-store connector') paths = resolve_prediction_output_paths( self.output_paths, output_types=self.artifact_type, output_table_name=self.table_name) + if not all( + is_file_destination(self.connector, path) + for path in paths.values()): + raise ValueError( + 'output_paths requires a file or object-store connector ' + '(or a volume or stage path it holds)') sink = self.connector.sink_backend if isinstance(sink, ObjectStoreSessions): # Reject aliases; retain caller paths. @@ -75,8 +77,8 @@ class TrainingTableOutputConfig: write_mode: WriteMode = WriteMode.OVERWRITE def __post_init__(self) -> None: - if (self.connector.source_type in _FILE_SOURCES - and not self.table_name.endswith('/')): + if (not self.table_name.endswith('/') + and is_file_destination(self.connector, self.table_name)): object.__setattr__(self, 'table_name', self.table_name + '/') diff --git a/packages/kumo-engine/test/test_filestore_sink.py b/packages/kumo-engine/test/test_filestore_sink.py new file mode 100644 index 000000000..21a59362d --- /dev/null +++ b/packages/kumo-engine/test/test_filestore_sink.py @@ -0,0 +1,155 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import os +from pathlib import Path +from types import SimpleNamespace + +import pandas as pd +import pytest +from kumo_core.chunked_fs import ChunkedFileHandler +from kumo_core.object_store_sink import FilesystemSink +from kumoapi.jobs import PredictionArtifactType +from kumoapi.plugin import FileStoreSink +from pyarrow.fs import PyFileSystem + +import kumo.engine as kumo +from kumo.engine._scratch import get_scratch_dir +from kumo.engine.connector import file_store, is_file_destination +from kumo.engine.trainer.output_config import ( + PredictionOutputConfig, + TrainingTableOutputConfig, +) + + +class _Bytes: + def __init__(self, data): + self.data = data + + def read_at(self, offset, length): + return self.data[offset:offset + length] + + def close(self): + pass + + +class MemoryStore(ChunkedFileHandler): + part_bytes = 4096 + + def __init__(self): + self.objects, self.spooled = {}, set() + + def get_type_name(self): + return 'memory' + + def _size(self, path): + data = self.objects.get(path) + return None if data is None else len(data) + + def _list(self, directory): + return { + path: len(data) + for path, data in self.objects.items() + if path.startswith(directory + '/') + } + + def _put(self, local, path): + self.spooled.add(Path(local).parent.parent) + self.objects[path] = Path(local).read_bytes() + + def _part(self, path, size): + return _Bytes(self.objects[path]) + + def _remove(self, path, *, tree): + for name in list(self.objects): + if name.startswith(path + '/') if tree else name == path: + del self.objects[name] + + +class StageSink: + r"""Not an object store: holds '@stage/...' paths through + ``filesystem``; anything else is a table.""" + def __init__(self): + self.store = MemoryStore() + + def holds_files(self, connector, path): + return path.startswith('@') + + def filesystem(self, connector, path): + assert self.holds_files(connector, path) + return PyFileSystem(self.store), path.rstrip('/') + + def write_files(self, *args, **kwargs): + return FilesystemSink(self).write_files(*args, **kwargs) + + +@pytest.fixture +def connector(tmp_path): + previous = get_scratch_dir() + kumo.configure(scratch_dir=tmp_path / 'scratch') + yield SimpleNamespace(source_type='SNOWFLAKE', sink_backend=StageSink()) + kumo.configure(scratch_dir=previous) + + +def test_sinks_with_a_filesystem_hold_files(tmp_path, connector): + local = kumo.FileConnector(root_dir=str(tmp_path)) + assert isinstance(connector.sink_backend, FileStoreSink) + assert isinstance(local.sink_backend, FileStoreSink) + assert is_file_destination(local, 'any/path') + assert is_file_destination(connector, '@stage/p') + assert not is_file_destination(connector, 'DB.SCHEMA.T') + assert file_store(connector, 'DB.SCHEMA.T') is None + assert not is_file_destination(SimpleNamespace(), 'x') + + +def test_output_configs_take_a_sinks_file_paths(connector): + assert TrainingTableOutputConfig(connector, '@stage/tt').table_name == ( + '@stage/tt/') + assert TrainingTableOutputConfig(connector, 'TT').table_name == 'TT' + PredictionOutputConfig(connector, + output_paths={'predictions': '@stage/p/'}) + with pytest.raises(ValueError, match='output_paths requires'): + PredictionOutputConfig(connector, output_paths={'predictions': 'T'}) + PredictionOutputConfig(connector, PredictionArtifactType.PREDICTIONS, + 'PREDS') + + +def test_bundles_round_trip_through_a_filesystem_only_sink( + tmp_path, connector): + bundle = tmp_path / 'bundle' + (bundle / 'sub').mkdir(parents=True) + (bundle / '_provenance.json').write_text('{}') + (bundle / 'sub' / 'weights.bin').write_bytes(os.urandom(10000)) + files = FilesystemSink(connector.sink_backend) + published = files.publish_model(connector, str(bundle), + path='@stage/models/m') + assert published.uri == '@stage/models/m' + store = connector.sink_backend.store + assert list(store.objects)[-1] == '@stage/models/m/_provenance.json' + assert '@stage/models/m/sub/weights.bin.chunks/_SUCCESS' in store.objects + files.fetch_model(connector, str(tmp_path / 'fetched'), + path='@stage/models/m') + for name in ('_provenance.json', 'sub/weights.bin'): + assert (tmp_path / 'fetched' / name).read_bytes() == ( + bundle / name).read_bytes() + + +def test_parquet_writes_stay_plain_up_to_one_part(tmp_path, connector): + files = FilesystemSink(connector.sink_backend) + frame = pd.DataFrame({'x': range(10)}) + files.write_table(connector, '@stage/small.parquet', frame) + files.write_table(connector, '@stage/big.parquet', + pd.DataFrame({'x': [os.urandom(50).hex() for _ in range(100)]})) + store = connector.sink_backend.store + assert '@stage/small.parquet' in store.objects + assert '@stage/big.parquet.chunks/_SUCCESS' in store.objects + assert pd.read_parquet( + '@stage/big.parquet', + filesystem=PyFileSystem(store)).shape == (100, 1) + + part = tmp_path / 'part-0.parquet' + frame.to_parquet(part) + assert files.write_files(connector, '@stage/out/', [str(part)]) + assert pd.read_parquet('@stage/out/part-0.parquet', + filesystem=PyFileSystem(store)).equals(frame) diff --git a/packages/kumo-engine/test/test_prediction_completion.py b/packages/kumo-engine/test/test_prediction_completion.py index a1ad9a53f..526d87a3e 100644 --- a/packages/kumo-engine/test/test_prediction_completion.py +++ b/packages/kumo-engine/test/test_prediction_completion.py @@ -261,6 +261,7 @@ def test_export_missing_artifact_does_not_publish_partial_outputs( writes = [] connector = SimpleNamespace( source_type='LOCAL_FILE', sink_backend=SimpleNamespace( + holds_files=lambda *a: True, filesystem=lambda *a: None, write_files=lambda *a, **kw: writes.append(a))) config = PredictionOutputConfig( connector, output_paths={ From 86fa3c6c46ba5c0dc521ad4a52547fd50bd92d18 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 22 Oct 2026 00:34:55 +0000 Subject: [PATCH 003/278] Expose published_fit_mode on the artifact test's fake trainer test_artifacts.py's fake trainer borrows Trainer methods; select_fit_mode now calls published_fit_mode, so the fake needs it too. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-core/tests/test_artifacts.py | 1 + 1 file changed, 1 insertion(+) diff --git a/packages/kumo-core/tests/test_artifacts.py b/packages/kumo-core/tests/test_artifacts.py index 35c2cef65..071f4339b 100644 --- a/packages/kumo-core/tests/test_artifacts.py +++ b/packages/kumo-core/tests/test_artifacts.py @@ -199,6 +199,7 @@ def test_vital_source_resolution_is_core_owned(tmp_path: Path, refit) -> None: class SourceTrainer(SimpleNamespace): _set_fit_mode_to_use = Trainer._set_fit_mode_to_use + published_fit_mode = Trainer.published_fit_mode @property def best_ckpt(self): From 551e31f1f9bb0a66f57a3e35398e3002f3953393 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 27 Oct 2026 23:49:53 +0000 Subject: [PATCH 004/278] Record provenance for every artifact under one MLOps artifact root Co-Authored-By: Claude Opus 5.5 --- 3rdparty/notices/inventory.json | 4 +- THIRD-PARTY.txt | 2 +- packages/diskgraph/Cargo.lock | 1 + .../diskgraph/diskgraph-ingest/Cargo.toml | 1 + .../diskgraph-ingest/src/snapshot_mirror.rs | 134 ++- packages/diskgraph/diskgraph-pybind/build.rs | 29 + .../python/diskgraph/__init__.py | 4 + .../python/tests/test_exception_exports.py | 24 + .../diskgraph/diskgraph-pybind/src/lib.rs | 18 + .../third_party/artifact-manifest.json | 9 +- .../src/kumo_bigquery/identity.py | 10 +- .../src/kumo_bigquery/session.py | 5 + packages/kumo-bigquery/tests/test_identity.py | 3 + .../src/kumo_core/artifact_export.py | 18 +- packages/kumo-core/src/kumo_core/artifacts.py | 109 ++- .../kumo-core/src/kumo_core/io/filesystem.py | 32 + .../src/kumo_core/object_store_sink.py | 65 +- .../src/kumo_core/sources/identity.py | 35 +- .../src/kumo_core/sources/object_store.py | 16 +- .../kumo_core/sources/object_store_backend.py | 2 +- .../sources/object_store_sessions.py | 13 +- .../src/kumo_core/train/_artifact_storage.py | 19 +- .../sources/parquet/test_object_store.py | 4 +- .../tests/sources/unit/test_sql_identity.py | 6 + .../kumo-core/tests/test_artifact_export.py | 32 + packages/kumo-core/tests/test_artifacts.py | 26 + .../src/kumo_databricks/identity.py | 8 +- .../src/kumo_databricks/session.py | 5 + .../kumo-databricks/tests/test_identity.py | 2 + packages/kumo-engine/README.md | 1 + packages/kumo-engine/docs/mlops.md | 88 ++ packages/kumo-engine/kumo/engine/__init__.py | 13 + packages/kumo-engine/kumo/engine/_manifest.py | 290 +++++++ packages/kumo-engine/kumo/engine/_mlops.py | 105 +++ .../kumo-engine/kumo/engine/_provenance.py | 137 +++ .../kumo-engine/kumo/engine/_snapshot_io.py | 557 +++++++++++++ .../kumo/engine/connector/__init__.py | 19 + .../kumo-engine/kumo/engine/connector/file.py | 5 + packages/kumo-engine/kumo/engine/graph.py | 374 ++++++++- .../kumo/engine/pquery/_custom_source.py | 8 +- .../kumo/engine/pquery/prediction_table.py | 3 +- .../kumo/engine/pquery/predictive_query.py | 77 +- .../kumo/engine/pquery/training_table.py | 182 +++- .../kumo/engine/trainer/__init__.py | 8 +- .../kumo/engine/trainer/_provenance.py | 86 ++ .../kumo-engine/kumo/engine/trainer/_spec.py | 90 +- .../kumo/engine/trainer/output_config.py | 9 + .../kumo/engine/trainer/trainer.py | 273 +++++- packages/kumo-engine/setup.py | 30 + packages/kumo-engine/test/conftest.py | 28 + .../test/connector/test_custom_source.py | 17 +- .../test/test_bundle_provenance.py | 164 ++++ packages/kumo-engine/test/test_custom_sql.py | 2 +- .../kumo-engine/test/test_filestore_sink.py | 25 + .../kumo-engine/test/test_graph_mismatch.py | 89 ++ packages/kumo-engine/test/test_holdout_df.py | 32 +- packages/kumo-engine/test/test_mlops.py | 86 ++ .../kumo-engine/test/test_predictive_query.py | 6 +- .../test/test_query_from_snapshot.py | 123 +++ .../test/test_semantic_carry_forward.py | 7 +- .../kumo-engine/test/test_snapshot_export.py | 786 ++++++++++++++++++ .../test/test_snapshot_export_live.py | 45 + .../test/test_snapshot_manifest.py | 256 ++++++ .../kumo-engine/test/test_trainer_registry.py | 16 +- .../test/test_training_table_export.py | 177 +++- packages/kumo-engine/test/test_warm_start.py | 2 + .../src/kumo_snowflake/identity.py | 9 +- .../src/kumo_snowflake/session.py | 10 + .../kumo-snowflake/tests/test_identity.py | 2 + 69 files changed, 4595 insertions(+), 278 deletions(-) create mode 100644 packages/kumo-engine/docs/mlops.md create mode 100644 packages/kumo-engine/kumo/engine/_manifest.py create mode 100644 packages/kumo-engine/kumo/engine/_mlops.py create mode 100644 packages/kumo-engine/kumo/engine/_provenance.py create mode 100644 packages/kumo-engine/kumo/engine/_snapshot_io.py create mode 100644 packages/kumo-engine/kumo/engine/trainer/_provenance.py create mode 100644 packages/kumo-engine/test/test_bundle_provenance.py create mode 100644 packages/kumo-engine/test/test_graph_mismatch.py create mode 100644 packages/kumo-engine/test/test_mlops.py create mode 100644 packages/kumo-engine/test/test_query_from_snapshot.py create mode 100644 packages/kumo-engine/test/test_snapshot_export.py create mode 100644 packages/kumo-engine/test/test_snapshot_export_live.py create mode 100644 packages/kumo-engine/test/test_snapshot_manifest.py diff --git a/3rdparty/notices/inventory.json b/3rdparty/notices/inventory.json index d0dc3d7bc..c79d73931 100644 --- a/3rdparty/notices/inventory.json +++ b/3rdparty/notices/inventory.json @@ -515,6 +515,7 @@ "rust:cpp-demangle@0.4.5", "rust:cpufeatures@0.2.17", "rust:cpufeatures@0.3.0", + "rust:crc32fast@1.5.0", "rust:crossbeam-deque@0.8.6", "rust:crossbeam-epoch@0.9.21", "rust:crossbeam-utils@0.8.21", @@ -3642,7 +3643,8 @@ }, { "artifact_ids": [ - "diskgraph-extension" + "diskgraph-extension", + "diskgraph-ingest" ], "copyright": [ "Copyright (c) 2018 Sam Rijs, Alex Crichton and contributors" diff --git a/THIRD-PARTY.txt b/THIRD-PARTY.txt index 1394a494e..df85abc32 100644 --- a/THIRD-PARTY.txt +++ b/THIRD-PARTY.txt @@ -5152,7 +5152,7 @@ cpufeatures 0.3.0 crc32fast 1.5.0 ------------------------------------------------------------------------------ License : MIT OR Apache-2.0 - Distributed in: diskgraph._c.abi3.so + Distributed in: diskgraph._c.abi3.so, diskgraph/bin/ingest Copyright : Copyright (c) 2018 Sam Rijs, Alex Crichton and contributors Legal text c6596eb7be85: LICENSE-APACHE diff --git a/packages/diskgraph/Cargo.lock b/packages/diskgraph/Cargo.lock index be75fbfdd..233fec7c1 100644 --- a/packages/diskgraph/Cargo.lock +++ b/packages/diskgraph/Cargo.lock @@ -1739,6 +1739,7 @@ dependencies = [ "bytes", "chrono", "clap", + "crc32fast", "crossterm", "dashmap", "dirs", diff --git a/packages/diskgraph/diskgraph-ingest/Cargo.toml b/packages/diskgraph/diskgraph-ingest/Cargo.toml index 1f29515e3..d068650d2 100644 --- a/packages/diskgraph/diskgraph-ingest/Cargo.toml +++ b/packages/diskgraph/diskgraph-ingest/Cargo.toml @@ -55,6 +55,7 @@ epserde = { workspace = true } memmap2 = "0.9" libc = "0.2" rayon = "1" +crc32fast = "1.5" jemallocator = { version = "0.5", optional = true, features = ["profiling"] } jemalloc-ctl = { version = "0.5", optional = true } sysinfo = "0.30" diff --git a/packages/diskgraph/diskgraph-ingest/src/snapshot_mirror.rs b/packages/diskgraph/diskgraph-ingest/src/snapshot_mirror.rs index 7ea4edbcc..bfba23cfb 100644 --- a/packages/diskgraph/diskgraph-ingest/src/snapshot_mirror.rs +++ b/packages/diskgraph/diskgraph-ingest/src/snapshot_mirror.rs @@ -36,6 +36,7 @@ use anyhow::{Context, Result}; use object_store::path::Path as ObjectPath; use parking_lot::Mutex; use serde::{Deserialize, Serialize}; +use std::collections::BTreeMap; use diskgraph_fs::{CloudTransfer, CloudTransferTicket}; @@ -488,21 +489,9 @@ pub(crate) fn validate_manifest_rel_path(rel: &str) -> Result<(), &'static str> Ok(()) } -/// Verify the snapshot rooted at `output_dir` against its `_MANIFEST.json`. -/// Used by the consumer side after a Python-driven dataplane→local copy -/// completes, to catch lifecycle-mid-delete races and partial copies -/// before the snapshot is handed to a workload. -/// -/// Pass on success; [`CorruptionError`] on any deviation from the -/// manifest contract. Errors are precise about what failed so the -/// caller's alerting carries actionable detail (which file is missing, -/// expected vs actual size, etc.). -pub fn verify_snapshot_integrity(output_dir: &std::path::Path) -> std::result::Result<(), CorruptionError> { - let success_path = output_dir.join(SUCCESS_MARKER_FILENAME); - if !success_path.exists() { - return Err(CorruptionError::MissingSuccess(success_path)); - } - +/// Read and validate `_MANIFEST.json` at `output_dir`: schema version and +/// every rel_path's confinement to the snapshot root. +fn read_manifest(output_dir: &std::path::Path) -> std::result::Result { let manifest_path = output_dir.join(MANIFEST_FILENAME); let manifest_bytes = match std::fs::read(&manifest_path) { Ok(b) => b, @@ -548,6 +537,25 @@ pub fn verify_snapshot_integrity(output_dir: &std::path::Path) -> std::result::R }); } } + Ok(manifest) +} + +/// Verify the snapshot rooted at `output_dir` against its `_MANIFEST.json`. +/// Used by the consumer side after a Python-driven dataplane→local copy +/// completes, to catch lifecycle-mid-delete races and partial copies +/// before the snapshot is handed to a workload. +/// +/// Pass on success; [`CorruptionError`] on any deviation from the +/// manifest contract. Errors are precise about what failed so the +/// caller's alerting carries actionable detail (which file is missing, +/// expected vs actual size, etc.). +pub fn verify_snapshot_integrity(output_dir: &std::path::Path) -> std::result::Result<(), CorruptionError> { + let success_path = output_dir.join(SUCCESS_MARKER_FILENAME); + if !success_path.exists() { + return Err(CorruptionError::MissingSuccess(success_path)); + } + + let manifest = read_manifest(output_dir)?; let mut missing: Vec = Vec::new(); let mut size_mismatches: Vec<(String, u64, u64)> = Vec::new(); @@ -584,6 +592,71 @@ pub fn verify_snapshot_integrity(output_dir: &std::path::Path) -> std::result::R Ok(()) } +/// CRC-32 of every `range_bytes` span of each file `_MANIFEST.json` lists, +/// and of `_MANIFEST.json` itself, keyed by rel_path (an empty file has one +/// zero entry). Lets a copy of the snapshot be checked for content as well +/// as size; the manifest format is unchanged. +pub fn snapshot_range_checksums( + output_dir: &std::path::Path, + range_bytes: u64, +) -> std::result::Result>, CorruptionError> { + use rayon::prelude::*; + use std::os::unix::fs::FileExt; + + const CHUNK: u64 = 1 << 20; + if range_bytes == 0 { + return Err(CorruptionError::Io { + path: output_dir.to_path_buf(), + detail: "range_bytes must be positive".into(), + }); + } + let mut files = read_manifest(output_dir)?.files; + let manifest_size = std::fs::metadata(output_dir.join(MANIFEST_FILENAME)) + .map(|meta| meta.len()) + .unwrap_or(0); + files.push(ManifestEntry { + rel_path: MANIFEST_FILENAME.into(), + size_bytes: manifest_size, + }); + let ranges: Vec<(usize, u64)> = files + .iter() + .enumerate() + .flat_map(|(index, entry)| { + let count = entry.size_bytes.div_ceil(range_bytes).max(1); + (0..count).map(move |range| (index, range)) + }) + .collect(); + let crcs = ranges + .par_iter() + .map(|&(index, range)| { + let entry = &files[index]; + let path = output_dir.join(&entry.rel_path); + let io = |e: std::io::Error| CorruptionError::Io { + path: path.clone(), + detail: e.to_string(), + }; + let file = std::fs::File::open(&path).map_err(io)?; + let start = range * range_bytes; + let end = entry.size_bytes.min(start + range_bytes); + let mut hasher = crc32fast::Hasher::new(); + let mut buffer = vec![0u8; CHUNK.min(end - start) as usize]; + let mut offset = start; + while offset < end { + let length = CHUNK.min(end - offset) as usize; + file.read_exact_at(&mut buffer[..length], offset).map_err(io)?; + hasher.update(&buffer[..length]); + offset += length as u64; + } + Ok(hasher.finalize()) + }) + .collect::, CorruptionError>>()?; + let mut out: BTreeMap> = BTreeMap::new(); + for (&(index, _), crc) in ranges.iter().zip(crcs) { + out.entry(files[index].rel_path.clone()).or_default().push(crc); + } + Ok(out) +} + #[cfg(test)] mod tests { use super::*; @@ -933,4 +1006,35 @@ mod tests { verify_snapshot_integrity(local_root.path())?; Ok(()) } + + #[test] + fn range_checksums_cover_every_listed_file_and_the_manifest() { + let tmp = tempfile::tempdir().unwrap(); + let big: Vec = (0..10u32).flat_map(|i| i.to_le_bytes()).collect(); + write_marker_files(tmp.path(), &[("a.bin", &big), ("nested/empty", b"")]); + write_manifest( + tmp.path(), + vec![ + ManifestEntry { + rel_path: "a.bin".into(), + size_bytes: 40, + }, + ManifestEntry { + rel_path: "nested/empty".into(), + size_bytes: 0, + }, + ], + ); + let crcs = snapshot_range_checksums(tmp.path(), 16).unwrap(); + let manifest = std::fs::read(tmp.path().join(MANIFEST_FILENAME)).unwrap(); + assert_eq!(crcs["a.bin"], [&big[..16], &big[16..32], &big[32..]].map(crc32fast::hash)); + assert_eq!(crcs["nested/empty"], vec![0]); + assert_eq!(crcs[MANIFEST_FILENAME], manifest.chunks(16).map(crc32fast::hash).collect::>()); + + std::fs::write(tmp.path().join("a.bin"), &big[..20]).unwrap(); + assert!(matches!( + snapshot_range_checksums(tmp.path(), 16).unwrap_err(), + CorruptionError::Io { .. } + )); + } } diff --git a/packages/diskgraph/diskgraph-pybind/build.rs b/packages/diskgraph/diskgraph-pybind/build.rs index 9bd833888..9b3c22cf0 100644 --- a/packages/diskgraph/diskgraph-pybind/build.rs +++ b/packages/diskgraph/diskgraph-pybind/build.rs @@ -9,8 +9,37 @@ mod native_cpu_guard { )); } +fn build_commit() -> String { + println!("cargo:rerun-if-env-changed=DISKGRAPH_BUILD_COMMIT"); + if let Ok(commit) = std::env::var("DISKGRAPH_BUILD_COMMIT") { + return commit; + } + let git = |args: &[&str]| { + std::process::Command::new("git") + .args(args) + .output() + .ok() + .filter(|out| out.status.success()) + .map(|out| String::from_utf8_lossy(&out.stdout).trim().to_owned()) + .filter(|out| !out.is_empty()) + }; + for name in ["HEAD", "packed-refs"] { + if let Some(path) = git(&["rev-parse", "--git-path", name]) { + println!("cargo:rerun-if-changed={path}"); + } + } + if let Some(head) = git(&["symbolic-ref", "-q", "HEAD"]) { + if let Some(path) = git(&["rev-parse", "--git-path", &head]) { + println!("cargo:rerun-if-changed={path}"); + } + } + git(&["rev-parse", "--short=12", "HEAD"]).unwrap_or_else(|| "unknown".into()) +} + fn main() { pyo3_build_config::use_pyo3_cfgs(); + println!("cargo:rerun-if-changed=build.rs"); + println!("cargo:rustc-env=DISKGRAPH_BUILD_COMMIT={}", build_commit()); if std::env::var_os("CARGO_FEATURE_EXTENSION_MODULE").is_some() { native_cpu_guard::reject_native_cpu_build("diskgraph Python extension"); diff --git a/packages/diskgraph/diskgraph-pybind/python/diskgraph/__init__.py b/packages/diskgraph/diskgraph-pybind/python/diskgraph/__init__.py index d157a5b0c..5b9f5d0b6 100644 --- a/packages/diskgraph/diskgraph-pybind/python/diskgraph/__init__.py +++ b/packages/diskgraph/diskgraph-pybind/python/diskgraph/__init__.py @@ -3,6 +3,8 @@ # SPDX-License-Identifier: Apache-2.0 from ._c import ( + __version__, + BUILD_COMMIT, export_tt_artifacts, convert_kumo_config, ingest, @@ -38,6 +40,7 @@ write_success_marker, touch_success_marker, verify_snapshot_integrity, + snapshot_range_checksums, has_resumable_ingest, validated_index_resume_progress, compute_col_stats_for_parquet, @@ -139,6 +142,7 @@ def __getattr__(name): "write_success_marker", "touch_success_marker", "verify_snapshot_integrity", + "snapshot_range_checksums", "has_resumable_ingest", "validated_index_resume_progress", "compute_col_stats_for_parquet", diff --git a/packages/diskgraph/diskgraph-pybind/python/tests/test_exception_exports.py b/packages/diskgraph/diskgraph-pybind/python/tests/test_exception_exports.py index 7744976d2..d4b29806e 100644 --- a/packages/diskgraph/diskgraph-pybind/python/tests/test_exception_exports.py +++ b/packages/diskgraph/diskgraph-pybind/python/tests/test_exception_exports.py @@ -11,3 +11,27 @@ def test_public_corruption_error_catches_snapshot_verifier(tmp_path): with pytest.raises(diskgraph.DiskgraphCorruptionError): diskgraph.verify_snapshot_integrity(str(tmp_path)) + + +def test_build_identity_is_exported(): + assert diskgraph.__version__ + assert diskgraph.BUILD_COMMIT + + +def test_snapshot_range_checksums_match_zlib(tmp_path): + import json + import zlib + + payload = bytes(range(256)) * 5 + (tmp_path / "a.bin").write_bytes(payload) + (tmp_path / "_MANIFEST.json").write_text(json.dumps({ + "schema_version": 1, + "files": [{"rel_path": "a.bin", "size_bytes": len(payload)}], + })) + crcs = diskgraph.snapshot_range_checksums(str(tmp_path), 512) + assert crcs["a.bin"] == [zlib.crc32(payload[i:i + 512]) + for i in range(0, len(payload), 512)] + assert crcs["_MANIFEST.json"] == [ + zlib.crc32((tmp_path / "_MANIFEST.json").read_bytes())] + with pytest.raises(diskgraph.DiskgraphCorruptionError): + diskgraph.snapshot_range_checksums(str(tmp_path / "missing"), 512) diff --git a/packages/diskgraph/diskgraph-pybind/src/lib.rs b/packages/diskgraph/diskgraph-pybind/src/lib.rs index c907055d1..a47daef4d 100644 --- a/packages/diskgraph/diskgraph-pybind/src/lib.rs +++ b/packages/diskgraph/diskgraph-pybind/src/lib.rs @@ -1553,6 +1553,21 @@ fn verify_snapshot_integrity(py: Python<'_>, output_dir: &str) -> PyResult<()> { .map_err(|e| DiskgraphCorruptionError::new_err(e.to_string())) } +/// CRC-32 of every `range_bytes` span of each file in the snapshot's +/// `_MANIFEST.json` and of the manifest itself, keyed by relative path. +#[pyfunction] +fn snapshot_range_checksums( + py: Python<'_>, + output_dir: &str, + range_bytes: u64, +) -> PyResult>> { + let output_dir = std::path::PathBuf::from(output_dir); + py.detach(|| { + diskgraph_ingest::snapshot_mirror::snapshot_range_checksums(&output_dir, range_bytes) + }) + .map_err(|e| DiskgraphCorruptionError::new_err(e.to_string())) +} + #[pyfunction] fn has_resumable_ingest(output_dir: &str) -> PyResult { diskgraph_ingest::initial_resume_phase(std::path::Path::new(output_dir)) @@ -1593,6 +1608,8 @@ fn diskgraph_pybind(m: &Bound<'_, PyModule>) -> PyResult<()> { "DiskgraphCorruptionError", m.py().get_type::(), )?; + m.add("__version__", env!("CARGO_PKG_VERSION"))?; + m.add("BUILD_COMMIT", env!("DISKGRAPH_BUILD_COMMIT"))?; m.add_function(wrap_pyfunction!(export_tt_artifacts, m)?)?; m.add_function(wrap_pyfunction!(convert_kumo_config, m)?)?; m.add_function(wrap_pyfunction!(ingest, m)?)?; @@ -1606,6 +1623,7 @@ fn diskgraph_pybind(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(success_marker::write_success_marker, m)?)?; m.add_function(wrap_pyfunction!(success_marker::touch_success_marker, m)?)?; m.add_function(wrap_pyfunction!(verify_snapshot_integrity, m)?)?; + m.add_function(wrap_pyfunction!(snapshot_range_checksums, m)?)?; m.add_function(wrap_pyfunction!(has_resumable_ingest, m)?)?; m.add_function(wrap_pyfunction!(validated_index_resume_progress, m)?)?; m.add_function(wrap_pyfunction!(compute_col_stats_for_parquet, m)?)?; diff --git a/packages/diskgraph/third_party/artifact-manifest.json b/packages/diskgraph/third_party/artifact-manifest.json index 1f62d43ef..d22f0bc33 100644 --- a/packages/diskgraph/third_party/artifact-manifest.json +++ b/packages/diskgraph/third_party/artifact-manifest.json @@ -3241,6 +3241,13 @@ "source_kind": "registry-or-git", "version": "0.3.0" }, + { + "license": "MIT OR Apache-2.0", + "name": "crc32fast", + "source": "registry+https://github.com/rust-lang/crates.io-index", + "source_kind": "registry-or-git", + "version": "1.5.0" + }, { "license": "MIT OR Apache-2.0", "name": "crossbeam-deque", @@ -5202,6 +5209,6 @@ ] } ], - "cargo_lock_sha256": "297851eb08303a13a97f091c68cf3ce5ccac633ad3d7fe09724c1506e486cdaf", + "cargo_lock_sha256": "a1c319a5eb83c1a60f056c213932c12e385f0cff468da4db0021947f11b7aa07", "schema": 1 } diff --git a/packages/kumo-bigquery/src/kumo_bigquery/identity.py b/packages/kumo-bigquery/src/kumo_bigquery/identity.py index dede046ed..d22d8aada 100644 --- a/packages/kumo-bigquery/src/kumo_bigquery/identity.py +++ b/packages/kumo-bigquery/src/kumo_bigquery/identity.py @@ -6,7 +6,11 @@ from typing import Any -from kumo_core.sources.identity import epoch_second, hour_version +from kumo_core.sources.identity import ( + SourceVersion, + epoch_second, + hour_version, +) LOGICAL_VIEW = "VIEW" UNDATED_VERSION = "-1" @@ -23,8 +27,8 @@ def stamp(table: Any) -> str: if is_logical_view(table): return hour_version() if table.modified is None: - return UNDATED_VERSION - return epoch_second(table.modified) + return SourceVersion(UNDATED_VERSION) + return SourceVersion(epoch_second(table.modified), "timestamp") def data_version(client: Any, fqn: str) -> str: diff --git a/packages/kumo-bigquery/src/kumo_bigquery/session.py b/packages/kumo-bigquery/src/kumo_bigquery/session.py index 36dada9c8..9e5283ab6 100644 --- a/packages/kumo-bigquery/src/kumo_bigquery/session.py +++ b/packages/kumo-bigquery/src/kumo_bigquery/session.py @@ -278,6 +278,11 @@ def _credentials(self, connector: SourceConnector) -> Credentials: _require_credentials(credentials) return credentials + def table_location(self, connector: SourceConnector, + table_name: str) -> str: + """Fully qualified name of ``table_name`` in the connector's target.""" + return _name(connector, table_name) + def _client(self, connector: SourceConnector) -> Any: return _clients.query_client( self._principal(connector), project=_billing_project(connector), diff --git a/packages/kumo-bigquery/tests/test_identity.py b/packages/kumo-bigquery/tests/test_identity.py index 28d98fc1d..82f74bbc2 100644 --- a/packages/kumo-bigquery/tests/test_identity.py +++ b/packages/kumo-bigquery/tests/test_identity.py @@ -20,6 +20,7 @@ def stamp(**overrides): def test_a_physical_table_is_dated_by_the_time_it_was_modified(): """A version carrying `.0` is a different path segment and a different key.""" assert stamp() == MODIFIED_VERSION + assert stamp().kind == "timestamp" assert stamp(modified=MODIFIED.replace( microsecond=900000)) == (MODIFIED_VERSION) assert stamp( @@ -28,6 +29,7 @@ def test_a_physical_table_is_dated_by_the_time_it_was_modified(): def test_a_logical_view_falls_back_to_the_hour_window(): assert stamp(table_type="VIEW") == hour_version() + assert stamp(table_type="VIEW").kind == "hour_bucket" @pytest.mark.parametrize("table_type", @@ -39,6 +41,7 @@ def test_only_the_exact_view_type_falls_back(table_type): def test_a_table_without_a_modification_time_is_reported_undated(): assert stamp(modified=None) == "-1" + assert stamp(modified=None).kind == "unknown" def test_the_stamp_reads_the_catalog_and_scans_nothing(): diff --git a/packages/kumo-core/src/kumo_core/artifact_export.py b/packages/kumo-core/src/kumo_core/artifact_export.py index 0d73620ca..311b972bc 100644 --- a/packages/kumo-core/src/kumo_core/artifact_export.py +++ b/packages/kumo-core/src/kumo_core/artifact_export.py @@ -199,18 +199,22 @@ def publish_export_destination(output_path, *, fs): def export_selected_files(files, output_path, *, fs, identity, file_type='parquet', convert_csv=None, - copy_parquet=None, temporary_directory=None): + copy_parquet=None, temporary_directory=None, + sidecars={}): """Publish the hosted selected-file layout and success marker last. Filesystem/authentication and retry wrappers stay with the caller. The directory ownership, output names and CSV retry decisions are shared. + ``sidecars`` (name to bytes) land in a directory destination before + ``_SUCCESS``. """ import os.path as osp import tempfile + output_is_dir = output_path.endswith('/') or len(files) > 1 output_path, _ = prepare_export_destination( - output_path, output_is_dir=output_path.endswith('/') or len(files) > 1, - file_type=file_type, fs=fs, identity=identity) + output_path, output_is_dir=output_is_dir, file_type=file_type, fs=fs, + identity=identity) if file_type == 'parquet': if copy_parquet is None: @@ -230,6 +234,14 @@ def export_selected_files(files, output_path, *, fs, identity, else: raise ValueError(f'Unsupported artifact file type: {file_type}') + if output_is_dir: + for name, payload in sidecars.items(): + try: + fs.write_bytes(osp.join(output_path, name), payload) + except Exception as error: + import warnings + warnings.warn(f'Exporting without {name}: {error}', + UserWarning, stacklevel=2) return publish_export_destination(output_path, fs=fs) diff --git a/packages/kumo-core/src/kumo_core/artifacts.py b/packages/kumo-core/src/kumo_core/artifacts.py index 7a4226974..10b33ee11 100644 --- a/packages/kumo-core/src/kumo_core/artifacts.py +++ b/packages/kumo-core/src/kumo_core/artifacts.py @@ -13,9 +13,16 @@ import tempfile from dataclasses import dataclass, field from pathlib import Path -from typing import Any, ClassVar, Mapping, Protocol, runtime_checkable - -from .io.filesystem import local_path +from typing import ( + Any, + ClassVar, + Iterable, + Mapping, + Protocol, + runtime_checkable, +) + +from .io.filesystem import atomic_write_bytes, fsync_directory, local_path from .io.json import load_json, save_json from .manifest import RunManifest, RunPaths @@ -66,36 +73,7 @@ def write_manifest(self, paths: RunPaths, manifest: RunManifest) -> None: ensure_ascii=False, allow_nan=False, ) + "\n" - temporary_path: Path | None = None - try: - with tempfile.NamedTemporaryFile( - mode="w", - encoding="utf-8", - dir=paths.manifest_path.parent, - prefix=f".{paths.manifest_path.name}.", - suffix=".tmp", - delete=False, - ) as temporary: - temporary_path = Path(temporary.name) - temporary.write(payload) - temporary.flush() - os.fsync(temporary.fileno()) - os.replace(temporary_path, paths.manifest_path) - temporary_path = None - self._fsync_directory(paths.manifest_path.parent) - finally: - if temporary_path is not None: - temporary_path.unlink(missing_ok=True) - - @staticmethod - def _fsync_directory(directory: Path) -> None: - if not hasattr(os, "O_DIRECTORY"): - return - descriptor = os.open(directory, os.O_RDONLY | os.O_DIRECTORY) - try: - os.fsync(descriptor) - finally: - os.close(descriptor) + atomic_write_bytes(paths.manifest_path, payload.encode("utf-8")) @dataclass(frozen=True) @@ -262,7 +240,7 @@ def publish( encoding="utf-8", ) os.replace(staging, destination) - FilesystemArtifactStore._fsync_directory(destination.parent) + fsync_directory(destination.parent) return publication finally: if staging.exists(): @@ -280,11 +258,10 @@ def _sha256(path: Path) -> str: def verify_model_bundle(bundle: Path) -> str: """Verify fetched artifacts before parsing or staging a model bundle. - Preserves the local SDK import verifier's error and content-identity - contract. Paths must stay inside the bundle after symlink resolution; - declared sizes and SHA-256 digests must match. Older manifests may omit - sizes or leave training-spec.json unlisted; its digest still participates - in the returned content identity. + Returns the bundle id, a SHA-256 over every verified file except the + manifest (``_provenance.json``). Paths must stay inside the bundle after symlink + resolution; declared sizes and SHA-256 digests must match. Older + manifests may omit sizes or leave training-spec.json unlisted. """ from .sources import layout from .sources.errors import PermanentSourceError @@ -302,7 +279,7 @@ def verify_model_bundle(bundle: Path) -> str: artifacts = manifest.get('artifacts') problems: list[str] = [] - verified: list[tuple[str, str, str]] = [] + verified: list[tuple[str, str]] = [] seen = set() claimed = set() if not artifacts: @@ -354,7 +331,7 @@ def verify_model_bundle(bundle: Path) -> str: if VitalArtifactPublicationPolicy._sha256(path) != expected: problems.append(f'{relative} does not match its digest') continue - verified.append((role, safe, expected)) + verified.append((safe, expected)) if problems: raise ValueError( f"The model bundle at {bundle} did not survive the transfer: " @@ -364,11 +341,53 @@ def verify_model_bundle(bundle: Path) -> str: if not spec_path.resolve().is_relative_to(bundle): raise ValueError('Training specification is not below the bundle') if 'training-spec.json' not in seen and spec_path.is_file(): - verified.append(('training_spec', 'training-spec.json', + verified.append(('training-spec.json', VitalArtifactPublicationPolicy._sha256(spec_path))) - payload = json.dumps(sorted(verified), sort_keys=True, - separators=(',', ':'), default=str) - return hashlib.sha256(payload.encode('utf-8')).hexdigest()[:16] + return _bundle_id(verified) + + +def bundle_id_from_manifest(bundle: Path) -> str: + """The bundle id from the digests a freshly published bundle lists.""" + manifest_path = Path(bundle) / VitalArtifactPublicationPolicy.MANIFEST_NAME + manifest = json.loads(manifest_path.read_text(encoding='utf-8')) + return _bundle_id((str(entry['relative_path']), str(entry['sha256'])) + for entry in manifest['artifacts']) + + +def _bundle_id(pairs: Iterable[tuple[str, str]]) -> str: + payload = json.dumps( + sorted(pair for pair in pairs if pair[0] != PROVENANCE_FILE), + separators=(',', ':')) + return hashlib.sha256(payload.encode('utf-8')).hexdigest() + + +def add_bundle_file(bundle_dir: Path, role: str, relative_path: str) -> None: + """List one file already written in a bundle in its publication manifest.""" + bundle_dir = Path(bundle_dir) + path = bundle_dir / relative_path + manifest_path = bundle_dir / VitalArtifactPublicationPolicy.MANIFEST_NAME + manifest = json.loads(manifest_path.read_text(encoding='utf-8')) + manifest['artifacts'] = [ + entry for entry in manifest.get('artifacts') or [] + if entry.get('relative_path') != relative_path + and entry.get('role') != role + ] + [ + PublishedArtifact( + role=role, + relative_path=relative_path, + size_bytes=path.stat().st_size, + sha256=VitalArtifactPublicationPolicy._sha256(path), + ).to_dict() + ] + _write_bundle_manifest(bundle_dir, manifest) + + +def _write_bundle_manifest(bundle_dir: Path, + manifest: Mapping[str, Any]) -> None: + payload = json.dumps(manifest, indent=2, sort_keys=True) + '\n' + atomic_write_bytes( + bundle_dir / VitalArtifactPublicationPolicy.MANIFEST_NAME, + payload.encode('utf-8')) @runtime_checkable diff --git a/packages/kumo-core/src/kumo_core/io/filesystem.py b/packages/kumo-core/src/kumo_core/io/filesystem.py index 7ac2de5be..c8c227148 100644 --- a/packages/kumo-core/src/kumo_core/io/filesystem.py +++ b/packages/kumo-core/src/kumo_core/io/filesystem.py @@ -7,6 +7,7 @@ import os import shutil +import tempfile from collections.abc import Callable from pathlib import Path, PurePosixPath from typing import Any @@ -85,6 +86,37 @@ def local_path(path: str) -> Path: return Path(fs_path) +def fsync_directory(directory: str | Path) -> None: + if not hasattr(os, "O_DIRECTORY"): + return + descriptor = os.open(directory, os.O_RDONLY | os.O_DIRECTORY) + try: + os.fsync(descriptor) + finally: + os.close(descriptor) + + +def atomic_write_bytes(path: str | Path, payload: bytes) -> None: + """Replace a local file so readers see the old or the new bytes.""" + path = Path(path) + temporary: Path | None = None + try: + with tempfile.NamedTemporaryFile(dir=path.parent, + prefix=f".{path.name}.", + suffix=".tmp", + delete=False) as out: + temporary = Path(out.name) + out.write(payload) + out.flush() + os.fsync(out.fileno()) + os.replace(temporary, path) + temporary = None + fsync_directory(path.parent) + finally: + if temporary is not None: + temporary.unlink(missing_ok=True) + + def listdir(path: str) -> list[str]: fs, fs_path = _resolve(path) return [fs.unstrip_protocol(item) for item in fs.ls(fs_path, detail=False)] diff --git a/packages/kumo-core/src/kumo_core/object_store_sink.py b/packages/kumo-core/src/kumo_core/object_store_sink.py index 3951ea5b3..b6750b69a 100644 --- a/packages/kumo-core/src/kumo_core/object_store_sink.py +++ b/packages/kumo-core/src/kumo_core/object_store_sink.py @@ -2,6 +2,7 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import logging import os import posixpath import shutil @@ -20,6 +21,7 @@ from pyarrow.fs import FileSystem _CHUNK = 8 * 1024 * 1024 +_LOG = logging.getLogger(__name__) _OBJECT_STORES = frozenset({'s3', 'gcs', 'abfs'}) @@ -59,6 +61,31 @@ def upload_bundle(filesystem: FileSystem, destination: str, shutil.copyfileobj(reader, writer, _CHUNK) +def holds_bundle(filesystem: FileSystem, destination: str, + bundle_dir: str) -> bool: + """Whether ``destination`` holds a complete copy of ``bundle_dir``: the + same files and sizes, and the same provenance file, written last.""" + from kumo_core.artifacts import PROVENANCE_FILE + from pyarrow.fs import FileSelector, FileType + + root = Path(bundle_dir) + local = { + path.relative_to(root).as_posix(): path.stat().st_size + for path in root.rglob('*') if path.is_file() + } + remote = { + info.path[len(destination) + 1:]: info.size + for info in filesystem.get_file_info( + FileSelector(destination, recursive=True, allow_not_found=True)) + if info.type is FileType.File + } + if remote != local or PROVENANCE_FILE not in local: + return False + with filesystem.open_input_stream( + f'{destination}/{PROVENANCE_FILE}') as reader: + return reader.readall() == (root / PROVENANCE_FILE).read_bytes() + + def download_bundle(filesystem: FileSystem, source: str, local_dir: str) -> bool: """Replace ``local_dir`` with the files under ``source``; False if there @@ -130,7 +157,8 @@ def write_table( def write_files(self, connector, table_name, files, write_mode=WriteMode.OVERWRITE, project_id=None, *, - export_id=None, extra_columns=None): + export_id=None, extra_columns=None, + sidecars: Mapping[str, bytes] = {}): """Export selected Parquet without collecting their rows.""" del project_id if write_mode is not WriteMode.OVERWRITE: @@ -140,11 +168,16 @@ def write_files(self, connector, table_name, files, from kumo_core.sources import layout filesystem, path = self._destination(connector, table_name) kind = layout.file_type(path.rstrip('/')) or layout.PARQUET + if export_id and holds_export(filesystem, path, export_id): + _LOG.info('%s already holds a complete copy of export %s; reused ' + 'it', path, export_id) + return f'{self.scheme}://{path.rstrip("/")}/' storage = _ArrowExportFiles(filesystem, extra_columns) _, output = export_selected_files(files, path, fs=storage, identity=export_id or str(files), file_type=kind, - convert_csv=storage.convert_csv) + convert_csv=storage.convert_csv, + sidecars=sidecars) return f'{self.scheme}://{output}' def publish_model( @@ -173,7 +206,11 @@ def publish_model( filesystem, destination = self._destination(connector, path) destination = destination.rstrip('/') - upload_bundle(filesystem, destination, bundle_dir) + if holds_bundle(filesystem, destination, bundle_dir): + _LOG.info('%s already holds this model bundle; reused it', + destination) + else: + upload_bundle(filesystem, destination, bundle_dir) return PublishedModel( name=posixpath.basename(destination), version=destination, @@ -309,6 +346,23 @@ def _prepared(filesystem: FileSystem, directory: str) -> None: filesystem.create_dir(directory, recursive=True) +def holds_export(filesystem: FileSystem, path: str, export_id: str) -> bool: + """Whether the directory ``path`` holds a complete export ``export_id``.""" + from kumo_core.artifact_export import EXPORT_MARKER_NAME + from pyarrow.fs import FileType + + directory = path.rstrip('/') + try: + if filesystem.get_file_info( + f'{directory}/_SUCCESS').type is not FileType.File: + return False + with filesystem.open_input_stream( + f'{directory}/{EXPORT_MARKER_NAME}') as reader: + return reader.readall().decode('utf-8') == _marker(export_id) + except (FileNotFoundError, OSError): + return False + + def _marker(export_id: Optional[str]) -> Optional[str]: from kumo_core.artifact_export import export_marker @@ -342,8 +396,11 @@ def read_text(self, path): return reader.readall().decode('utf-8') def write_text(self, path, value): + self.write_bytes(path, value.encode('utf-8')) + + def write_bytes(self, path, value): with self.filesystem.open_output_stream(path) as writer: - writer.write(value.encode('utf-8')) + writer.write(value) def makedirs(self, path, exist_ok=True): _prepared(self.filesystem, path) diff --git a/packages/kumo-core/src/kumo_core/sources/identity.py b/packages/kumo-core/src/kumo_core/sources/identity.py index ac69d6b92..0e3d8e664 100644 --- a/packages/kumo-core/src/kumo_core/sources/identity.py +++ b/packages/kumo-core/src/kumo_core/sources/identity.py @@ -20,6 +20,26 @@ CUSTOM_SQL_MAX_STALENESS_HOURS = 24 +class SourceVersion(str): + """A version string that also says what kind of evidence it is: + ``exact``, ``timestamp``, ``listing``, ``hour_bucket`` or ``unknown``.""" + + kind: str + dependencies: Mapping[str, str] | None + + def __new__(cls, value: str, kind: str = "unknown", + dependencies: Mapping[str, str] | None = None + ) -> "SourceVersion": + version = super().__new__(cls, value) + version.kind = kind + version.dependencies = dependencies + return version + + +def version_kind(version: object) -> str: + return getattr(version, "kind", "unknown") + + def epoch_second(moment: datetime) -> str: """Truncated by string rather than through ``int``: the two disagree just below the epoch, and this value is a raw path segment.""" @@ -29,8 +49,10 @@ def epoch_second(moment: datetime) -> str: def hour_version() -> str: """Fallback for a source with no readable version; it expires itself rather than needing invalidation.""" - return epoch_second( - datetime.now(timezone.utc).replace(minute=0, second=0, microsecond=0)) + return SourceVersion( + epoch_second( + datetime.now(timezone.utc).replace(minute=0, second=0, + microsecond=0)), "hour_bucket") def sql_data_version( @@ -83,10 +105,11 @@ def sql_data_version( versions[leaf] = _indeterminate() else: versions[leaf] = str(stamped) - return custom_sql_data_version( - sql, connector_type, versions, cols or {}, - max_staleness_bucket=int(time()) // - (CUSTOM_SQL_MAX_STALENESS_HOURS * 3600)) + return SourceVersion( + custom_sql_data_version( + sql, connector_type, versions, cols or {}, + max_staleness_bucket=int(time()) // + (CUSTOM_SQL_MAX_STALENESS_HOURS * 3600)), "hour_bucket", versions) #: The prefix every sentinel carries, so one can be recognized without being diff --git a/packages/kumo-core/src/kumo_core/sources/object_store.py b/packages/kumo-core/src/kumo_core/sources/object_store.py index 4b7ac0a37..bf1058d07 100644 --- a/packages/kumo-core/src/kumo_core/sources/object_store.py +++ b/packages/kumo-core/src/kumo_core/sources/object_store.py @@ -52,7 +52,7 @@ TransientSourceError, UnsupportedReadModeError, ) -from .identity import epoch_second +from .identity import SourceVersion, epoch_second from .object_store_access import ObjectFetch, ObjectIdentity from .parallel import run_in_parallel from .prepare import ReadOutput, run_prepare @@ -101,14 +101,16 @@ def _source_version(path: str, entries: Sequence[Any]) -> str: if len(entries) == 1 and entries[0].path.rstrip("/") == path.rstrip("/"): latest = entries[0].mtime - return epoch_second(latest) if latest is not None else "-1" + return (SourceVersion(epoch_second(latest), "timestamp") + if latest is not None else SourceVersion("-1")) # A maximum mtime misses changes to older files and added/deleted parts. metadata = [ (entry.path, entry.size, getattr(entry, "mtime_ns", None), entry.mtime.isoformat() if entry.mtime is not None else None) for entry in sorted(entries, key=lambda entry: entry.path) ] - return hashlib.sha256(json.dumps(metadata).encode()).hexdigest() + return SourceVersion( + hashlib.sha256(json.dumps(metadata).encode()).hexdigest(), "listing") class _CsvSource(NamedTuple): @@ -489,7 +491,7 @@ def _copy(self, path: str, found: Sequence[Any], dest: str, if self.fetch is not None: return self._download_objects(path, found, dest, worker_ceiling, parquet=True) - with _copies(found, worker_ceiling) as workers: + with copies(found, worker_ceiling) as workers: return fetch_all([ StagedFile( entry.path, entry.size @@ -519,7 +521,7 @@ def download(entry: Any) -> str: return target try: - with _copies(found, worker_ceiling) as workers, ThreadPoolExecutor( + with copies(found, worker_ceiling) as workers, ThreadPoolExecutor( max_workers=workers) as pool: futures = [pool.submit(download, entry) for entry in found] try: @@ -814,8 +816,8 @@ def _parquet_batches(fragment: Any, @contextmanager -def _copies(found: Sequence[Any], - worker_ceiling: int | None) -> Iterator[int]: +def copies(found: Sequence[Any], + worker_ceiling: int | None = None) -> Iterator[int]: """Arrow-backed downloaders copy on Arrow's IO pool.""" with memory_share( COPY_WORKER_BYTES, lambda ceiling: min( diff --git a/packages/kumo-core/src/kumo_core/sources/object_store_backend.py b/packages/kumo-core/src/kumo_core/sources/object_store_backend.py index 69f3025f7..f5f903d0a 100644 --- a/packages/kumo-core/src/kumo_core/sources/object_store_backend.py +++ b/packages/kumo-core/src/kumo_core/sources/object_store_backend.py @@ -130,7 +130,7 @@ def get_data_version( ) -> str: del project_id reader, name = self._reader(connector, table_name) - return str(reader.data_version(name)) + return reader.data_version(name) def prepare_ingest( self, diff --git a/packages/kumo-core/src/kumo_core/sources/object_store_sessions.py b/packages/kumo-core/src/kumo_core/sources/object_store_sessions.py index 1e9dd307a..951b780bb 100644 --- a/packages/kumo-core/src/kumo_core/sources/object_store_sessions.py +++ b/packages/kumo-core/src/kumo_core/sources/object_store_sessions.py @@ -6,7 +6,7 @@ from abc import ABC, abstractmethod from dataclasses import dataclass from typing import Any, Callable, Optional, Tuple -from urllib.parse import unquote, urlsplit +from urllib.parse import quote, unquote, urlsplit from kumoapi.plugin import SourceConnector from pyarrow.fs import FileSystem @@ -16,6 +16,8 @@ from .object_store_access import ObjectStoreAccess from .principals import IdentityRef, Principal, resolve +_KEY_SAFE = "/-._~!$&'()*+,;=:@" + @dataclass(frozen=True) class ObjectStoreRoot: @@ -104,6 +106,15 @@ def _over(self, connector: SourceConnector, bucket: str, root: str, fetch=access.fetch, object_identity=access.object_identity) + def table_location(self, connector: SourceConnector, + table_name: str) -> str: + """Full ``scheme://bucket/key`` of ``table_name``, without credentials.""" + root = self._root_for(connector) + if root is None: + return self.paths.source_uri(table_name).rstrip('/') + key = quote(posixpath.join(root.prefix, table_name), safe=_KEY_SAFE) + return f'{self.scheme}://{root.bucket}/{key}'.rstrip('/') + def _placement(self, connector: SourceConnector, table_name: str) -> Tuple[str, str, str, Optional[str]]: root = self._root_for(connector) diff --git a/packages/kumo-core/src/kumo_core/train/_artifact_storage.py b/packages/kumo-core/src/kumo_core/train/_artifact_storage.py index 6e6cc5ebd..06153d6f7 100644 --- a/packages/kumo-core/src/kumo_core/train/_artifact_storage.py +++ b/packages/kumo-core/src/kumo_core/train/_artifact_storage.py @@ -16,6 +16,8 @@ from pathlib import Path from typing import Any, Callable +from kumo_core.io.filesystem import atomic_write_bytes, fsync_directory + class _LocalFiles: @staticmethod @@ -54,7 +56,7 @@ def write_text_if_absent(path, text): os.link(temporary, path) except FileExistsError: return False - _fsync_directory(Path(path).parent) + fsync_directory(Path(path).parent) return True @@ -74,20 +76,9 @@ def _temporary_text(path, text): os.unlink(temporary) -def _fsync_directory(directory): - if not hasattr(os, "O_DIRECTORY"): - return - descriptor = os.open(directory, os.O_RDONLY | os.O_DIRECTORY) - try: - os.fsync(descriptor) - finally: - os.close(descriptor) - - def _save_json_atomic(value, path): - with _temporary_text(path, json.dumps(value, indent=2)) as temporary: - os.replace(temporary, path) - _fsync_directory(Path(path).parent) + Path(path).parent.mkdir(parents=True, exist_ok=True) + atomic_write_bytes(path, json.dumps(value, indent=2).encode('utf-8')) @dataclass(frozen=True) diff --git a/packages/kumo-core/tests/sources/parquet/test_object_store.py b/packages/kumo-core/tests/sources/parquet/test_object_store.py index 4c2001d9a..5989a09f5 100644 --- a/packages/kumo-core/tests/sources/parquet/test_object_store.py +++ b/packages/kumo-core/tests/sources/parquet/test_object_store.py @@ -144,9 +144,9 @@ def get_file_info(self, target): reader = ObjectStoreSourceReader(MetadataOnlyFilesystem(), root=str(root)) version = reader.data_version(table) if table in ("events", "hive"): - assert len(version) == 64 + assert (len(version), version.kind) == (64, "listing") else: - assert int(version) > 0 + assert int(version) > 0 and version.kind == "timestamp" def test_version_of_exact_root_does_not_probe_other_suffixes(root): diff --git a/packages/kumo-core/tests/sources/unit/test_sql_identity.py b/packages/kumo-core/tests/sources/unit/test_sql_identity.py index 7f3ad3993..2d18419d9 100644 --- a/packages/kumo-core/tests/sources/unit/test_sql_identity.py +++ b/packages/kumo-core/tests/sources/unit/test_sql_identity.py @@ -36,6 +36,12 @@ def test_a_referenced_table_that_moved_gives_a_new_version(): assert stamped(version="v1") != stamped(version="v2") +def test_the_version_keeps_the_dependency_versions_it_folds(): + version = stamped(version="v1") + assert version.kind == "hour_bucket" + assert set(version.dependencies.values()) == {"v1"} + + def test_a_formatting_only_edit_does_not_move_the_version(): assert stamped( sql="select ID,\n AMOUNT from ORDERS where AMOUNT > 0") == (stamped()) diff --git a/packages/kumo-core/tests/test_artifact_export.py b/packages/kumo-core/tests/test_artifact_export.py index 39aba03b2..b3701f00f 100644 --- a/packages/kumo-core/tests/test_artifact_export.py +++ b/packages/kumo-core/tests/test_artifact_export.py @@ -222,3 +222,35 @@ def interrupted(source, destination, **kwargs): pd.testing.assert_frame_equal(pd.read_csv(target / 'part-00000.csv'), expected) assert (target / '_SUCCESS').exists() + + + +def test_sidecars_are_best_effort_and_precede_success(tmp_path): + import pyarrow.parquet as pq + from kumo_core.artifact_export import export_selected_files + from kumo_core.object_store_sink import _ArrowExportFiles + from pyarrow.fs import LocalFileSystem + + source = tmp_path / 'part-0.parquet' + pq.write_table(pa.table({'id': [1]}), source) + target = tmp_path / 'output' + fs = _ArrowExportFiles(LocalFileSystem(), None) + write_bytes, touch = fs.write_bytes, fs.touch + + def denied(path, value): + if path.endswith('bad.json'): + raise PermissionError('denied') + write_bytes(path, value) + + def checked_touch(path): + assert (target / 'meta.json').read_bytes() == b'{}' + touch(path) + + fs.write_bytes, fs.touch = denied, checked_touch + with pytest.warns(UserWarning, match='denied'): + export_selected_files([str(source)], f'{target}/', fs=fs, + identity='job-1', sidecars={ + 'meta.json': b'{}', + 'bad.json': b'{}' + }) + assert (target / '_SUCCESS').exists() diff --git a/packages/kumo-core/tests/test_artifacts.py b/packages/kumo-core/tests/test_artifacts.py index 071f4339b..17e5255b0 100644 --- a/packages/kumo-core/tests/test_artifacts.py +++ b/packages/kumo-core/tests/test_artifacts.py @@ -15,7 +15,10 @@ FilesystemVitalArtifactPublisher, VitalArtifactPublicationPolicy, VitalArtifactSources, + add_bundle_file, + bundle_id_from_manifest, resolve_vital_artifact_sources, + verify_model_bundle, ) from kumo_core.config import NativeRuntimeConfig from kumo_core.manifest import RunManifest, RunPaths @@ -237,3 +240,26 @@ def best_ckpt(self): assert sources.selected_checkpoint == checkpoints[selected_mode] assert sources.graph_manifest == bundle / "graph.json" assert json.loads(sources.metrics.read_text()) == {"test_accuracy": 0.8} + + +def test_bundle_id_ignores_provenance_only(tmp_path: Path) -> None: + sources, _ = _vital_sources(tmp_path) + bundle = tmp_path / "bundle" + VitalArtifactPublicationPolicy().publish(sources, bundle) + bundle_id = verify_model_bundle(bundle) + + for name, payload in (("_provenance.json", '{"a": 1}'), + ("_provenance.json", '{"a": 2}')): + (bundle / name).write_text(payload) + add_bundle_file(bundle, "provenance", name) + assert verify_model_bundle(bundle) == bundle_id + assert bundle_id_from_manifest(bundle) == bundle_id + roles = [ + entry["role"] for entry in json.loads( + (bundle / "publication-manifest.json").read_text())["artifacts"] + ] + assert roles.count("provenance") == 1 + + (bundle / "metrics.json").write_text('{"auc": 0}') + add_bundle_file(bundle, "metrics", "metrics.json") + assert verify_model_bundle(bundle) != bundle_id diff --git a/packages/kumo-databricks/src/kumo_databricks/identity.py b/packages/kumo-databricks/src/kumo_databricks/identity.py index ae6023a93..c4ea11088 100644 --- a/packages/kumo-databricks/src/kumo_databricks/identity.py +++ b/packages/kumo-databricks/src/kumo_databricks/identity.py @@ -8,7 +8,11 @@ from databricks.sql.client import Cursor from kumo_core.sources.errors import SourceNotFoundError -from kumo_core.sources.identity import epoch_second, hour_version +from kumo_core.sources.identity import ( + SourceVersion, + epoch_second, + hour_version, +) from kumo_databricks.errors import is_not_found DESCRIBE = "DESCRIBE EXTENDED {source}" @@ -65,5 +69,5 @@ def table_version(cursor: Cursor, source: str) -> str: rows = cursor.fetchall() if rows: commit = epoch_second(pd.Timestamp(rows[0][1])) - return f"{rows[0][0]}_{commit}" + return SourceVersion(f"{rows[0][0]}_{commit}", "exact") return hour_version() diff --git a/packages/kumo-databricks/src/kumo_databricks/session.py b/packages/kumo-databricks/src/kumo_databricks/session.py index af5a9209e..b274cd234 100644 --- a/packages/kumo-databricks/src/kumo_databricks/session.py +++ b/packages/kumo-databricks/src/kumo_databricks/session.py @@ -194,6 +194,11 @@ def _principal(self, connector: SourceConnector) -> Principal: f'accepted, as {self._default_hint}.'), ) + def table_location(self, connector: SourceConnector, + table_name: str) -> str: + """Fully qualified name of ``table_name`` in the connector's target.""" + return self._qualified(connector, self._bare(connector, table_name)) + @staticmethod def _qualified(connector: SourceConnector, table_name: Optional[str] = None) -> str: diff --git a/packages/kumo-databricks/tests/test_identity.py b/packages/kumo-databricks/tests/test_identity.py index 28bd35f48..03aeee1ad 100644 --- a/packages/kumo-databricks/tests/test_identity.py +++ b/packages/kumo-databricks/tests/test_identity.py @@ -49,6 +49,7 @@ def stamp(recorder=None, source=SOURCE, **overrides): def test_a_delta_table_is_stamped_with_its_version_and_commit_time(): """The first two history columns are the stamp; a fraction is dropped.""" assert stamp() == COMMIT_VERSION + assert stamp().kind == "exact" assert stamp(history=[(947, f"{COMMITTED}.9", "WRITE", "extra")]) == (COMMIT_VERSION) @@ -60,6 +61,7 @@ def test_the_stamp_changes_when_the_same_version_is_recommitted(): def test_a_table_with_no_history_falls_back_to_the_hour_window(): assert stamp(history=[]) == hour_version() + assert stamp(history=[]).kind == "hour_bucket" @pytest.mark.parametrize("refusal", [ diff --git a/packages/kumo-engine/README.md b/packages/kumo-engine/README.md index 6466a8d08..8596f7262 100644 --- a/packages/kumo-engine/README.md +++ b/packages/kumo-engine/README.md @@ -103,6 +103,7 @@ To use your own data, point `FileConnector` at a directory of CSV or Parquet fil ## Documentation - [Kumo Engine documentation](https://docs.nvidia.com/sdgm/kumo-engine/get-started-with-kumo-engine/overview): prerequisites, GPU sizing, installation, and end-to-end workflows. Its User Code Guide covers connecting data, tables and graphs, training, and prediction. +- [MLOps](docs/mlops.md): lineage manifests, the artifact root, tracking and registries. - [Predictive query examples](https://docs.nvidia.com/sdgm/kumo-engine/get-started-with-kumo-engine/workflows-and-solutions/predictive-query-examples): common prediction tasks written as predictive queries. ## Support diff --git a/packages/kumo-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md new file mode 100644 index 000000000..77bc7f609 --- /dev/null +++ b/packages/kumo-engine/docs/mlops.md @@ -0,0 +1,88 @@ +# MLOps with Kumo Engine + +## Turning MLOps on + +`kumo.MLOps` says where artifacts go. Use it as a block, which notebooks and scheduled scripts alike can do, or pass it as `mlops=` on a single call: + +```python +import kumo.engine as kumo + +cfg = kumo.MLOps(artifact_root="s3://bucket/kumo") # or (connector, path), e.g. a volume or stage +with kumo.mlops(cfg): + snapshot = graph.snapshot() +``` + +- **`artifact_root`** is required. It is a URI (`s3://`, `gs://`, `file://`, or a local path, including `/Volumes/...` on a Databricks cluster) or a `(connector, path)` pair for a store a URI cannot reach. +- **Volumes and stages:** `(DatabricksConnector(...), "/Volumes////kumo")` and `(SnowflakeConnector(...), "@DB.SCHEMA.STAGE/kumo")` work as roots once the installed kumo-databricks or kumo-snowflake sink holds files at those paths; with an older plugin, building `kumo.MLOps` raises. +- **`export`** names the artifact kinds stored under the root when they are produced: `snapshot`, `training_table` and `model` (default: all). +- **Precedence:** a call's `mlops=` replaces the enclosing block, and an inner block replaces an outer one. A block covers the calling thread or asyncio task. +- **Off by default.** Without a block or `mlops=`, every call behaves as it always has. Plain exports (`export(config)`, `export_config=`, `output_config=`) never write to the artifact root. A path export that finds a complete copy of the same artifact at its destination reuses it instead of uploading again (logged at info): a snapshot with the same bytes, a training table or prediction directory with `_SUCCESS` and the same `_KUMO_EXPORT` identity, or a model bundle with the same files and `_provenance.json`. + +## Lineage + +Every artifact carries a `_provenance.json` manifest, metadata only. The leading underscore keeps Spark and Hive Parquet readers from treating it as data. + +A graph snapshot's manifest, next to its index, records each source table's connector, location and version, plus the engine build that wrote it. `graph.snapshot(record_sources=False)` keeps only connector and table names, plus a SHA-256 of each full location so a source moved elsewhere is still noticed, and keeps Custom SQL only as a hash. + +### Training tables + +A generated or caller-supplied training table carries `_provenance.json` in its staging directory: +- the query, the effective generation plan and the task type; +- the split seed; +- the snapshot it was generated from (that snapshot's manifest, embedded); +- for `TrainingTable.from_source_table`, the source table's connector, location and version. + +`TrainingTable.id` is deterministic: the same snapshot, query, plan, split seed and engine build give the same id. A caller-supplied source without a version gets a random id. An export to files (`TrainingTable.export`) writes the manifest next to the Parquet parts. `TrainingTable.load_from_connector(connector, path=..., pq=..., snapshot=...)` materializes an export again. + +With MLOps on, each table is stored once under `/training_tables//`. + +### Models + +Every fitted bundle carries `_provenance.json` (`ModelManifest`): +- its `id`, the bundle id (a content hash that leaves the manifest file out); +- when it was fitted and with which fit mode; +- its parents: the training table's manifest, embedded with that table's snapshot, and the warm-start model's manifest, if any. + +An embedded model keeps only the ids of its own model parents, so a long chain of warm-start retrains does not grow each manifest. `fit` returns a new `TrainingResult` per call (`.id`, `.manifest`); `Trainer.result` is the latest fit's result or the loaded bundle's (`trainer.result.id`, `trainer.result.manifest`). + +With MLOps on, each bundle is stored under `/models//`. A complete copy that is already there is kept; a partial one is replaced. + +The bundle also carries the published model's holdout predictions (`holdout.parquet`, with its truncation flag), so `result.holdout_df()` works on a model loaded with `Trainer.load` or `Trainer.load_from_connector` (a path or a registry), not only on the result of `fit`. Bundles fitted before this have no copy, and `holdout_df()` raises `FileNotFoundError` for them. + +### Moving snapshots between machines + +To build an index on one machine and train or score on another, export it to a file destination and load it there: + +```python +from kumo.engine import GraphSnapshot, S3Connector, SnapshotOutputConfig + +location = graph.snapshot().export(SnapshotOutputConfig(S3Connector(), "s3://bucket/snapshots")) +# On the training machine: +snapshot = GraphSnapshot.load_from_connector(S3Connector(), path=location) +``` + +- **Layout:** `export` writes `///`, the two halves of `snapshot.id`, and returns that location. It holds `kumo_snapshot.json` and one directory per export generation (`index/` and `_provenance.json`). +- **Transfer:** files move in parallel with a CRC32 per range. Each export uploads into a new generation and commits by writing `kumo_snapshot.json` last, naming that generation. Until then the previous export stays loadable, and concurrent exports never mix files: the last commit wins whole. Superseded generations, and abandoned ones over a day old, are removed best-effort. A load either completes or leaves nothing; if a newer export commits while it downloads, it switches to that one. +- **Reuse:** loading again reuses the local copy only if it holds this export's bytes (the same range checksums) and passes DiskGraph's integrity check; otherwise it is downloaded again. Exporting again skips the upload when the destination already holds a complete export (its `kumo_snapshot.json` and `_DONE`) with the local index's bytes, and logs that it reused it. +- **Staleness:** `snapshot.source_changes(graph)` reports, per table of `graph`, whether its source is `unchanged`, `changed` or `unknown` since the snapshot was built. `load_from_connector(..., graph=graph)` warns when the snapshot was taken of another graph, naming the schema changes (tables, columns, types, keys, edges), or when any of its tables changed. A graph id also covers each table's connector and location, so a graph that differs only there is logged at info level, not warned about. +- **Manifest:** the snapshot's `_provenance.json` travels with the export. +- **Trust:** load only from paths whose writers you trust. Checksums detect corruption, not tampering. +- **Destinations:** any destination whose sink holds files at the path works: S3, GCS, local files, or a volume or stage once those sinks are installed. A warehouse connector raises. + +With MLOps on (`graph.snapshot(mlops=...)` or a `kumo.mlops` block), each snapshot is exported under `/snapshots/`, unless a complete export of it is already there. + +### Scoring without the graph + +A scoring job does not need the training job's `Graph` code. It loads the snapshot the training job exported and the model, then scores: + +```python +snapshot = GraphSnapshot.load_from_connector(store, path=location) +trainer = Trainer.load_from_connector(store, path="models/") +result = trainer.predict(snapshot, binary_classification_threshold=0.5) +``` + +- `predict(snapshot)` reuses the trained query and table-generation plan, and scores its population at `prediction_time`. +- `entities=` scores given rows instead. +- It warns when the snapshot was taken of a different graph than the one the model was trained on, naming the schema changes, for example `Schema changes: events edges.` A graph that differs only in table connectors or locations (a dev graph and a prod graph of the same schema) is logged at info level instead. The model records its training graph from this release on; for an older model the warning cannot say what differs. +- `predict(prediction_table, snapshot=...)` scores a prediction table generated beforehand, as before. +- `fit` results score too: `result.predict(snapshot, ...)`. diff --git a/packages/kumo-engine/kumo/engine/__init__.py b/packages/kumo-engine/kumo/engine/__init__.py index 4d89af512..db8c0743e 100644 --- a/packages/kumo-engine/kumo/engine/__init__.py +++ b/packages/kumo-engine/kumo/engine/__init__.py @@ -25,6 +25,12 @@ from kumoapi.plugin import ModelRegistryLocation from kumoapi.table import Column +from kumo.engine._manifest import ( + ModelManifest, + SnapshotManifest, + TrainingTableManifest, +) +from kumo.engine._mlops import MLOps, mlops from kumo.engine._version import __version__ from kumo.engine.config import configure from kumo.engine.connector import ( @@ -51,6 +57,7 @@ ModelOutputConfig, PredictionOutputConfig, PredictionResult, + SnapshotOutputConfig, Trainer, TrainingResult, TrainingTableOutputConfig, @@ -83,6 +90,11 @@ "Edge", "Graph", "GraphSnapshot", + "MLOps", + "mlops", + "SnapshotManifest", + "TrainingTableManifest", + "ModelManifest", "PredictiveQuery", "TrainingTable", "PredictionTable", @@ -92,6 +104,7 @@ "ModelOutputConfig", "PredictionOutputConfig", "TrainingTableOutputConfig", + "SnapshotOutputConfig", ] ensure_nofile() diff --git a/packages/kumo-engine/kumo/engine/_manifest.py b/packages/kumo-engine/kumo/engine/_manifest.py new file mode 100644 index 000000000..a6e8a871d --- /dev/null +++ b/packages/kumo-engine/kumo/engine/_manifest.py @@ -0,0 +1,290 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared JSON manifest format for engine artifacts.""" + +import json +import types +import warnings +from collections.abc import Mapping as AbcMapping +from dataclasses import dataclass, field, fields, replace +from datetime import datetime, timezone +from pathlib import Path +from typing import ( + Any, + Dict, + List, + Mapping, + Optional, + Type, + TypeVar, + Union, + get_args, + get_origin, + get_type_hints, +) + +from kumo_core.io.filesystem import atomic_write_bytes + +SCHEMA_VERSION = 1 +_OMIT_NONE = 'omit_none' + +R = TypeVar('R', bound='Record') +M = TypeVar('M', bound='Manifest') + + +def omit_none() -> Any: + """A field written only when it is not None, or read as an explicit + null (kept in ``extra``).""" + return field(default=None, metadata={_OMIT_NONE: True}) + + +@dataclass(frozen=True, kw_only=True) +class Record: + """A JSON object; keys this version does not know are kept in ``extra``.""" + + extra: Mapping[str, Any] = field(default_factory=dict, repr=False, + compare=False) + + def to_dict(self) -> Dict[str, Any]: + data = dict(self.extra) + for item in fields(self): + value = getattr(self, item.name) + if item.name != 'extra' and not (value is None + and item.metadata.get(_OMIT_NONE)): + data[item.name] = _plain(value) + return data + + @classmethod + def from_dict(cls: Type[R], data: Mapping[str, Any]) -> R: + hints = get_type_hints(cls) + names = {item.name for item in fields(cls)} - {'extra'} + nullable = { + item.name + for item in fields(cls) if item.metadata.get(_OMIT_NONE) + } + return cls( + **{ + name: _typed(hints[name], data[name]) + for name in names if name in data + }, extra={ + key: value + for key, value in data.items() + if key not in names or (key in nullable and value is None) + }) + + +def _plain(value: Any) -> Any: + if isinstance(value, Record): + return value.to_dict() + if isinstance(value, AbcMapping): + return {key: _plain(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_plain(item) for item in value] + return value + + +def _typed(hint: Any, value: Any) -> Any: + origin, args = get_origin(hint), get_args(hint) + if value is None: + return None + if origin in (Union, types.UnionType): + return _typed(next(arg for arg in args if arg is not type(None)), value) + if isinstance(hint, type) and issubclass(hint, Manifest): + return _parent(hint, value) + if isinstance(hint, type) and issubclass(hint, Record): + return hint.from_dict(value) + if origin in (dict, AbcMapping) and args: + return {key: _typed(args[1], item) for key, item in value.items()} + if origin is list and args: + return [_typed(args[0], item) for item in value] + return value + + +@dataclass(frozen=True, kw_only=True) +class Environment(Record): + kumo_engine: Optional[str] + build_commit: Optional[str] + kumo_core: Optional[str] + diskgraph: Optional[str] + diskgraph_build_commit: Optional[str] + torch: Optional[str] + python: Optional[str] + + +@dataclass(frozen=True, kw_only=True) +class SourceTableEntry(Record): + connector: Optional[str] + location: Optional[str] + version: Optional[str] + version_kind: str + dependencies: Optional[Dict[str, str]] = omit_none() + location_fingerprint: Optional[str] = omit_none() + normalized_sql_fingerprint: Optional[str] = omit_none() + sql: Optional[str] = omit_none() + dialect: Optional[str] = omit_none() + namespace: Optional[List[Any]] = omit_none() + + +#: The artifact kinds, as manifests and the artifact root name them. +KINDS = ('snapshot', 'training_table', 'model') + + +@dataclass(frozen=True, kw_only=True) +class Manifest(Record): + """An artifact's record. A stub (a parent with no record of its own) + sets only ``kind``, ``id`` and ``reconstructed``.""" + + kind: str = '' + schema_version: int = SCHEMA_VERSION + id: str + created_at: Optional[str] = None + environment: Optional[Environment] = None + reconstructed: Optional[bool] = omit_none() + + #: Fields left out where this manifest is embedded in a child's. + EMBED_OMITS = () + + +def stamp() -> Dict[str, Any]: + """Envelope fields for a manifest written now.""" + from kumo.engine._provenance import environment + + return dict( + created_at=datetime.now(timezone.utc).isoformat(timespec='seconds'), + environment=environment()) + + +#: Envelope fields for a manifest rebuilt after the fact. +RECONSTRUCTED: Dict[str, Any] = dict(created_at=None, environment=None, + reconstructed=True) + + +def _kinds() -> Dict[str, Type['Manifest']]: + found: Dict[str, Type[Manifest]] = {} + pending = list(Manifest.__subclasses__()) + while pending: + cls = pending.pop() + found[next(str(item.default) for item in fields(cls) + if item.name == 'kind')] = cls + pending.extend(cls.__subclasses__()) + return found + + +def stub(kind: str, id: str) -> 'Manifest': + """The manifest of a ``kind`` artifact known only by its ``id``.""" + return _kinds()[kind](id=id, reconstructed=True) + + +def embedded(manifest: M) -> M: + """``manifest`` as a child embeds it: without its statistics and the + fields that change after it is written.""" + return replace(manifest, **{ + name: None + for name in manifest.EMBED_OMITS + }) + + +def _parent(hint: Type[M], value: Any) -> M: + try: + return hint.from_dict(value) + except (AttributeError, TypeError, ValueError, KeyError) as error: + identifier = value.get('id') if isinstance(value, AbcMapping) else None + warnings.warn(f"Reading an unreadable parent manifest as a stub: " + f"{error}", UserWarning, stacklevel=2) + return hint(id=str(identifier), reconstructed=True) + + +@dataclass(frozen=True, kw_only=True) +class SnapshotParents(Record): + source_tables: Dict[str, SourceTableEntry] + + +@dataclass(frozen=True, kw_only=True) +class SnapshotManifest(Manifest): + """What a graph snapshot was built from: each source table's connector, + location and version, and the engine build that wrote it.""" + + kind: str = 'snapshot' + parents: Optional[SnapshotParents] = None + record_sources: Optional[bool] = omit_none() + + +@dataclass(frozen=True, kw_only=True) +class TrainingTableParents(Record): + snapshot: SnapshotManifest + source_table: Optional[SourceTableEntry] = omit_none() + + +@dataclass(frozen=True, kw_only=True) +class TrainingTableManifest(Manifest): + """How a training table was generated: from which snapshot (embedded) + and, for caller-supplied rows, which source table.""" + + kind: str = 'training_table' + parents: Optional[TrainingTableParents] = None + query: Optional[str] = None + generation_plan: Optional[Dict[str, Any]] = None + task_type: Optional[str] = None + split_seed: Optional[int] = None + + +@dataclass(frozen=True, kw_only=True) +class ModelParents(Record): + training_table: Optional[TrainingTableManifest] = None + warm_start: Optional['ModelManifest'] = None + + +@dataclass(frozen=True, kw_only=True) +class FitTiming(Record): + started_at: str + finished_at: str + seconds: float + + +@dataclass(frozen=True, kw_only=True) +class ModelManifest(Manifest): + """How a model bundle was fitted: from which training table (embedded, + with its snapshot) and warm-start model, when, and with which fit + mode. Its ``id`` is the bundle id.""" + + kind: str = 'model' + parents: Optional[ModelParents] = None + fit: Optional[FitTiming] = None + fit_mode: Optional[str] = None + + +def manifest_bytes(manifest: Manifest) -> bytes: + return (json.dumps(manifest.to_dict(), indent=2, sort_keys=True) + + '\n').encode('utf-8') + + +def write_manifest(path: Path, manifest: Manifest) -> None: + atomic_write_bytes(path, manifest_bytes(manifest)) + + +def read_manifest(path: Path, kind: Type[M]) -> Optional[M]: + try: + manifest = json.loads(path.read_text(encoding='utf-8')) + except FileNotFoundError: + return None + except (OSError, ValueError) as error: + warnings.warn(f"Ignoring unreadable manifest {path}: {error}", + UserWarning, stacklevel=3) + return None + version = (manifest.get('schema_version') + if isinstance(manifest, dict) else None) + if type(version) is not int or not 1 <= version <= SCHEMA_VERSION: + warnings.warn( + f"Ignoring {path}: manifest schema_version {version!r} is not " + f"supported by this kumo-engine (up to {SCHEMA_VERSION}).", + UserWarning, + stacklevel=3, + ) + return None + try: + return kind.from_dict(manifest) + except (AttributeError, TypeError, ValueError) as error: + warnings.warn(f"Ignoring unreadable manifest {path}: {error}", + UserWarning, stacklevel=3) + return None diff --git a/packages/kumo-engine/kumo/engine/_mlops.py b/packages/kumo-engine/kumo/engine/_mlops.py new file mode 100644 index 000000000..3c7b16463 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/_mlops.py @@ -0,0 +1,105 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""The MLOps switch: :class:`MLOps` settings and the :func:`mlops` block.""" + +import contextlib +import contextvars +from dataclasses import dataclass, field +from typing import Any, FrozenSet, Iterable, Iterator, Optional, Tuple, Union + +from kumo.engine._manifest import KINDS + + +@dataclass(frozen=True) +class _Root: + connector: Any + path: str + _uri: str + + def location(self, relative: str) -> str: + r"""``relative`` as a path on :attr:`connector`.""" + base = self.path.rstrip('/') + return f'{base}/{relative}' if base else relative + + def uri(self, relative: str) -> str: + r"""``relative`` as a full URI or path, for messages and records.""" + return f"{self._uri.rstrip('/')}/{relative}" + + +@dataclass(frozen=True) +class MLOps: + r"""Where artifacts go, for :func:`mlops` blocks and ``mlops=`` + arguments. + + Args: + artifact_root: an ``s3://``, ``gs://``, ``file://`` or local path + (``/Volumes/...`` on a Databricks cluster), or a ``(connector, + path)`` pair for a store a URI cannot reach, such as a volume + through a ``DatabricksConnector`` or a Snowflake stage. + export: the artifact kinds stored under the root when produced + (``snapshot``, ``training_table``, ``model``); + ``True`` for all, ``False`` for none. + """ + + artifact_root: Union[str, Tuple[Any, str]] + export: Union[bool, Iterable[str]] = field(default=True, kw_only=True) + root: _Root = field(init=False, repr=False, compare=False) + exports: FrozenSet[str] = field(init=False, repr=False, compare=False) + + def __post_init__(self) -> None: + from kumo.engine.connector import ( + connector_for_root, + is_file_destination, + ) + + if not self.artifact_root: + raise ValueError('MLOps needs an artifact_root.') + if isinstance(self.artifact_root, str): + root = _Root(connector_for_root(self.artifact_root), '', + self.artifact_root) + else: + connector, path = self.artifact_root + if not is_file_destination(connector, path): + raise ValueError( + f'{type(connector).__name__} cannot hold files at ' + f'{path!r}; artifact_root needs a file location.') + root = _Root(connector, path, path) + kinds = (frozenset(KINDS) if self.export is True else + frozenset() if self.export is False else frozenset( + self.export)) + unknown = sorted(kinds - set(KINDS)) + if unknown: + raise ValueError(f'Unknown export kinds {unknown}; choose from ' + f'{list(KINDS)}.') + object.__setattr__(self, 'root', root) + object.__setattr__(self, 'exports', kinds) + + def wants(self, kind: str) -> bool: + r"""Whether artifacts of ``kind`` are stored under the root.""" + return kind in self.exports + + +_CURRENT: contextvars.ContextVar[Optional[MLOps]] = contextvars.ContextVar( + 'kumo_mlops', default=None) + + +@contextlib.contextmanager +def mlops(settings: MLOps) -> Iterator[MLOps]: + r"""Turns MLOps on for the calls in the block (in this thread or asyncio + task): each snapshot, training table and model is stored + under ``settings.artifact_root``. An inner block replaces an outer one; + an ``mlops=`` argument replaces both for its call.""" + if not isinstance(settings, MLOps): + raise TypeError('kumo.mlops() takes a kumo.MLOps') + token = _CURRENT.set(settings) + try: + yield settings + finally: + _CURRENT.reset(token) + + +def resolve(settings: Optional[MLOps] = None) -> Optional[MLOps]: + r"""The settings a call runs with: its ``mlops=`` argument, else the + enclosing block's, else None (off).""" + return settings if settings is not None else _CURRENT.get() diff --git a/packages/kumo-engine/kumo/engine/_provenance.py b/packages/kumo-engine/kumo/engine/_provenance.py new file mode 100644 index 000000000..7f0bef9eb --- /dev/null +++ b/packages/kumo-engine/kumo/engine/_provenance.py @@ -0,0 +1,137 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Runtime versions and source tables recorded in manifests.""" + +import hashlib +import platform +import subprocess +import warnings +from functools import cache +from pathlib import Path +from typing import Any, Dict, Optional + +from kumo.engine._manifest import Environment, SourceTableEntry + +_LOCATED_CONNECTORS = frozenset( + {'LOCAL_FILE', 'S3', 'GCS', 'SNOWFLAKE', 'DATABRICKS', 'BIGQUERY'}) + + +@cache +def _build_commit(version: str) -> Optional[str]: + if '+' in version: + return version.split('+', 1)[1] + from kumo.engine import _version + stamped = getattr(_version, 'BUILD_COMMIT', None) + if stamped: + return stamped + root = Path(__file__).resolve().parents[2] + if not (root / 'pyproject.toml').is_file(): + return None + command = ['git', '-C', str(root), 'rev-parse', '--short', 'HEAD'] + try: + return subprocess.run(command, capture_output=True, text=True, + check=True, timeout=10).stdout.strip() or None + except (OSError, subprocess.SubprocessError): + return None + + +def _diskgraph() -> Any: + try: + import diskgraph + return diskgraph + except ImportError: + return None + + +def environment() -> Environment: + from kumo_core.version import ( + get_kumo_version, + installed_distribution_version, + ) + + from kumo.engine._version import __version__ + + diskgraph = _diskgraph() + return Environment( + kumo_engine=get_kumo_version(), + build_commit=_build_commit(__version__), + kumo_core=installed_distribution_version('kumo-core'), + diskgraph=getattr(diskgraph, '__version__', None), + diskgraph_build_commit=getattr(diskgraph, 'BUILD_COMMIT', None), + torch=installed_distribution_version('torch'), + python=platform.python_version(), + ) + + +def source_location(connector: Any, table_name: str, + sql: Optional[str] = None) -> str: + if sql is not None: + return 'sql:' + hashlib.sha256(sql.encode()).hexdigest() + if _connector_kind(connector) not in _LOCATED_CONNECTORS: + return table_name + return connector.source_backend.table_location(connector, table_name) + + +def _connector_kind(connector: Any) -> str: + kind = connector.source_type + return str(getattr(kind, 'value', kind)) + + +def source_table(connector: Any, table_name: str, sql: Optional[str], + version: Any, *, + record_sources: bool = True) -> SourceTableEntry: + from kumo_core.sources.identity import version_kind + + entry: Dict[str, Any] = {'connector': None, 'location': None} + dependencies = getattr(version, 'dependencies', None) + try: + entry['connector'] = _connector_kind(connector) + if record_sources or sql is not None: + entry['location'] = source_location(connector, table_name, sql) + else: + entry['location'] = table_name + entry['location_fingerprint'] = _location_fingerprint( + entry['connector'], connector, table_name) + if sql is not None: + _record_sql(entry, connector, sql, record_sources) + except Exception as error: + warnings.warn(f"Could not record a source table's location: {error}", + UserWarning, stacklevel=3) + return SourceTableEntry( + version=None if version is None else str(version), + version_kind=version_kind(version), + dependencies=None if dependencies is None else { + name: str(value) + for name, value in dependencies.items() + }, + **entry, + ) + + +def _location_fingerprint(kind: str, connector: Any, + table_name: str) -> Optional[str]: + try: + location = source_location(connector, table_name, None) + except Exception: + return None + return hashlib.sha256(f'{kind}:{location}'.encode()).hexdigest() + + +def _record_sql(entry: Dict[str, Any], connector: Any, sql: str, + record_sources: bool) -> None: + from kumo_core.custom_sql import ( + CUSTOM_SQL_NAMESPACE_ATTRS, + custom_sql_synthetic_name, + ) + from kumo_core.types import ConnectorType + + kind = ConnectorType[entry['connector']] + entry['normalized_sql_fingerprint'] = custom_sql_synthetic_name(sql, kind) + if record_sources: + entry['sql'] = sql + entry['dialect'] = entry['connector'].lower() + entry['namespace'] = [ + getattr(connector.config, attr) + for attr in CUSTOM_SQL_NAMESPACE_ATTRS[kind] + ] diff --git a/packages/kumo-engine/kumo/engine/_snapshot_io.py b/packages/kumo-engine/kumo/engine/_snapshot_io.py new file mode 100644 index 000000000..7f83970c2 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/_snapshot_io.py @@ -0,0 +1,557 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import hashlib +import json +import logging +import os +import re +import threading +import time +import uuid +from concurrent.futures import ThreadPoolExecutor +from functools import partial +from pathlib import Path +from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple + +from pyarrow.fs import FileSystem + +_FORMAT = 1 +_MANIFEST = 'kumo_snapshot.json' +_INDEX_MANIFEST = '_MANIFEST.json' +_RANGE_BYTES = 64 << 20 +_CHUNK = 8 << 20 +_ATTEMPTS = 3 +_KEY = re.compile('[0-9a-f]{64}') +_GENERATION = re.compile('[0-9a-f]{32}') +_EXPORT_IDENTITY = '.kumo_export' +_PROGRESS_SECONDS = 30 +_DONE = '_DONE' +_ABANDONED_SECONDS = 24 * 3600 + +logger = logging.getLogger(__name__) + + +def _location(connector: Any, path: str) -> Tuple[FileSystem, str]: + from kumo.engine.connector import file_store + + store = file_store(connector, path) + if store is None: + raise NotImplementedError( + f'{type(connector).__name__} cannot hold a snapshot export: export ' + 'needs a file destination (artifact root), such as an S3Connector, ' + 'GCSConnector or a FileConnector on a volume.') + return store + + +def _files(filesystem: FileSystem) -> Any: + from kumo_core.object_store_sink import _ArrowExportFiles + + return _ArrowExportFiles(filesystem, None) + + +def _check_key(graph_key: Any, data_version: Any, snapshot_id: Any) -> None: + from kumo.engine.graph import _snapshot_id + + if not (isinstance(graph_key, str) and isinstance(data_version, str) + and _KEY.fullmatch(graph_key) and _KEY.fullmatch(data_version) + and snapshot_id == _snapshot_id(graph_key, data_version)): + raise ValueError(f'{snapshot_id!r} is not a snapshot id.') + + +class _Client: + def __init__(self, connector: Any, path: str) -> None: + self._make = lambda: _location(connector, path)[0] + self.current = self._make() + self._lock = threading.Lock() + + def renew(self, stale: FileSystem) -> None: + with self._lock: + if self.current is stale: + self.current = self._make() + + +class _Progress: + def __init__(self, action: str, path: str, total: int) -> None: + self.action, self.path, self.total = action, path, total + self.moved = 0 + self.started = self.logged = time.monotonic() + self._lock = threading.Lock() + + def add(self, size: int) -> None: + with self._lock: + self.moved += size + now = time.monotonic() + if (now - self.logged < _PROGRESS_SECONDS + and self.moved < self.total): + return + self.logged = now + logger.info('Snapshot %s of %s: %d of %d bytes, %.2f GB/s', + self.action, self.path, self.moved, self.total, + self.moved / max(now - self.started, 1e-9) / 1e9) + + +def _transfer(connector: Any, path: str, items: Sequence[Any], + sizes: Sequence[int], work: Callable[..., Any], + action: str) -> List[Any]: + from kumo_core.sources.object_store import copies + from kumo_core.sources.parallel import run_in_parallel + + if not items: + return [] + stop = threading.Event() + with copies(items) as workers: + # After `copies`: Arrow's IO pool size sets the connection pool. + client = _Client(connector, path) + logger.info('Snapshot %s of %s: %d files, %d bytes, %d workers', + action, path, len(items), sum(sizes), workers) + progress = _Progress(action, path, sum(sizes)) + + def run(item: Any, size: int) -> Any: + result = _retried(client, work, item, stop) + progress.add(size) + return result + + return run_in_parallel( + [partial(run, item, size) for item, size in zip(items, sizes)], + workers=workers, stop=stop) + + +def _retried(client: _Client, work: Callable[..., Any], item: Any, + stop: threading.Event) -> Any: + attempt = 0 + while True: + _stopped(stop) + current = client.current + try: + return work(current, item, stop) + except (FileNotFoundError, PermissionError): + raise + except OSError as error: + if attempt + 1 == _ATTEMPTS: + raise + logger.warning('Retrying a snapshot transfer after: %s', error) + # Arrow keeps the credentials a client was built with. + client.renew(current) + stop.wait(2**attempt) + attempt += 1 + + +def _stopped(stop: threading.Event) -> None: + from kumo_core.sources.parallel import WorkCancelled + + if stop.is_set(): + raise WorkCancelled('abandoned after another transfer failed') + + +def _read_json(filesystem: FileSystem, path: str) -> Optional[Any]: + try: + return json.loads(_files(filesystem).read_text(path)) + except FileNotFoundError: + return None + + +def _open_output(filesystem: FileSystem, path: str, size: int) -> Any: + from kumo_core.object_store_sink import open_sized + + _files(filesystem).makedirs(path.rsplit('/', 1)[0]) + return open_sized(filesystem, path, size) + + +def _write(filesystem: FileSystem, path: str, payload: bytes) -> None: + with _open_output(filesystem, path, len(payload)) as writer: + writer.write(payload) + + +def _fingerprint(index_manifest: bytes) -> str: + return hashlib.sha256(index_manifest).hexdigest() + + +def _snapshot_manifest(snapshot: Any) -> bytes: + from kumo.engine._manifest import ( + RECONSTRUCTED, + SnapshotManifest, + SnapshotParents, + ) + from kumo.engine.graph import SNAPSHOT_MANIFEST_FILE + + try: + return (Path(snapshot.index_path) / SNAPSHOT_MANIFEST_FILE).read_bytes() + except OSError: + manifest = SnapshotManifest( + id=snapshot.id, parents=SnapshotParents(source_tables={}), + **RECONSTRUCTED) + return (json.dumps(manifest.to_dict(), indent=2, sort_keys=True) + + '\n').encode() + + +def export_location(snapshot: Any, path: str) -> str: + graph_key, data_version = snapshot._graph_key, snapshot._data_version + _check_key(graph_key, data_version, snapshot.id) + return f"{path.rstrip('/')}/{graph_key}/{data_version[:16]}" + + +def store(snapshot: Any, settings: Any) -> None: + r"""Export ``snapshot`` under ``/snapshots/`` unless a complete + export of it is there; failures warn.""" + import warnings + + try: + connector = settings.root.connector + path = settings.root.location('snapshots') + filesystem, root = _location(connector, + export_location(snapshot, path)) + try: + marker = _read_json(filesystem, f'{root}/{_MANIFEST}') + except ValueError: + marker = None + if not (isinstance(marker, dict) and marker.get('format') == _FORMAT + and marker.get('snapshot_id') == snapshot.id): + export(snapshot, connector, path) + except Exception as error: + warnings.warn( + f"Could not store snapshot {snapshot.id} under the artifact " + f"root: {error}", UserWarning, stacklevel=3) + + +def export(snapshot: Any, connector: Any, path: str) -> str: + from diskgraph import snapshot_range_checksums + + from kumo.engine._provenance import environment + from kumo.engine.graph import SNAPSHOT_MANIFEST_FILE + + source = Path(snapshot.index_path) + graph_key, data_version = snapshot._graph_key, snapshot._data_version + location = export_location(snapshot, path) + snapshot_manifest = _snapshot_manifest(snapshot) + if _holds_copy(connector, location, snapshot.id, source, + snapshot_manifest): + logger.info('%s already holds an identical export of snapshot %s; ' + 'reused it', location, snapshot.id) + return location + index_manifest = (source / _INDEX_MANIFEST).read_bytes() + sizes = { + entry['rel_path']: int(entry['size_bytes']) + for entry in json.loads(index_manifest)['files'] + } + # Each export writes its own generation; the marker, written last, + # names the one that loads. Concurrent exports never share a file. + generation = uuid.uuid4().hex + filesystem, root = _location(connector, location) + base = f'{root}/{generation}' + + def upload(client: FileSystem, name: str, stop: threading.Event) -> None: + size = sizes[name] + with open(source / name, 'rb', buffering=0) as reader, _open_output( + client, f'{base}/index/{name}', size) as writer: + for offset in range(0, size, _CHUNK): + _stopped(stop) + chunk = reader.read(min(_CHUNK, size - offset)) + if len(chunk) != min(_CHUNK, size - offset): + raise OSError(f'{source / name} changed while it was ' + 'exported') + writer.write(chunk) + + names = sorted(sizes, key=lambda name: -sizes[name]) + with ThreadPoolExecutor(1) as pool: + checksums = pool.submit(snapshot_range_checksums, str(source), + _RANGE_BYTES) + _transfer(connector, location, names, + [sizes[name] for name in names], upload, 'export') + crc32 = checksums.result() + _write(filesystem, f'{base}/index/{_INDEX_MANIFEST}', index_manifest) + _write(filesystem, f'{base}/{SNAPSHOT_MANIFEST_FILE}', snapshot_manifest) + marker = { + 'format': _FORMAT, + 'snapshot_id': snapshot.id, + 'generation': generation, + 'graph_key': graph_key, + 'data_version': data_version, + 'index_manifest': _fingerprint(index_manifest), + 'graph_config': snapshot.graph_config, + 'environment': environment().to_dict(), + 'range_bytes': _RANGE_BYTES, + 'crc32': crc32, + } + _commit(filesystem, root, generation, json.dumps(marker).encode()) + try: + (source / _EXPORT_IDENTITY).write_text(export_identity(marker)) + except OSError: + pass + return location + + +def _holds_copy(connector: Any, location: str, snapshot_id: str, + index: Path, snapshot_manifest: bytes) -> bool: + r"""Whether ``location`` holds a complete export of this snapshot: the + local index's bytes, every index file at its size, and the same + ``_provenance.json``.""" + from pyarrow.fs import FileSelector, FileType + + from kumo.engine.graph import SNAPSHOT_MANIFEST_FILE + + try: + root, manifest = read(connector, location) + if (manifest.get('snapshot_id') != snapshot_id + or not holds_export(index, manifest)): + return False + filesystem, _ = _location(connector, location) + remote = { + info.path[len(root) + 1:]: info.size + for info in filesystem.get_file_info( + FileSelector(root, recursive=True, allow_not_found=True)) + if info.type is FileType.File + } + listed = json.loads((index / _INDEX_MANIFEST).read_bytes())['files'] + if (_DONE not in remote + or f'index/{_INDEX_MANIFEST}' not in remote) or any( + remote.get(f"index/{entry['rel_path']}") != int( + entry['size_bytes']) for entry in listed): + return False + with filesystem.open_input_stream( + f'{root}/{SNAPSHOT_MANIFEST_FILE}') as reader: + return reader.readall() == snapshot_manifest + except Exception: + return False + + +def _commit(filesystem: FileSystem, root: str, generation: str, + marker: bytes) -> None: + r"""Writes the marker in one step: a rename on a local filesystem (an + in-place write there truncates first), a single object put elsewhere.""" + if filesystem.type_name == 'local': + staged = f'{root}/{generation}/{_MANIFEST}' + _write(filesystem, staged, marker) + filesystem.move(staged, f'{root}/{_MANIFEST}') + else: + _write(filesystem, f'{root}/{_MANIFEST}', marker) + _write(filesystem, f'{root}/{generation}/{_DONE}', b'') + _clean(filesystem, root) + + +def _clean(filesystem: FileSystem, root: str) -> None: + r"""Best effort: remove generations that a later commit superseded, and + abandoned ones. A generation is marked done only after its commit, so + one listed done before the marker is re-read lost to a later one.""" + from pyarrow.fs import FileSelector, FileType + + try: + newest: Dict[str, float] = {} + done = set() + for info in filesystem.get_file_info( + FileSelector(root, recursive=True, allow_not_found=True)): + name, _, rest = info.path[len(root) + 1:].partition('/') + if (not rest or info.type is not FileType.File + or not _GENERATION.fullmatch(name)): + continue + mtime = info.mtime.timestamp() if info.mtime else time.time() + newest[name] = max(newest.get(name, 0.0), mtime) + if rest == _DONE: + done.add(name) + marker = _read_json(filesystem, f'{root}/{_MANIFEST}') or {} + keep = marker.get('generation') + stale = time.time() - _ABANDONED_SECONDS + for name, mtime in newest.items(): + if name != keep and (name in done or mtime < stale): + filesystem.delete_dir(f'{root}/{name}') + except Exception as error: + logger.warning('Could not remove old snapshot exports under %s: %s', + root, error) + + +def _corrupt(path: str, detail: str) -> ValueError: + return ValueError(f'The snapshot export at {path} is corrupt ({detail}); ' + 'delete it and run GraphSnapshot.export again.') + + +def read(connector: Any, path: str) -> Tuple[str, dict]: + filesystem, root = _location(connector, path) + try: + manifest = _read_json(filesystem, f'{root}/{_MANIFEST}') + except ValueError as error: + raise _corrupt(path, f'unreadable {_MANIFEST}: {error}') from error + if manifest is None: + raise FileNotFoundError(f'No complete snapshot export at {path}; ' + 're-run export.') + if not isinstance(manifest, dict): + raise _corrupt(path, 'manifest is not a JSON object') + if manifest.get('format') != _FORMAT: + raise ValueError( + f"Snapshot export at {path} has format " + f"{manifest.get('format')!r}; this kumo-engine reads format " + f"{_FORMAT}. Load it with the kumo-engine that exported it.") + _check_key(manifest.get('graph_key'), manifest.get('data_version'), + manifest.get('snapshot_id')) + span, crcs = manifest.get('range_bytes'), manifest.get('crc32') + generation = manifest.get('generation') + if (type(span) is not int or span <= 0 or not isinstance(crcs, dict) + or not isinstance(manifest.get('index_manifest'), str) + or not isinstance(manifest.get('graph_config'), dict) + or not isinstance(generation, str) + or not _GENERATION.fullmatch(generation)): + raise _corrupt(path, 'malformed manifest') + return f'{root}/{generation}', manifest + + +def export_identity(manifest: dict) -> str: + r"""What makes one export's bytes differ from another's of the same + snapshot: its range checksums and index manifest.""" + return hashlib.sha256(json.dumps( + [manifest['range_bytes'], manifest['crc32'], + manifest['index_manifest']], sort_keys=True).encode()).hexdigest() + + +def holds_export(index: Path, manifest: dict) -> bool: + r"""Whether the local index at ``index`` has the export's bytes. An + index without a recorded export identity (built here, or loaded by an + older engine) is checksummed once.""" + from diskgraph import snapshot_range_checksums + + identity = export_identity(manifest) + try: + return (index / _EXPORT_IDENTITY).read_text() == identity + except FileNotFoundError: + pass + if snapshot_range_checksums(str(index), + manifest['range_bytes']) != manifest['crc32']: + return False + (index / _EXPORT_IDENTITY).write_text(identity) + return True + + +def _download(root: str, staging: Path, client: FileSystem, + item: Tuple[str, int, int], stop: threading.Event) -> None: + name, offset, length = item + # Each open is a HEAD: pyarrow cannot open a file from a known size. + with client.open_input_file(f'{root}/index/{name}') as reader: + fd = os.open(staging / name, os.O_WRONLY) + try: + for start in range(offset, offset + length, _CHUNK): + _stopped(stop) + want = min(_CHUNK, offset + length - start) + chunk = reader.read_at(want, start) + if len(chunk) != want: + raise ValueError(f'{name} is truncated') + os.pwrite(fd, chunk, start) + getattr(os, 'fdatasync', os.fsync)(fd) + finally: + os.close(fd) + + +def _index_files( + path: str, root: str, manifest: dict, connector: Any +) -> Tuple[bytes, List[Tuple[str, int]], Optional[bytes]]: + from kumo_core.sources.layout import confined + + from kumo.engine.graph import SNAPSHOT_MANIFEST_FILE + + filesystem, _ = _location(connector, path) + with filesystem.open_input_stream( + f'{root}/index/{_INDEX_MANIFEST}') as reader: + index_manifest = reader.readall() + if _fingerprint(index_manifest) != manifest['index_manifest']: + raise _corrupt(path, f'{_INDEX_MANIFEST} is not the exported one') + try: + files = [(confined(entry['rel_path']), int(entry['size_bytes'])) + for entry in json.loads(index_manifest)['files']] + except Exception as error: + raise _corrupt(path, f'bad {_INDEX_MANIFEST}: {error}') from error + if len({name for name, _ in files}) != len(files): + raise _corrupt(path, f'{_INDEX_MANIFEST} lists a file twice') + try: + with filesystem.open_input_stream( + f'{root}/{SNAPSHOT_MANIFEST_FILE}') as reader: + snapshot_manifest = reader.readall() + except FileNotFoundError: + snapshot_manifest = None + return index_manifest, files, snapshot_manifest + + +def load_into(connector: Any, path: str, root: str, manifest: dict, + target: Path) -> None: + from diskgraph import ( + DiskgraphCorruptionError, + snapshot_range_checksums, + verify_snapshot_integrity, + ) + from kumo_core.diskgraph.snapshot import ( + DiskgraphIngestResult, + write_diskgraph_success_marker, + ) + from kumo_core.io.filesystem import fsync_directory + from kumo_core.sources.transfer import discard_destination + + from kumo.engine.graph import SNAPSHOT_MANIFEST_FILE + + span = manifest['range_bytes'] + index_manifest, files, snapshot_manifest = _index_files( + path, root, manifest, connector) + staging = target.with_name(f'{target.name}.loading') + discard_destination(str(target)) + discard_destination(str(staging)) + try: + staging.mkdir(parents=True) + needed = sum(size for _, size in files) + disk = os.statvfs(staging) + if disk.f_bavail * disk.f_frsize < needed: + raise OSError(f'Loading this snapshot needs {needed} free bytes ' + f'in {staging}.') + for name, size in files: + (staging / name).parent.mkdir(parents=True, exist_ok=True) + fd = os.open(staging / name, os.O_WRONLY | os.O_CREAT | os.O_EXCL) + try: + if size: + os.posix_fallocate(fd, 0, size) + finally: + os.close(fd) + ranges = sorted(((name, offset, min(span, size - offset)) + for name, size in files + for offset in range(0, size, span)), + key=lambda item: -item[2]) + try: + _transfer(connector, path, ranges, [item[2] for item in ranges], + partial(_download, root, staging), 'load') + except ValueError as error: + raise _corrupt(path, str(error)) from error + (staging / _INDEX_MANIFEST).write_bytes(index_manifest) + write_diskgraph_success_marker(str(staging), + DiskgraphIngestResult.for_index(target)) + try: + verify_snapshot_integrity(str(staging)) + loaded = snapshot_range_checksums(str(staging), span) + except DiskgraphCorruptionError as error: + raise _corrupt(path, str(error)) from error + for name, crcs in sorted(loaded.items()): + expected = manifest['crc32'].get(name) + if crcs != expected: + offset = next((index * span for index, crc in enumerate(crcs) + if not expected or index >= len(expected) + or crc != expected[index]), 0) + raise _corrupt( + path, f'{name} failed its checksum at offset {offset}') + _check_readable(staging, manifest) + (staging / _EXPORT_IDENTITY).write_text(export_identity(manifest)) + if snapshot_manifest is not None: + (staging / SNAPSHOT_MANIFEST_FILE).write_bytes(snapshot_manifest) + for directory in sorted({(staging / name).parent + for name, _ in files} | {staging}): + fsync_directory(directory) + os.rename(staging, target) + fsync_directory(target.parent) + except BaseException: + discard_destination(str(staging)) + raise + + +def _check_readable(index: Path, manifest: dict) -> None: + from diskgraph import SamplerEngine + + try: + SamplerEngine(str(index), num_threads=1) + except Exception as error: + built = (manifest.get('environment') or {}).get('kumo_engine') + raise RuntimeError( + f'This DiskGraph cannot open the snapshot exported by kumo-engine ' + f'{built} ({error}). Rebuild it with this kumo-engine or load it ' + 'with the exporting one.') from error diff --git a/packages/kumo-engine/kumo/engine/connector/__init__.py b/packages/kumo-engine/kumo/engine/connector/__init__.py index 93265951d..fd745e0af 100644 --- a/packages/kumo-engine/kumo/engine/connector/__init__.py +++ b/packages/kumo-engine/kumo/engine/connector/__init__.py @@ -37,6 +37,24 @@ from kumo.engine.connector.source_table import SourceTable, SQLSourceTable +def connector_for_root(uri: str) -> Connector: + r"""A connector rooted at ``uri``: ``s3://``, ``gs://``, ``file://`` or a + local path.""" + from urllib.parse import urlparse + from urllib.request import url2pathname + + parsed = urlparse(uri) + if parsed.scheme == 's3': + return S3Connector(root_dir=uri) + if parsed.scheme == 'gs': + return GCSConnector(root_dir=uri) + if parsed.scheme not in ('', 'file'): + raise ValueError(f'{uri!r} is not s3://, gs://, file:// or a local ' + f'path.') + return FileConnector( + root_dir=url2pathname(parsed.path) if parsed.scheme else uri) + + def _file_sink(connector: Any, path: str) -> Any: from kumoapi.plugin import FileStoreSink @@ -93,4 +111,5 @@ def model_sink(connector: Any, path: Optional[str]) -> Any: "GCSIdentity", "BigQueryIdentity", "DatabricksIdentity", + "connector_for_root", ] diff --git a/packages/kumo-engine/kumo/engine/connector/file.py b/packages/kumo-engine/kumo/engine/connector/file.py index d9098a74e..be3839d2b 100644 --- a/packages/kumo-engine/kumo/engine/connector/file.py +++ b/packages/kumo-engine/kumo/engine/connector/file.py @@ -121,6 +121,11 @@ def get_table_config(self, connector: Union[Connector, SourceConnector], LocalFileSourceTable(path=str(Path(path).resolve()), file_type=FileType(kind.upper()))) + def table_location(self, connector: Union[Connector, SourceConnector], + table_name): + reader, name = self._reader(connector, table_name) + return str(Path(reader.resolve(name)[0]).resolve()) + def prepare_ingest(self, connector, table_name, dest_dir, *args, **kwargs): _disjoint(Path(connector.config.root_dir), Path(dest_dir)) return super().prepare_ingest(connector, table_name, dest_dir, *args, diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index b5ecf3299..04df470a2 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -60,24 +60,45 @@ """ import json +import logging +import threading import warnings -from dataclasses import dataclass, field -from functools import cached_property +from dataclasses import dataclass, field, replace +from functools import cached_property, partial from pathlib import Path -from typing import Any, Dict, Mapping, Sequence +from typing import TYPE_CHECKING, Any, Dict, List, Mapping, Optional, Sequence +from kumo_core.artifacts import PROVENANCE_FILE from kumo_core.pquery import NativeGraph, normalize_graph from kumoapi.graph import Edge +from kumo.engine._manifest import ( + RECONSTRUCTED, + SnapshotManifest, + SnapshotParents, + SourceTableEntry, + read_manifest, + stamp, + write_manifest, +) +from kumo.engine._provenance import source_table from kumo.engine._scratch import ( fingerprint, get_scratch_dir, snapshot_scratch_dir, ) from kumo.engine._storage import network_storage +from kumo.engine.connector.base import Connector + +if TYPE_CHECKING: + from kumo.engine.trainer.output_config import SnapshotOutputConfig from kumo.engine.table import Table +SNAPSHOT_MANIFEST_FILE = PROVENANCE_FILE + _warned_scratch: set[Path] = set() +_warned_backfill: set[Path] = set() +logger = logging.getLogger(__name__) def _warn_if_network_scratch(root: Path) -> None: @@ -122,6 +143,14 @@ def _warn_low_fk_match_rates(stats: Dict[str, Any], ) +def _ingest_stats(path: Any) -> Dict[str, Any]: + try: + stats = json.loads(Path(path).read_text())['stats'] + except (OSError, ValueError, KeyError, TypeError): + return {} + return stats if isinstance(stats, dict) else {} + + def _warn_ingest_data_loss(result: Any, graph_config: Dict[str, Any]) -> None: from kumo_core.diskgraph.validation import parse_failure_warning @@ -145,10 +174,8 @@ def _warn_ingest_data_loss(result: Any, graph_config: Dict[str, Any]) -> None: warnings.warn(parse_failure_warning(name, column, count), UserWarning, stacklevel=3) - try: - report = json.loads(Path(result.ingest_stats_path).read_text()) - stats = report['stats'] - except (AttributeError, OSError, ValueError, KeyError, TypeError): + stats = _ingest_stats(getattr(result, 'ingest_stats_path', None)) + if not stats: return tables = {table['table_name']: table for table in graph_config['tables']} duplicates: Dict[str, int] = {} @@ -183,6 +210,138 @@ def _warn_ingest_data_loss(result: Any, graph_config: Dict[str, Any]) -> None: ) +_LOCATION_KEYS = frozenset({'table_name', 'source_table_name', 'connector_id'}) +_SCHEMA_PARTS = { + 'cols': 'columns', + 'fkeys': 'edges', + 'pkey': 'primary key', + 'time_col': 'time column', + 'end_time_col': 'end time column', +} + + +def _schema_tables(config: Mapping[str, Any]) -> Dict[str, Dict[str, Any]]: + return { + table['table_name']: { + key: (sorted(json.dumps(item, sort_keys=True, default=str) + for item in value) if isinstance(value, list) else + value) + for key, value in table.items() if key not in _LOCATION_KEYS + } + for table in config.get('tables', []) + } + + +def schema_changes(before: Mapping[str, Any], + after: Mapping[str, Any]) -> List[str]: + r"""How two graph configs differ in tables, columns, types, keys and + edges; empty when only connectors or table locations differ.""" + old, new = _schema_tables(before), _schema_tables(after) + changes = [ + f"table {name} {'removed' if name in old else 'added'}" + for name in sorted(old.keys() ^ new.keys()) + ] + for name in sorted(old.keys() & new.keys()): + parts = sorted( + _SCHEMA_PARTS.get(key, key) + for key in old[name].keys() | new[name].keys() + if old[name].get(key) != new[name].get(key)) + if parts: + changes.append(f"{name} {', '.join(parts)}") + return changes + + +def report_other_graph(message: str, before: Optional[Mapping[str, Any]], + after: Mapping[str, Any], stacklevel: int) -> None: + if before is None: + warnings.warn( + f"{message} The model does not record its training graph, so " + "what differs is unknown; a different connector or table " + "location alone also changes the graph id.", UserWarning, + stacklevel=stacklevel + 1) + return + changes = schema_changes(before, after) + if changes: + warnings.warn(f"{message} Schema changes: {'; '.join(changes)}.", + UserWarning, stacklevel=stacklevel + 1) + else: + logger.info('%s The schema is the same; only table connectors or ' + 'locations differ.', message) + + +def _store() -> Any: + from kumo_core.diskgraph.snapshot import ( + DiskgraphLocalConfig, + LocalSnapshotStore, + ) + + return LocalSnapshotStore( + DiskgraphLocalConfig(data_root=str(get_scratch_dir() / 'snapshots'))) + + +def _lock(graph_key: str) -> Any: + from filelock import FileLock + + persist_dir = snapshot_scratch_dir(graph_key) / 'source_tables' + persist_dir.parent.mkdir(parents=True, exist_ok=True) + return FileLock(f'{persist_dir}.lock') + + +def _loaded(store: Any, handle: Any) -> Optional[Any]: + from diskgraph import DiskgraphCorruptionError, read_success_marker + from kumo_core.diskgraph.snapshot import rebased + + path = store.path(handle) + if read_success_marker(path) is None: + return None + try: + store.verify(handle) + return rebased(store, handle) + except (DiskgraphCorruptionError, ValueError) as error: + warnings.warn(f"Reloading the damaged graph index at {path}: {error}", + UserWarning, stacklevel=3) + return None + + +def _snapshot_id(graph_key: str, data_version: str) -> str: + return f'{graph_key}@{data_version[:16]}' + + +_COMPARABLE = ('exact', 'timestamp', 'listing') + + +def _source_change(table: Table, recorded: Optional[SourceTableEntry], + record_sources: bool) -> Dict[str, Any]: + from kumoapi.source_table import CustomSQLSourceTable + + before = None if recorded is None else recorded.version + current, same_source = None, False + comparable = (recorded is not None + and recorded.version_kind in _COMPARABLE) + try: + source = table.definition.source_table + if not isinstance(source, CustomSQLSourceTable): + connector = table.connector + stamp = connector.source_backend.get_data_version( + connector, table.table_name) + now = source_table(connector, table.table_name, None, stamp, + record_sources=record_sources) + current = now.version + comparable = comparable and now.version_kind in _COMPARABLE + same_source = recorded is not None and ( + recorded.connector, recorded.location, + recorded.location_fingerprint) == (now.connector, now.location, + now.location_fingerprint) + except Exception: + current = None + if not (same_source and comparable and before is not None + and current is not None): + status = 'unknown' + else: + status = 'unchanged' if before == current else 'changed' + return {'status': status, 'recorded': before, 'current': current} + + @dataclass(frozen=True) class GraphSnapshot: r"""The materialized artifact produced by `Graph.snapshot()`: an on-disk @@ -209,11 +368,164 @@ class GraphSnapshot: `TrainingWorkflowRequest.from_spec`/`materialize(...)`). Computed once here, at snapshot time, rather than being rebuilt on every `Trainer.fit`/`predict` call. + record_sources: the `record_sources` choice this snapshot was taken + with; `None` when unknown (built elsewhere or loaded). """ id: str index_path: Path graph_config: Dict[str, Any] + record_sources: bool | None = field(default=None, compare=False) + + @property + def _graph_key(self) -> str: + """The `Graph.id` this snapshot was built from.""" + return self.id.partition('@')[0] + + @property + def _data_version(self) -> str: + """The full fingerprint of the source versions the index was built + from; `id` carries its first 16 characters.""" + return Path(self.index_path).name + + @cached_property + def _native_graph(self) -> NativeGraph: + """`graph_config` in `kumo_core`'s normalized form, as + `Graph.native_graph`.""" + return normalize_graph(self.graph_config) + + def export(self, output_config: 'SnapshotOutputConfig') -> str: + r"""Copy the index to `//` and + return that location. An earlier export there stays loadable until + the new one commits, then is replaced. The connector must hold files + at the path: an S3, GCS or file connector, or a volume or stage + whose sink offers a filesystem. + """ + from kumo.engine._snapshot_io import export + + return export(self, output_config.connector, output_config.path) + + @classmethod + def load_from_connector(cls, connector: Connector, *, path: str, + graph: Optional['Graph'] = None) -> 'GraphSnapshot': + r"""Copy the snapshot exported at `path` into the local store. + + Args: + path: the location `export` returned. + graph: if given, warn when its source tables changed since the + snapshot was built (see `source_changes`). + """ + from kumo_core.diskgraph.snapshot import DiskgraphSnapshotHandle + + from kumo.engine._snapshot_io import holds_export, load_into, read + + root, manifest = read(connector, path) + graph_key = manifest['graph_key'] + if graph is not None and graph.id != graph_key: + report_other_graph( + f"Snapshot at {path} was taken of another graph.", + graph._graph_config, manifest['graph_config'], stacklevel=2) + store = _store() + handle = DiskgraphSnapshotHandle( + graph_key=graph_key, graph_data_version=manifest['data_version']) + with _lock(graph_key): + result = _loaded(store, handle) + if result is not None and not holds_export( + Path(result.index_path), manifest): + result = None + if result is None: + try: + load_into(connector, path, root, manifest, + Path(store.path(handle))) + except FileNotFoundError: + # A newer export of this snapshot committed meanwhile. + newer_root, newer = read(connector, path) + if (newer_root == root + or newer['snapshot_id'] != manifest['snapshot_id']): + raise + root, manifest = newer_root, newer + load_into(connector, path, root, manifest, + Path(store.path(handle))) + result = _loaded(store, handle) + if result is None: + raise RuntimeError(f'The snapshot loaded from {path} did not ' + 'pass its integrity check.') + snapshot = cls(id=manifest['snapshot_id'], + index_path=Path(result.index_path), + graph_config=manifest['graph_config']) + if graph is not None: + changed = [ + alias + for alias, change in snapshot.source_changes(graph).items() + if change['status'] == 'changed' + ] + if changed: + warnings.warn( + f"Snapshot at {path} predates source table(s) " + f"{', '.join(changed)}; see `snapshot.source_changes`.", + UserWarning, stacklevel=2) + return snapshot + + def source_changes(self, graph: 'Graph') -> Dict[str, Dict[str, Any]]: + r"""Each of `graph`'s tables' source status since this snapshot: + `unchanged`, `changed` or `unknown`. Never raises.""" + try: + manifest = self.manifest + recorded = ({} if manifest is None or manifest.reconstructed else + dict(manifest.parents.source_tables)) + record_sources = (manifest is None + or manifest.record_sources is not False) + except Exception: + recorded, record_sources = {}, True + from kumo_core.sources.parallel import run_in_parallel + + tables = graph._named_tables + changes = run_in_parallel( + [ + partial(_source_change, table, recorded.get(alias), + record_sources) + for alias, table in tables.items() + ], workers=max(1, min(len(tables), 16)), stop=threading.Event()) + return dict(zip(tables, changes)) + + @property + def manifest(self) -> SnapshotManifest | None: + """The snapshot manifest stored next to the index, if any.""" + return read_manifest(Path(self.index_path) / SNAPSHOT_MANIFEST_FILE, + SnapshotManifest) + + +def _write_snapshot_manifest(result: Any, snapshot_id: str, + source_tables: Dict[str, SourceTableEntry], + record_sources: bool, *, fresh: bool) -> None: + path = Path(result.index_path) / SNAPSHOT_MANIFEST_FILE + try: + if not fresh and path.exists(): + existing = read_manifest(path, SnapshotManifest) + if existing is None or (existing.record_sources + is not False) == record_sources: + return + parents = (SnapshotParents(source_tables=source_tables) + if existing.parents is None else replace( + existing.parents, source_tables=source_tables)) + manifest = replace(existing, record_sources=record_sources, + parents=parents) + else: + manifest = SnapshotManifest( + id=snapshot_id, + parents=SnapshotParents(source_tables=source_tables), + record_sources=record_sources, + **(stamp() if fresh else RECONSTRUCTED)) + write_manifest(path, manifest) + except Exception as error: + if fresh or path not in _warned_backfill: + if not fresh: + _warned_backfill.add(path) + warnings.warn( + f"Could not write the snapshot manifest at {path}: {error}", + UserWarning, + stacklevel=4, + ) @dataclass(frozen=True) @@ -345,7 +657,8 @@ def validate(self) -> None: """ self.native_graph - def snapshot(self) -> GraphSnapshot: + def snapshot(self, *, record_sources: bool = True, + mlops: Optional[Any] = None) -> GraphSnapshot: r"""Materialize this graph: read each table's connector, run the result through `kumo-core` to build column stats and the on-disk index. The index is stored under a directory keyed by both `self.id` @@ -354,9 +667,15 @@ def snapshot(self) -> GraphSnapshot: and `GraphSnapshot`'s docstring for the resulting `id` shape. If a valid build already exists for that exact key, ingest is skipped and it's reused as is. - """ - from filelock import FileLock + Args: + mlops: a `kumo.MLOps`; with MLOps on (this or a `kumo.mlops` + block), the snapshot is also exported under the artifact + root's `snapshots/`, unless it is there already. + record_sources: if False, the snapshot manifest records each + source as connector and table name only (no bucket, path or + account) and Custom SQL as a hash only (no SQL text). + """ self.validate() for table in self._named_tables.values(): table.connector.validate_scratch_dir(get_scratch_dir()) @@ -364,7 +683,6 @@ def snapshot(self) -> GraphSnapshot: graph_key = self.id persist_dir = snapshot_scratch_dir(graph_key) / 'source_tables' - persist_dir.parent.mkdir(parents=True, exist_ok=True) # Guards two things, both shared across every `snapshot()` call for # this `graph_key` -- different processes, or different threads in @@ -382,11 +700,20 @@ def snapshot(self) -> GraphSnapshot: # mid-build -- a plain in-memory lock wouldn't coordinate across # processes, and a lock-file-existence mutex would leave a stale # lock behind after a crash. - with FileLock(f'{persist_dir}.lock'): - return self._snapshot_no_lock(graph_key, persist_dir) + from kumo.engine._mlops import resolve - def _snapshot_no_lock(self, graph_key: str, persist_dir: Path) \ - -> GraphSnapshot: + settings = resolve(mlops) + with _lock(graph_key): + snapshot = self._snapshot_no_lock(graph_key, persist_dir, + record_sources) + if settings is not None and settings.wants('snapshot'): + from kumo.engine._snapshot_io import store + + store(snapshot, settings) + return snapshot + + def _snapshot_no_lock(self, graph_key: str, persist_dir: Path, + record_sources: bool = True) -> GraphSnapshot: r"""`snapshot()`'s body, run under its `FileLock`. Not meant to be called on its own -- `persist_dir`, unlike everything else here, isn't a pure function of `self` alone (it's `snapshot()`'s lock key, @@ -400,9 +727,7 @@ def _snapshot_no_lock(self, graph_key: str, persist_dir: Path) \ from kumo_core.diskgraph.runner import IngestExecutionConfig from kumo_core.diskgraph.snapshot import ( DiskgraphIngestResult, - DiskgraphLocalConfig, DiskgraphSnapshotHandle, - LocalSnapshotStore, check_diskgraph_output, ) from kumo_core.sources.layout import staging_dir @@ -412,12 +737,11 @@ def _snapshot_no_lock(self, graph_key: str, persist_dir: Path) \ from kumo.engine.connector.source_table import checked_sql graph_config = self._graph_config - store = LocalSnapshotStore( - DiskgraphLocalConfig(data_root=str(get_scratch_dir() / - 'snapshots'))) + store = _store() table_data_versions: Dict[str, str] = {} source_paths: Dict[str, str] = {} + source_tables: Dict[str, SourceTableEntry] = {} try: for alias, table in self._named_tables.items(): @@ -451,6 +775,9 @@ def _snapshot_no_lock(self, graph_key: str, persist_dir: Path) \ "limitation in kumoai/backend/local/graph.py's " "_stage_sources, which raises the same way.") table_data_versions[alias] = prepared.version + source_tables[alias] = source_table( + connector, name, sql, prepared.version, + record_sources=record_sources) paths = prepared.source.get('paths') or [] if paths: @@ -499,11 +826,16 @@ def _snapshot_no_lock(self, graph_key: str, persist_dir: Path) \ ), ) + snapshot_id = _snapshot_id(graph_key, data_version) + _write_snapshot_manifest(result, snapshot_id, source_tables, + record_sources, fresh=existing is None) + _warn_ingest_data_loss(result, graph_config) return GraphSnapshot( - id=f'{graph_key}@{data_version[:16]}', + id=snapshot_id, index_path=Path(result.index_path), graph_config=graph_config, + record_sources=record_sources, ) finally: discard_destination(str(persist_dir)) diff --git a/packages/kumo-engine/kumo/engine/pquery/_custom_source.py b/packages/kumo-engine/kumo/engine/pquery/_custom_source.py index a98c74975..74ea2e9d9 100644 --- a/packages/kumo-engine/kumo/engine/pquery/_custom_source.py +++ b/packages/kumo-engine/kumo/engine/pquery/_custom_source.py @@ -8,6 +8,7 @@ import os import shutil from pathlib import Path +from typing import Optional, Tuple import pyarrow.dataset as ds import pyarrow.parquet as pq @@ -15,8 +16,9 @@ from kumo.engine.connector.source_table import SourceTable -def stage_source_table(source_table: SourceTable, destination: Path) -> Path: - """Read ``source_table`` through its connector and retain its parquet parts.""" +def stage_source_table(source_table: SourceTable, + destination: Path) -> Tuple[Path, Optional[str]]: + """Stage ``source_table``'s parquet parts; returns them and their version.""" if not isinstance(source_table, SourceTable): raise TypeError("source_table must be a kumo.engine SourceTable") destination.mkdir(parents=True, exist_ok=False) @@ -62,7 +64,7 @@ def stage_source_table(source_table: SourceTable, destination: Path) -> Path: errno.EPERM, errno.EACCES, errno.EOPNOTSUPP, errno.ENOSYS): raise shutil.copyfile(part, target) - return output + return output, prepared.version def _validate_source(path, query, graph, *, weight_col, split_col, diff --git a/packages/kumo-engine/kumo/engine/pquery/prediction_table.py b/packages/kumo-engine/kumo/engine/pquery/prediction_table.py index d0f96bc01..f64912109 100644 --- a/packages/kumo-engine/kumo/engine/pquery/prediction_table.py +++ b/packages/kumo-engine/kumo/engine/pquery/prediction_table.py @@ -104,7 +104,8 @@ def from_source_table( pq._require_matching_snapshot(snapshot) run_id = new_run_id() run_dir = prediction_table_scratch_dir(run_id) - rows_path = stage_source_table(source_table, run_dir / "source-table") + rows_path, _ = stage_source_table(source_table, + run_dir / "source-table") request, _ = _build_prediction_table_request( snapshot, pq._parsed, diff --git a/packages/kumo-engine/kumo/engine/pquery/predictive_query.py b/packages/kumo-engine/kumo/engine/pquery/predictive_query.py index 20fb51a9e..dab929730 100644 --- a/packages/kumo-engine/kumo/engine/pquery/predictive_query.py +++ b/packages/kumo-engine/kumo/engine/pquery/predictive_query.py @@ -16,6 +16,9 @@ `to_portable(definition)` -- the schema-level conversion, not an ingested snapshot -- that this doesn't need a `GraphSnapshot` either. +`PredictiveQuery._from_snapshot(snapshot, query)` binds a snapshot instead and +parses against `snapshot.graph_config`. + Construction itself stays I/O-free, consistent with `Table`/`Graph`: nothing parses until `_parsed` is actually touched (by `validate()`, `get_task_type()`, or a caller directly). A failed parse is *not* cached -- `functools. @@ -25,8 +28,8 @@ `suggest_model_plan` uses resolved graph/query metadata. Methods that generate training or prediction tables also read the materialized index. These methods -accept an explicit `GraphSnapshot` produced by `self.graph.snapshot()` and -never trigger graph ingest themselves. +accept an explicit `GraphSnapshot` of this query's graph and never trigger +graph ingest themselves. Supported, matching `kumo.engine.pquery.predictive_query.PredictiveQuery` in the cloud SDK: `validate`, `get_task_type`, `suggest_model_plan`, @@ -57,6 +60,7 @@ ) from kumo_core.model_plan_conversion import public_model_plan_from_internal from kumo_core.pquery import ( + NativeGraph, ValidatedNativeQuery, parse_and_validate_query, query_is_temporal, @@ -90,22 +94,34 @@ PredictionTable, _compile_prediction_table, ) -from kumo.engine.pquery.training_table import TrainingTable +from kumo.engine.pquery.training_table import TrainingTable, _finish logger = logging.getLogger(__name__) @dataclass(frozen=True) class PredictiveQuery: - r"""An immutable PQL query bound to a `Graph`.""" + r"""An immutable PQL query bound to a `Graph` or a `GraphSnapshot`.""" - graph: Graph + graph: Graph | GraphSnapshot query: str + @classmethod + def _from_snapshot(cls, snapshot: GraphSnapshot, + query: str) -> 'PredictiveQuery': + r"""Bind `query` to `snapshot`'s graph without its `Graph` object. + + Snapshots passed to this query's methods must come from the same + graph (same `graph_key`). That only checks the snapshot belongs to + this query's graph; it says nothing about compatibility with a + trained model. + """ + return cls(graph=snapshot, query=query) + @cached_property def _parsed(self) -> ValidatedNativeQuery: - r"""Parses and validates `self.query` against `self.graph. - native_graph`, via `kumo_core.pquery.parse_and_validate_query` -- + r"""Parses and validates `self.query` against `self.graph.native_graph`, + via `kumo_core.pquery.parse_and_validate_query` -- the same parser `Trainer._predict` already calls directly, and what `kumo_core.training.workflow.TrainingWorkflowRequest.from_spec` uses internally for `fit`. No materialized `GraphSnapshot` needed @@ -125,23 +141,27 @@ def _parsed(self) -> ValidatedNativeQuery: `task_type`/`task_family`/`query_type` on the returned `ValidatedNativeQuery` are all plain `str`; `get_task_type` below wraps `task_type` into the shared, public `kumoapi.task.TaskType`. - - Reuses `self.graph.native_graph` (itself a `cached_property`) - rather than building its own `graph_config`/`NativeGraph` -- one - canonical conversion per `Graph`, shared by `Graph.validate()` and - every `PredictiveQuery` bound to it. """ - return parse_and_validate_query(self.query, self.graph.native_graph) + return parse_and_validate_query(self.query, self._native_graph) + + @property + def _native_graph(self) -> NativeGraph: + if isinstance(self.graph, GraphSnapshot): + return self.graph._native_graph + return self.graph.native_graph def _require_matching_snapshot(self, snapshot: GraphSnapshot) -> None: - # snapshot.id is 'graph_key@data_version_prefix' (see - # Graph.snapshot()); self.graph.id is the graph_key alone, so - # membership is a prefix check, not equality. - if not snapshot.id.startswith(f'{self.graph.id}@'): + if isinstance(self.graph, GraphSnapshot): + if snapshot.id == self.graph.id: + return + graph_key = self.graph._graph_key + else: + graph_key = self.graph.id + if snapshot._graph_key != graph_key: raise ValueError( f"snapshot {snapshot.id!r} was not materialized from " - f"self.graph (id {self.graph.id!r}) -- call " - "self.graph.snapshot() and pass its result.") + f"this query's graph (id {graph_key!r}) -- pass a " + "snapshot of that graph.") def validate(self, verbose: bool = True) -> ValidationResponse: r"""Return validation warnings and information for a valid query. @@ -173,7 +193,7 @@ def get_task_type(self) -> TaskType: def _model_plan_validation_facts(self, run_mode: RunMode) -> dict[str, Any]: parsed = self._parsed - graph = self.graph.native_graph + graph = self._native_graph entity_tables = [parsed.entity_table] if parsed.rhs_node_type is not None: entity_tables.append(parsed.rhs_node_type) @@ -195,7 +215,7 @@ def suggest_model_plan( train_table_spec: Optional[TrainingTableSpec] = None, ) -> ModelPlan: r"""Suggests a modeling plan given this query and a materialized - `snapshot` of `self.graph`. Mirrors the cloud SDK's + `snapshot` of this query's graph. Mirrors the cloud SDK's `PredictiveQuery.suggest_model_plan`. Uses resolved graph/query metadata without generating a training table @@ -219,9 +239,11 @@ def generate_training_table( self, snapshot: GraphSnapshot, plan: Optional[TrainingTableGenerationPlan] = None, + *, + mlops: Optional[Any] = None, ) -> TrainingTable: r"""Compiles this query into a `TrainingTable` against a materialized - `snapshot` of `self.graph`. Mirrors the cloud SDK's + `snapshot` of this query's graph. Mirrors the cloud SDK's `PredictiveQuery.generate_training_table`. `Trainer.fit` consumes the returned `TrainingTable` directly. Its @@ -229,10 +251,12 @@ def generate_training_table( settings are retained without keeping a graph-bound task plan. Args: - snapshot: a snapshot of `self.graph` to compile against. + snapshot: a snapshot of this query's graph to compile against. plan: if not given, uses `kumo_core`'s generation defaults, matching the local backend's existing `suggest_training_table_plan` default. + mlops: a `kumo.MLOps`; with MLOps on (this or a `kumo.mlops` + block), the table is also stored under the artifact root. """ self._require_matching_snapshot(snapshot) for name in ('train_table_path_override', 'task_path_override'): @@ -308,7 +332,7 @@ def generate_training_table( or {})), ) - return TrainingTable( + table = TrainingTable( query=self.query, train_table_path=Path(generated.train_table_path), manifest_path=Path(generated.manifest_path), @@ -316,6 +340,7 @@ def generate_training_table( required_stages=task_plan.task_dataset.required_stages, generation_plan=generation_plan, ) + return _finish(table, snapshot, task_plan.train_table, mlops) def generate_prediction_table( self, @@ -325,7 +350,7 @@ def generate_prediction_table( plan: Optional[TrainingTableGenerationPlan] = None, ) -> PredictionTable: r"""Compiles this query into a `PredictionTable` against a - materialized `snapshot` of `self.graph`. Mirrors the cloud SDK's + materialized `snapshot` of this query's graph. Mirrors the cloud SDK's `PredictiveQuery.generate_prediction_table`, and parallels `generate_training_table` above. @@ -337,7 +362,7 @@ def generate_prediction_table( docstring). Args: - snapshot: a snapshot of `self.graph` to compile against. + snapshot: a snapshot of this query's graph to compile against. prediction_time: the anchor timestamp for a temporal query; when omitted, uses the latest source-time maximum across the query's target aggregations in the snapshot. diff --git a/packages/kumo-engine/kumo/engine/pquery/training_table.py b/packages/kumo-engine/kumo/engine/pquery/training_table.py index ae58aac63..d8e88689e 100644 --- a/packages/kumo-engine/kumo/engine/pquery/training_table.py +++ b/packages/kumo-engine/kumo/engine/pquery/training_table.py @@ -2,21 +2,38 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -from dataclasses import dataclass, replace +import hashlib +import json +import warnings +from dataclasses import dataclass, field, replace from pathlib import Path -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Any, Dict, Optional +from kumo_core.artifacts import PROVENANCE_FILE from kumo_core.model_plan import TrainingTableGenerationPlan from kumoapi.model_plan import \ TrainingTableGenerationPlan as PublicGenerationPlan from kumoapi.train import TrainingTableSpec +from kumo.engine import _provenance as provenance +from kumo.engine._manifest import ( + SourceTableEntry, + TrainingTableManifest, + TrainingTableParents, + read_manifest, + stamp, + write_manifest, +) from kumo.engine.connector.source_table import SourceTable from kumo.engine.graph import GraphSnapshot if TYPE_CHECKING: from kumo.engine.trainer.output_config import TrainingTableOutputConfig +STAGING_MANIFEST_FILE = PROVENANCE_FILE +EXPORT_MANIFEST_FILE = PROVENANCE_FILE +_RANDOM_SPLITS = ('RandomSplit', 'RandomLinkSplit') + @dataclass(frozen=True) class TrainingTable: @@ -37,6 +54,22 @@ class TrainingTable: required_stages: tuple[str, ...] #: Effective generation settings carried forward into ``Trainer.fit``. generation_plan: TrainingTableGenerationPlan + #: The artifact id: the same snapshot, query, plan, split seed and + #: engine build give the same id. None for a hand-built table. + id: Optional[str] = field(default=None, compare=False) + + @property + def tt_id(self) -> str: + """The id DiskGraph assigned to this table at generation.""" + from kumo_core.training.export import read_spec + + return read_spec(str(self.staging_dir))['tt_id'] + + @property + def manifest(self) -> Optional[TrainingTableManifest]: + """Lineage recorded at generation; None for older tables.""" + return read_manifest(self.staging_dir / STAGING_MANIFEST_FILE, + TrainingTableManifest) def export(self, output_config: "TrainingTableOutputConfig") -> str: r"""Write the staged rows with a ``SPLIT`` column (``0`` train, ``1`` @@ -63,7 +96,57 @@ def export(self, output_config: "TrainingTableOutputConfig") -> str: with tempfile.TemporaryDirectory(dir=root) as output: files = export_train_table(str(self.staging_dir), output) return write_files(connector, output_config.table_name, files, - output_config.write_mode, export_id=export_id) + output_config.write_mode, export_id=export_id, + **self._export_sidecars(connector, + output_config)) + + def _export_sidecars(self, connector, output_config) -> Dict[str, Any]: + from kumo_core.sources import layout + + from kumo.engine.connector import is_file_destination + + manifest = self.staging_dir / STAGING_MANIFEST_FILE + name = output_config.table_name.rstrip('/') + try: + if (not is_file_destination(connector, output_config.table_name) + or layout.file_type(name) == layout.CSV + or not manifest.is_file()): + return {} + return {'sidecars': {EXPORT_MANIFEST_FILE: manifest.read_bytes()}} + except Exception as error: + warnings.warn( + f"Exporting without {EXPORT_MANIFEST_FILE}: {error}", + UserWarning, stacklevel=3) + return {} + + @classmethod + def load_from_connector(cls, connector: Any, *, path: str, + pq: "PredictiveQuery", # noqa: F821 + snapshot: GraphSnapshot) -> "TrainingTable": + r"""Materialize a table written by :meth:`export` (a directory with a + ``SPLIT`` column) against ``snapshot``, as :meth:`from_source_table` + with ``split_col='SPLIT'``.""" + return cls.from_source_table(connector[path], + TrainingTableSpec(split_col='SPLIT'), + pq=pq, snapshot=snapshot) + + def _store(self, settings: Any) -> None: + from pyarrow.fs import FileType + + from kumo.engine.connector import file_store + from kumo.engine.trainer.output_config import TrainingTableOutputConfig + + path = settings.root.location(f'training_tables/{self.id}/') + try: + filesystem, root = file_store(settings.root.connector, path) + if filesystem.get_file_info( + f'{root}/_SUCCESS').type is not FileType.File: + self.export( + TrainingTableOutputConfig(settings.root.connector, path)) + except Exception as error: + warnings.warn( + f"Could not store training table {self.id} under the " + f"artifact root: {error}", UserWarning, stacklevel=3) @classmethod def from_source_table( @@ -74,6 +157,7 @@ def from_source_table( pq: "PredictiveQuery", # noqa: F821 snapshot: GraphSnapshot, plan: Optional[PublicGenerationPlan] = None, + mlops: Optional[Any] = None, ) -> "TrainingTable": r"""Materialize caller-supplied labeled rows as a training table. @@ -86,7 +170,9 @@ def from_source_table( ``snapshot`` supplies the index used to validate and normalize the rows. The returned object retains only the materialized rows and generation metadata; :meth:`Trainer.fit` independently selects its - graph through its own required ``snapshot=`` argument. + graph through its own required ``snapshot=`` argument. With MLOps on + (``mlops=`` or a ``kumo.mlops`` block), the table is also stored + under the artifact root. """ # Local imports keep the immutable descriptor cheap to import and # avoid pquery -> trainer -> pquery cycles. @@ -126,7 +212,8 @@ def from_source_table( ) run_id = new_run_id() run_dir = train_table_scratch_dir(run_id) - rows_path = stage_source_table(source_table, run_dir / "source-table") + rows_path, version = stage_source_table(source_table, + run_dir / "source-table") generation_plan = from_json_dict( to_json_dict(public_plan), TrainingTableGenerationPlan, @@ -183,7 +270,7 @@ def from_source_table( year_over_year=("yoy_offsets_secs" in (task_plan.train_table.compiler_options or {})), ) - return cls( + table = cls( query=pq.query, train_table_path=Path(generated.train_table_path), manifest_path=Path(generated.manifest_path), @@ -191,3 +278,86 @@ def from_source_table( required_stages=task_plan.task_dataset.required_stages, generation_plan=generation_plan, ) + return _finish( + table, snapshot, task_plan.train_table, mlops, + source_table=provenance.source_table( + source_table.connector, source_table.table_name, + source_table.sql, version, + record_sources=_record_sources(snapshot)), + source_version=version) + + +def _record_sources(snapshot: GraphSnapshot) -> bool: + if snapshot.record_sources is not None: + return snapshot.record_sources + manifest = snapshot.manifest + return manifest is not None and manifest.record_sources is True + + +def _finish(table: TrainingTable, snapshot: GraphSnapshot, request: Any, + mlops: Optional[Any], *, + source_table: Optional[SourceTableEntry] = None, + source_version: Any = None) -> TrainingTable: + from kumo.engine._mlops import resolve + + seed = None + if getattr(request, 'split_family', None) in _RANDOM_SPLITS: + seed = (request.split_config or {}).get('rnd_seed') or 0 + table = replace(table, id=_table_id(table, snapshot, seed, source_version, + custom=source_table is not None)) + _record_manifest(table, snapshot, request, seed, source_table=source_table) + settings = resolve(mlops) + if settings is not None and settings.wants('training_table'): + table._store(settings) + return table + + +def _table_id(table: TrainingTable, snapshot: GraphSnapshot, + seed: Optional[int], source_version: Any, *, + custom: bool) -> str: + from kumo_core.io.json import to_json_dict + + from kumo.engine._provenance import environment + + if custom and source_version is None: + from kumo.engine._scratch import new_run_id + return f'tt-{new_run_id()}' + build = environment() + key = json.dumps([ + snapshot.id, table.query, + to_json_dict(table.generation_plan), seed, + None if source_version is None else str(source_version), + build.kumo_engine, build.build_commit, build.diskgraph, + build.diskgraph_build_commit + ], sort_keys=True, default=str) + return f'tt-{hashlib.sha256(key.encode()).hexdigest()[:32]}' + + +def _record_manifest(table: TrainingTable, snapshot: GraphSnapshot, + request: Any, seed: Optional[int], *, + source_table: Optional[SourceTableEntry] = None) -> None: + from kumo.engine._manifest import embedded, stub + + path = table.staging_dir / STAGING_MANIFEST_FILE + try: + from kumo_core.io.json import to_json_dict + + parent = snapshot.manifest + write_manifest( + path, + TrainingTableManifest( + id=table.id, + parents=TrainingTableParents( + snapshot=embedded(parent) if parent is not None else + stub('snapshot', snapshot.id), + source_table=source_table), + query=table.query, + generation_plan=to_json_dict(table.generation_plan), + task_type=request.task_type, + split_seed=seed, + **stamp(), + )) + except Exception as error: + warnings.warn( + f"Could not write the training-table manifest at {path}: " + f"{error}", UserWarning, stacklevel=3) diff --git a/packages/kumo-engine/kumo/engine/trainer/__init__.py b/packages/kumo-engine/kumo/engine/trainer/__init__.py index 3301763b5..0f4c1cfff 100644 --- a/packages/kumo-engine/kumo/engine/trainer/__init__.py +++ b/packages/kumo-engine/kumo/engine/trainer/__init__.py @@ -5,9 +5,14 @@ from kumo.engine.trainer.output_config import ( ModelOutputConfig, PredictionOutputConfig, + SnapshotOutputConfig, TrainingTableOutputConfig, ) -from kumo.engine.trainer.trainer import PredictionResult, Trainer, TrainingResult +from kumo.engine.trainer.trainer import ( + PredictionResult, + Trainer, + TrainingResult, +) __all__ = [ "Trainer", @@ -15,5 +20,6 @@ "PredictionResult", "ModelOutputConfig", "PredictionOutputConfig", + "SnapshotOutputConfig", "TrainingTableOutputConfig", ] diff --git a/packages/kumo-engine/kumo/engine/trainer/_provenance.py b/packages/kumo-engine/kumo/engine/trainer/_provenance.py new file mode 100644 index 000000000..67bb84d56 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -0,0 +1,86 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""A model bundle's ``_provenance.json``; excluded from the bundle id.""" + +from dataclasses import replace +from datetime import datetime, timezone +from pathlib import Path +from typing import Optional + +from kumo.engine._manifest import ( + FitTiming, + ModelManifest, + ModelParents, + embedded, + read_manifest, + stamp, + stub, + write_manifest, +) +from kumo.engine.pquery.training_table import TrainingTable +from kumo.engine.trainer._spec import PROVENANCE_ROLE, listed_file + + +def write_provenance( + bundle_dir: Path, + *, + bundle_id: str, + train_table: TrainingTable, + warm_start: Optional[ModelManifest], + started_at: datetime, + fit_mode: Optional[str], +) -> None: + from kumo_core.artifacts import PROVENANCE_FILE, add_bundle_file + + finished_at = datetime.now(timezone.utc) + table = train_table.manifest + write_manifest( + bundle_dir / PROVENANCE_FILE, + ModelManifest( + id=bundle_id, + parents=ModelParents( + training_table=embedded(table) if table is not None else + stub('training_table', train_table.id or train_table.tt_id), + warm_start=None + if warm_start is None else embedded_model(warm_start)), + fit=FitTiming( + started_at=started_at.isoformat(timespec='seconds'), + finished_at=finished_at.isoformat(timespec='seconds'), + seconds=round((finished_at - started_at).total_seconds(), 3), + ), + fit_mode=fit_mode, + **stamp(), + )) + add_bundle_file(bundle_dir, PROVENANCE_ROLE, PROVENANCE_FILE) + + +def embedded_model(manifest: ModelManifest) -> ModelManifest: + r"""``manifest`` as a child model embeds it: its own model parents cut + to stubs, so a chain of retrains does not grow each manifest.""" + copy = embedded(manifest) + parents = copy.parents + if parents is None: + return copy + + def cut(model: Optional[ModelManifest]) -> Optional[ModelManifest]: + return None if model is None else ModelManifest(id=model.id, + reconstructed=True) + + return replace(copy, + parents=replace(parents, warm_start=cut(parents.warm_start))) + + +def parent_model(bundle_dir: Path, bundle_id: Optional[str] + ) -> Optional[ModelManifest]: + r"""The manifest of the bundle at ``bundle_dir``, or a stub.""" + manifest = read_provenance(bundle_dir) + if manifest is not None: + return manifest + return None if bundle_id is None else ModelManifest(id=bundle_id, + reconstructed=True) + + +def read_provenance(bundle_dir: Path) -> Optional[ModelManifest]: + path = listed_file(bundle_dir, PROVENANCE_ROLE) + return None if path is None else read_manifest(path, ModelManifest) diff --git a/packages/kumo-engine/kumo/engine/trainer/_spec.py b/packages/kumo-engine/kumo/engine/trainer/_spec.py index 0115532b7..6c7d4a660 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_spec.py +++ b/packages/kumo-engine/kumo/engine/trainer/_spec.py @@ -9,14 +9,12 @@ separate and must be supplied through a compatible prediction snapshot. """ -import hashlib import json -import os -import tempfile from dataclasses import asdict, dataclass from pathlib import Path -from typing import Any, Dict +from typing import Any, Dict, Optional, Tuple +from kumo_core.io.filesystem import atomic_write_bytes from kumoapi.model_plan import ModelPlan, TrainingTableGenerationPlan #: Bundle metadata filename and schema version. @@ -24,10 +22,12 @@ SPEC_VERSION = 1 #: What `kumo_core`'s `VitalArtifactPublicationPolicy` names its own -#: manifest. `write_training_spec` amends this file's `artifacts` list so -#: `training-spec.json` is covered the same way every other bundle file is. +#: manifest; `training-spec.json` is listed in it like every other file. MANIFEST_FILE = 'publication-manifest.json' SPEC_ROLE = 'training_spec' +PROVENANCE_ROLE = 'provenance' +HOLDOUT_ROLE = 'holdout' +HOLDOUT_META_ROLE = 'holdout_meta' @dataclass(frozen=True) @@ -44,36 +44,7 @@ class TrainingSpec: model_plan: Dict[str, Any] snapshot_id: str training_table_generation: Dict[str, Any] - - -def _sha256(path: Path) -> str: - digest = hashlib.sha256() - digest.update(path.read_bytes()) - return digest.hexdigest() - - -def _atomic_write(path: Path, payload: bytes) -> None: - """Replace one metadata file without exposing partial JSON.""" - temporary = None - try: - with tempfile.NamedTemporaryFile(dir=path.parent, - prefix=f'.{path.name}.', - delete=False) as out: - temporary = Path(out.name) - out.write(payload) - out.flush() - os.fsync(out.fileno()) - os.replace(temporary, path) - temporary = None - if hasattr(os, 'O_DIRECTORY'): - descriptor = os.open(path.parent, os.O_RDONLY | os.O_DIRECTORY) - try: - os.fsync(descriptor) - finally: - os.close(descriptor) - finally: - if temporary is not None: - temporary.unlink(missing_ok=True) + graph_config: Optional[Dict[str, Any]] = None def write_training_spec( @@ -83,8 +54,10 @@ def write_training_spec( model_plan: ModelPlan, snapshot_id: str, training_table_generation: TrainingTableGenerationPlan, + graph_config: Optional[Dict[str, Any]] = None, ) -> None: r"""Write an integrity-covered training spec into a completed bundle.""" + from kumo_core.artifacts import add_bundle_file from kumo_core.io.json import to_json_dict spec = TrainingSpec( @@ -93,30 +66,15 @@ def write_training_spec( model_plan=to_json_dict(model_plan), snapshot_id=snapshot_id, training_table_generation=to_json_dict(training_table_generation), + graph_config=graph_config, ) spec_path = bundle_dir / SPEC_FILE - _atomic_write( + atomic_write_bytes( spec_path, (json.dumps(asdict(spec), indent=2, sort_keys=True) + '\n').encode('utf-8'), ) - - manifest_path = bundle_dir / MANIFEST_FILE - with open(manifest_path, encoding='utf-8') as handle: - manifest = json.load(handle) - manifest['artifacts'] = [ - entry for entry in manifest.get('artifacts') or [] - if entry.get('relative_path') != SPEC_FILE - ] + [{ - 'role': SPEC_ROLE, - 'relative_path': SPEC_FILE, - 'size_bytes': spec_path.stat().st_size, - 'sha256': _sha256(spec_path), - }] - _atomic_write( - manifest_path, - (json.dumps(manifest, indent=2, sort_keys=True) + '\n').encode('utf-8'), - ) + add_bundle_file(bundle_dir, SPEC_ROLE, SPEC_FILE) def read_training_spec(bundle_dir: Path) -> TrainingSpec: @@ -147,14 +105,27 @@ def read_training_spec(bundle_dir: Path) -> TrainingSpec: model_plan=data['model_plan'], snapshot_id=data['snapshot_id'], training_table_generation=data['training_table_generation'], + graph_config=data.get('graph_config'), ) -def read_bundle_artifact_paths(bundle_dir: Path) -> Dict[str, str]: - r"""Rebuild the verified ``role -> path`` artifact mapping.""" +def listed_file(bundle_dir: Path, role: str) -> Optional[Path]: + try: + manifest = json.loads( + (bundle_dir / MANIFEST_FILE).read_text(encoding='utf-8')) + entries = manifest.get('artifacts') or [] + except (OSError, ValueError, AttributeError): + return None + return next((bundle_dir / entry['relative_path'] + for entry in entries if entry.get('role') == role), None) + + +def read_bundle_artifact_paths( + bundle_dir: Path) -> Tuple[Dict[str, str], str]: + r"""Rebuild the verified ``role -> path`` mapping and the bundle id.""" from kumo_core.artifacts import verify_model_bundle - verify_model_bundle(bundle_dir) + bundle_id = verify_model_bundle(bundle_dir) manifest_path = bundle_dir / MANIFEST_FILE if not manifest_path.is_file(): raise ValueError( @@ -175,5 +146,6 @@ def read_bundle_artifact_paths(bundle_dir: Path) -> Dict[str, str]: return { entry['role']: str(bundle_dir / entry['relative_path']) for entry in artifacts - if entry['role'] != SPEC_ROLE - } + if entry['role'] not in (SPEC_ROLE, PROVENANCE_ROLE, HOLDOUT_ROLE, + HOLDOUT_META_ROLE) + }, bundle_id diff --git a/packages/kumo-engine/kumo/engine/trainer/output_config.py b/packages/kumo-engine/kumo/engine/trainer/output_config.py index d947fe775..c828ab2a2 100644 --- a/packages/kumo-engine/kumo/engine/trainer/output_config.py +++ b/packages/kumo-engine/kumo/engine/trainer/output_config.py @@ -90,6 +90,15 @@ def __post_init__(self) -> None: object.__setattr__(self, 'table_name', self.table_name + '/') +@dataclass(frozen=True) +class SnapshotOutputConfig: + """Export a graph snapshot under ``path`` via ``connector``, a file + destination (`GraphSnapshot.export`).""" + + connector: Connector + path: str + + @dataclass(frozen=True) class ModelOutputConfig: """Publish a model to exactly one ``path`` or registry ``location``.""" diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 3e87b954e..45be9d35f 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -8,9 +8,10 @@ import hashlib import json import logging +import shutil import time import warnings -from dataclasses import dataclass, replace +from dataclasses import dataclass, field, replace from datetime import datetime, timezone from pathlib import Path from typing import ( @@ -21,6 +22,7 @@ Mapping, Optional, TypeVar, + Union, cast, ) @@ -52,13 +54,18 @@ require_embedding_models, require_trained_embedding_models, ) +from kumo.engine._manifest import ModelManifest from kumo.engine._scratch import new_run_id, trainer_scratch_dir from kumo.engine.connector import model_sink from kumo.engine.connector.base import Connector -from kumo.engine.graph import GraphSnapshot +from kumo.engine.connector.source_table import SourceTable +from kumo.engine.graph import GraphSnapshot, report_other_graph from kumo.engine.pquery.prediction_table import PredictionTable from kumo.engine.pquery.training_table import TrainingTable from kumo.engine.trainer._spec import ( + HOLDOUT_META_ROLE, + HOLDOUT_ROLE, + listed_file, read_bundle_artifact_paths, read_training_spec, write_training_spec, @@ -79,11 +86,12 @@ class TrainingResult: when `fit()` exports the bundle. It is `None` without an export. `holdout_path` is the published model's holdout Parquet file written by - `fit()`, or `None` when it has none (no holdout task, a model refit on the - full data, or results of `Trainer.load`). It is not part of the published - model bundle. The file is capped at 1 GiB in memory; `holdout_truncated` is - true when rows were dropped to fit, and `None` when the size record could - not be read. + `fit()`, or `None` when it has none (no holdout task, or a model refit on + the full data). The bundle carries a copy (`holdout.parquet`), so a model + loaded with `Trainer.load` or `Trainer.load_from_connector` has it too; + bundles from before this copy was added have none. The file is capped at + 1 GiB in memory; `holdout_truncated` is true when rows were dropped to fit, + and `None` when the size record could not be read. """ metrics: Dict[str, float] @@ -92,6 +100,38 @@ class TrainingResult: published_model: Optional[PublishedModel] = None holdout_path: Optional[Path] = None holdout_truncated: Optional[bool] = False + bundle_id: Optional[str] = None + _trainer: Optional['Trainer'] = field(default=None, repr=False, + compare=False) + + @property + def id(self) -> Optional[str]: + """The model's id: its bundle id.""" + return self.bundle_id + + @property + def manifest(self) -> Optional[ModelManifest]: + """The bundle's ``_provenance.json``; None for older bundles.""" + from kumo.engine.trainer._provenance import read_provenance + + return read_provenance(self.bundle_dir) + + def predict(self, *args: Any, **kwargs: Any) -> 'PredictionResult': + """Score with this model, as :meth:`Trainer.predict`.""" + return self._model('predict').predict(*args, **kwargs) + + def _model(self, call: str) -> 'Trainer': + trainer = self._trainer + if trainer is None: + raise RuntimeError(f'This result has no model to {call} with; ' + 'use Trainer.load on its bundle_dir.') + current = trainer._training_result + if current is None or current.bundle_dir != self.bundle_dir: + trainer = Trainer._load_bundle(self.bundle_dir, + trainer._embedding_models, + accelerator=trainer._accelerator) + self._trainer = trainer + return trainer def holdout_df(self) -> pd.DataFrame: r"""Read the holdout predictions as a pandas DataFrame. @@ -187,6 +227,72 @@ def _find_holdout(log_dir: str) -> tuple[Optional[Path], Optional[bool]]: return path, truncated +HOLDOUT_FILE = 'holdout.parquet' +HOLDOUT_META_FILE = 'holdout-meta.json' + + +def _bundle_holdout(result: TrainingResult) -> TrainingResult: + r"""Copy the holdout and its size record into the bundle, so a model + loaded from a registry or path still has them.""" + from kumo_core.artifacts import add_bundle_file + + if result.holdout_path is None: + return result + try: + shutil.copyfile(result.holdout_path, result.bundle_dir / HOLDOUT_FILE) + add_bundle_file(result.bundle_dir, HOLDOUT_ROLE, HOLDOUT_FILE) + if result.holdout_truncated is not None: + (result.bundle_dir / HOLDOUT_META_FILE).write_text( + json.dumps({'truncated': result.holdout_truncated})) + add_bundle_file(result.bundle_dir, HOLDOUT_META_ROLE, + HOLDOUT_META_FILE) + except Exception as error: + warnings.warn( + f"Could not add the holdout to the model bundle: {error}; a " + f"loaded copy of this model has no holdout_df().", UserWarning, + stacklevel=4) + return result + + +def _listed_holdout(bundle_dir: Path) -> tuple[Optional[Path], Optional[bool]]: + path = listed_file(bundle_dir, HOLDOUT_ROLE) + if path is None or not path.is_file(): + return None, False + meta = listed_file(bundle_dir, HOLDOUT_META_ROLE) + try: + truncated = (None if meta is None else bool( + json.loads(meta.read_text())['truncated'])) + except (OSError, ValueError, KeyError, TypeError): + truncated = None + return path, truncated + + +def _store_model(result: TrainingResult, settings: Any) -> None: + r"""Copy the bundle to ``/models//`` unless a complete + copy is there; failures warn.""" + from kumo_core.object_store_sink import holds_bundle, upload_bundle + from pyarrow.fs import FileSelector + + from kumo.engine.connector import file_store + + if result.bundle_id is None: + return + try: + filesystem, root = file_store( + settings.root.connector, + settings.root.location(f'models/{result.bundle_id}')) + if holds_bundle(filesystem, root, str(result.bundle_dir)): + return + if filesystem.get_file_info( + FileSelector(root, allow_not_found=True)): + filesystem.delete_dir_contents(root, missing_dir_ok=True) + upload_bundle(filesystem, root, str(result.bundle_dir)) + except Exception as error: + warnings.warn( + f"Could not store model {result.bundle_id} under the artifact " + f"root: {error}", UserWarning, stacklevel=4) + + def _published_holdout(trainer: Any) -> tuple[Optional[Path], Optional[bool]]: from kumo_core.training.config import FitMode @@ -275,6 +381,7 @@ def __init__( self._id = new_run_id() self._training_result: Optional[TrainingResult] = None self._snapshot_id: Optional[str] = None + self._graph_config: Optional[Dict[str, Any]] = None self._training_query: Optional[str] = None self._training_table_generation_plan: Optional[ TrainingTableGenerationPlan] = None @@ -298,6 +405,11 @@ def training_table_generation_plan( def scratch_dir(self) -> Path: return trainer_scratch_dir(self._id) + @property + def result(self) -> Optional[TrainingResult]: + """The last fit's result, or the loaded bundle's.""" + return self._training_result + @_without_torch_scatter_hint def fit( self, @@ -306,6 +418,7 @@ def fit( snapshot: GraphSnapshot, export_config: Optional[ModelOutputConfig] = None, warm_start: Optional[Self] = None, + mlops: Optional[Any] = None, ) -> TrainingResult: r"""Train a model against `train_table`. @@ -334,6 +447,11 @@ def fit( may contain an AutoML search space instead of the fitted architecture; they support prediction but are unreliable warm-start sources. + mlops: a `kumo.MLOps`; with MLOps on (this or a `kumo.mlops` + block), the bundle is also stored under the artifact root. + + Each call returns a new `TrainingResult`; `Trainer.result` holds the + latest. """ from kumo_core.automl.runtime import ( inherit_warm_start_model_plan, @@ -363,18 +481,25 @@ def fit( raise NotImplementedError( f'{type(export_config.connector).__name__} cannot publish ' f'models') + from kumo.engine._mlops import resolve + from kumo.engine.trainer._provenance import parent_model + + settings = resolve(mlops) + started_at = datetime.now(timezone.utc) run_id = f'{self._id}-fit-{new_run_id()}' run_dir = self.scratch_dir / 'runs' / run_id model_plan = _to_core_model_plan(self._model_plan) warm_start_artifacts = None + warm_start_bundle_id = None effective_plan = self._model_plan if warm_start is not None: if warm_start._training_result is None: raise ValueError( "warm_start requires a fitted or loaded Trainer") - warm_start_artifacts = read_bundle_artifact_paths( - warm_start._training_result.bundle_dir) + warm_start_artifacts, warm_start_bundle_id = ( + read_bundle_artifact_paths( + warm_start._training_result.bundle_dir)) inherit_warm_start_model_plan( model_plan, load_json(warm_start_artifacts['resolved_model_plan'], @@ -508,6 +633,8 @@ def fit( holdout_path, holdout_truncated = _published_holdout( result.best_trainer) + fit_mode = result.best_trainer.published_fit_mode() + fit_mode = None if fit_mode is None else fit_mode.value training_result = TrainingResult( metrics=dict(result.metrics), artifact_paths=dict(result.artifact_paths), @@ -521,9 +648,17 @@ def fit( training_result, predictive_query=train_table.query, snapshot_id=snapshot.id, + graph_config=snapshot.graph_config, training_table_generation=train_table.generation_plan, export_config=export_config, model_plan=effective_plan, + provenance=dict( + train_table=train_table, + warm_start=None if warm_start is None else parent_model( + warm_start._training_result.bundle_dir, + warm_start_bundle_id), + started_at=started_at, fit_mode=fit_mode), + settings=settings, ) def _finish_fit( @@ -534,9 +669,16 @@ def _finish_fit( snapshot_id: str, training_table_generation: TrainingTableGenerationPlan, export_config: Optional[ModelOutputConfig], + graph_config: Optional[Dict[str, Any]] = None, model_plan: Optional[ModelPlan] = None, + provenance: Optional[Mapping[str, Any]] = None, + settings: Optional[Any] = None, ) -> TrainingResult: """Commit durable fit outputs before making this trainer usable.""" + from kumo_core.artifacts import bundle_id_from_manifest + + from kumo.engine.trainer._provenance import write_provenance + effective_plan = model_plan if model_plan is not None else self._model_plan write_training_spec( training_result.bundle_dir, @@ -544,7 +686,25 @@ def _finish_fit( model_plan=effective_plan, snapshot_id=snapshot_id, training_table_generation=training_table_generation, + graph_config=graph_config, ) + training_result = _bundle_holdout(training_result) + try: + training_result = replace( + training_result, + bundle_id=bundle_id_from_manifest(training_result.bundle_dir)) + except Exception as error: + warnings.warn(f"Could not compute the bundle id: {error}", + UserWarning, stacklevel=3) + if provenance is not None and training_result.bundle_id is not None: + try: + write_provenance(training_result.bundle_dir, + bundle_id=training_result.bundle_id, + **provenance) + except Exception as error: + warnings.warn( + f"Publishing the model bundle without _provenance.json: " + f"{error}", UserWarning, stacklevel=3) if export_config is not None: spec = read_training_spec(training_result.bundle_dir) sink = model_sink(export_config.connector, export_config.path) @@ -562,9 +722,13 @@ def _finish_fit( ) training_result = replace(training_result, published_model=published) + if settings is not None and settings.wants('model'): + _store_model(training_result, settings) + training_result = replace(training_result, _trainer=self) self._training_result = training_result self._snapshot_id = snapshot_id + self._graph_config = graph_config self._training_query = predictive_query self._training_table_generation_plan = training_table_generation @@ -573,20 +737,31 @@ def _finish_fit( @_without_torch_scatter_hint def predict( self, - prediction_table: PredictionTable, + source: Union[GraphSnapshot, PredictionTable, None] = None, + /, *, - snapshot: GraphSnapshot, + snapshot: Optional[GraphSnapshot] = None, + entities: Optional[SourceTable] = None, + prediction_time: Optional[datetime] = None, output_config: Optional[PredictionOutputConfig] = None, num_classes_to_return: Optional[int] = None, num_workers: int = 1, binary_classification_threshold: Optional[float] = None, + prediction_table: Optional[PredictionTable] = None, ) -> PredictionResult: - r"""Run batch prediction and return durable file-backed results. + r"""Score a snapshot and return durable file-backed results. + + ``predict(snapshot)`` scores the trained query's population at + ``prediction_time`` (as :meth:`PredictiveQuery.generate_prediction_table`); + ``entities`` is a source table of ``ENTITY1`` (and ``TIMESTAMP`` + where the task needs one) rows to score instead. It warns when the + snapshot was taken of another graph than the one the model was + trained on. ``predict(prediction_table, snapshot=...)`` scores a + prediction table generated beforehand, as before. The trainer must contain fitted state from :meth:`fit`, :meth:`load`, - or :meth:`load_from_connector`. ``snapshot`` selects graph data and - features independently of the table's retained population and anchor - times. Without ``output_config``, predictions remain on local disk. + or :meth:`load_from_connector`. Without ``output_config``, + predictions remain on local disk. ``num_workers`` controls scorer processes, not DataLoader workers. The default stays in-process; values above one spawn independent @@ -608,6 +783,65 @@ def predict( validate_prediction_workers(num_workers) validate_num_classes_to_return(num_classes_to_return) + if source is not None and not isinstance(source, GraphSnapshot): + if prediction_table is not None: + raise TypeError('predict() got two prediction tables') + source, prediction_table = None, source + if source is not None and snapshot is not None: + raise TypeError('predict() got two snapshots') + snapshot = snapshot if snapshot is not None else source + if snapshot is None: + raise TypeError('predict() needs a snapshot') + if prediction_table is not None and (entities is not None + or prediction_time is not None): + raise ValueError('A prediction table already holds its entities ' + 'and times; pass it or entities=/' + 'prediction_time=, not both.') + table = (prediction_table if prediction_table is not None else + self._prediction_table(snapshot, entities, prediction_time)) + return self._predict( + table, snapshot=snapshot, output_config=output_config, + num_classes_to_return=num_classes_to_return, + num_workers=num_workers, + binary_classification_threshold=binary_classification_threshold) + + def _prediction_table(self, snapshot: GraphSnapshot, + entities: Optional[SourceTable], + prediction_time: Optional[datetime] + ) -> PredictionTable: + from kumo.engine.pquery.predictive_query import PredictiveQuery + + if self._training_query is None: + raise RuntimeError( + "Trainer.predict() requires a prior Trainer.fit() or " + "Trainer.load() call on this same Trainer instance.") + if entities is not None and prediction_time is not None: + raise ValueError("Pass prediction_time or entities, not both; " + "put anchor times in the entities' TIMESTAMP.") + trained = (self._snapshot_id or '').partition('@')[0] + if trained and snapshot._graph_key != trained: + report_other_graph( + f"Snapshot {snapshot.id} was taken of graph " + f"{snapshot._graph_key}, but this model was trained on a " + f"snapshot of graph {trained}.", self._graph_config, + snapshot.graph_config, stacklevel=3) + query = PredictiveQuery._from_snapshot(snapshot, self._training_query) + plan = self._training_table_generation_plan + return (query.generate_prediction_table( + snapshot, prediction_time=prediction_time, plan=plan) + if entities is None else PredictionTable.from_source_table( + entities, pq=query, snapshot=snapshot, plan=plan)) + + def _predict( + self, + prediction_table: PredictionTable, + *, + snapshot: GraphSnapshot, + output_config: Optional[PredictionOutputConfig] = None, + num_classes_to_return: Optional[int] = None, + num_workers: int = 1, + binary_classification_threshold: Optional[float] = None, + ) -> PredictionResult: if self._training_result is None: raise RuntimeError( "Trainer.predict() requires a prior Trainer.fit() call on " @@ -831,7 +1065,7 @@ def _load_bundle( accelerator: Literal['cpu', 'gpu', 'mps', 'auto'] = 'auto', ) -> Self: bundle_dir = bundle_dir.resolve() - paths = read_bundle_artifact_paths(bundle_dir) + paths, bundle_id = read_bundle_artifact_paths(bundle_dir) spec = read_training_spec(bundle_dir) from kumo_core.io.json import from_json_dict plan = from_json_dict(spec.model_plan, ModelPlan) @@ -840,8 +1074,13 @@ def _load_bundle( encoding='utf-8')) if metrics_path.is_file() else {}) trainer = cls(plan, embedding_models, accelerator=accelerator) trainer._snapshot_id = spec.snapshot_id + trainer._graph_config = spec.graph_config trainer._training_query = spec.predictive_query trainer._training_table_generation_plan = from_json_dict( spec.training_table_generation, TrainingTableGenerationPlan) - trainer._training_result = TrainingResult(metrics, paths, bundle_dir) + holdout_path, holdout_truncated = _listed_holdout(bundle_dir) + trainer._training_result = TrainingResult( + metrics, paths, bundle_dir, holdout_path=holdout_path, + holdout_truncated=holdout_truncated, bundle_id=bundle_id, + _trainer=trainer) return trainer diff --git a/packages/kumo-engine/setup.py b/packages/kumo-engine/setup.py index f468be561..26d9a1f75 100644 --- a/packages/kumo-engine/setup.py +++ b/packages/kumo-engine/setup.py @@ -2,10 +2,13 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import os import re +import subprocess from pathlib import Path from setuptools import setup +from setuptools.command.build_py import build_py kumo_api_release_ver = '3.0.0' api_pyproject = Path(__file__).parent.parent / 'kumo-api' / 'pyproject.toml' @@ -22,7 +25,34 @@ kumo_api = f'kumo-api=={api_ver}' + +def _build_commit(): + commit = os.environ.get('KUMO_ENGINE_BUILD_COMMIT') + if commit: + return commit + command = ['git', '-C', str(Path(__file__).parent), 'rev-parse', + '--short=12', 'HEAD'] + try: + return subprocess.run(command, capture_output=True, text=True, + check=True, timeout=10).stdout.strip() or None + except (OSError, subprocess.SubprocessError): + return None + + +class _StampBuildCommit(build_py): + def run(self): + super().run() + target = Path(self.build_lib) / 'kumo' / 'engine' / '_version.py' + if target.is_file(): + text = target.read_text(encoding='utf-8') + target.unlink() + target.write_text( + f'{text}BUILD_COMMIT = {_build_commit()!r}\n', + encoding='utf-8') + + setup( + cmdclass={'build_py': _StampBuildCommit}, install_requires=[ "pandas", "pyarrow>=8.0.0", diff --git a/packages/kumo-engine/test/conftest.py b/packages/kumo-engine/test/conftest.py index e7601b972..4b60addea 100644 --- a/packages/kumo-engine/test/conftest.py +++ b/packages/kumo-engine/test/conftest.py @@ -59,6 +59,34 @@ def identity(filesystem, path): return probes +@pytest.fixture +def graph(tmp_path): + rng = np.random.default_rng(0) + data = tmp_path / 'data' + data.mkdir() + pd.DataFrame({'ID': np.arange(50)}).to_parquet(data / 'users.parquet') + pd.DataFrame({ + 'EID': + np.arange(500), + 'USER_ID': + rng.integers(0, 50, 500), + 'TS': + pd.Timestamp('2024-01-01') + + pd.to_timedelta(rng.integers(0, 120, 500), unit='D'), + }).to_parquet(data / 'events.parquet') + connector = kumo.FileConnector(root_dir=str(data)) + return kumo.Graph( + tables={ + 'users': + kumo.Table.from_source_table(connector['users.parquet'], + primary_key='ID'), + 'events': + kumo.Table.from_source_table( + connector['events.parquet'], primary_key='EID', + time_column='TS').with_column('USER_ID', stype=Stype.ID), + }, edges=[Edge(src_table='events', fkey='USER_ID', dst_table='users')]) + + @pytest.fixture def snapshot(tmp_path): config = { diff --git a/packages/kumo-engine/test/connector/test_custom_source.py b/packages/kumo-engine/test/connector/test_custom_source.py index aaa2de4e1..87080104a 100644 --- a/packages/kumo-engine/test/connector/test_custom_source.py +++ b/packages/kumo-engine/test/connector/test_custom_source.py @@ -40,11 +40,12 @@ def prepare_ingest(connector, table_name, destination, sql): "paths": [str(first), str(second)], }, credentials=None, + version=None, ) connector = SimpleNamespace(source_backend=Backend()) source = SourceTable("curated", connector) - output = stage_source_table(source, tmp_path / "staged") + output, _ = stage_source_table(source, tmp_path / "staged") assert output == tmp_path / "staged" / "parts" assert pq.read_table(output)["ENTITY1"].to_pylist() == [1, 2] @@ -62,7 +63,7 @@ def test_stage_source_table_accepts_partitioned_file_connector_directory( pq.write_table(pa.table({"ENTITY1": [2]}), parts / "two.parquet") connector = FileConnector(root_dir=str(root)) - output = stage_source_table(connector["parts"], tmp_path / "staged") + output, _ = stage_source_table(connector["parts"], tmp_path / "staged") assert output == tmp_path / "staged" / "parts" assert sorted(pq.read_table(output)["ENTITY1"].to_pylist()) == [1, 2] @@ -96,10 +97,10 @@ def refuse_link(*args, **kwargs): source={ "type": "parquet", "paths": [str(p) for p in paths] - }, credentials=None) + }, credentials=None, version=None) source = SourceTable("curated", Mock(source_backend=backend)) - output = stage_source_table(source, tmp_path / "staged") + output, _ = stage_source_table(source, tmp_path / "staged") retained = sorted(output.glob("*.parquet")) assert len(retained) == 2 @@ -124,7 +125,7 @@ def test_stage_source_table_casts_to_what_the_table_reader_reads(tmp_path): "TARGET1": pa.array(["a", "b"]).dictionary_encode(), }), root / "curated.parquet") - output = stage_source_table( + output, _ = stage_source_table( FileConnector(root_dir=str(root))["curated"], tmp_path / "staged") assert pq.read_table(output).schema == pa.schema([("ENTITY1", pa.int64()), @@ -137,7 +138,7 @@ def test_stage_source_table_retains_single_file_in_parts_directory(tmp_path): pq.write_table(pa.table({"ENTITY1": [1, 2]}), root / "table.parquet") connector = FileConnector(root_dir=str(root)) - output = stage_source_table(connector["table.parquet"], tmp_path / "staged") + output, _ = stage_source_table(connector["table.parquet"], tmp_path / "staged") assert output == tmp_path / "staged" / "parts" parts = list(output.glob("*.parquet")) @@ -171,8 +172,8 @@ def test_stage_source_table_retains_reencoded_float64_weight_parts( pq.write_table(table, root / "weighted.parquet") connector = FileConnector(root_dir=str(root)) - output = stage_source_table(connector["weighted.parquet"], - tmp_path / "staged") + output, _ = stage_source_table(connector["weighted.parquet"], + tmp_path / "staged") assert output == tmp_path / "staged" / "parts" assert len(list((tmp_path / "staged" / "source").rglob("*.parquet"))) > 1 diff --git a/packages/kumo-engine/test/test_bundle_provenance.py b/packages/kumo-engine/test/test_bundle_provenance.py new file mode 100644 index 000000000..dab92e4e7 --- /dev/null +++ b/packages/kumo-engine/test/test_bundle_provenance.py @@ -0,0 +1,164 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json +import re +import shutil +from datetime import datetime, timezone +from types import SimpleNamespace + +import pandas as pd +import pytest +from kumo_core.artifacts import add_bundle_file +from kumoapi.model_plan import ModelPlan, TrainingTableGenerationPlan + +import kumo.engine as kumo +from kumo.engine._manifest import ( + ModelManifest, + ModelParents, + read_manifest, + write_manifest, +) +from kumo.engine.trainer._provenance import embedded_model +from kumo.engine.trainer._spec import MANIFEST_FILE +from kumo.engine.trainer.trainer import TrainingResult + + +def test_fit_records_a_self_contained_manifest(fitted): + trainer, _, snapshot = fitted + result = trainer.result + manifest = result.manifest + assert manifest.kind == 'model' + assert re.fullmatch('[0-9a-f]{64}', trainer.result.id) + assert trainer.result.id == result.id == manifest.id + assert trainer.result.manifest == manifest + table = manifest.parents.training_table + assert table.kind == 'training_table' and table.id.startswith('tt-') + assert table.parents.snapshot == snapshot.manifest + assert manifest.parents.warm_start is None + assert manifest.fit_mode == 'refit_trainval' + assert (result.bundle_dir / '_provenance.json').is_file() + loaded = kumo.Trainer.load(result.bundle_dir) + assert loaded.result.id == manifest.id + assert loaded.result.manifest == manifest + + +def test_manifest_round_trips_unchanged(fitted, tmp_path): + trainer, _, _ = fitted + newer = json.loads( + (trainer.result.bundle_dir / '_provenance.json').read_text()) + newer['future'] = True + newer['fit']['gpu_seconds'] = 1.5 + source = tmp_path / 'newer.json' + source.write_text(json.dumps(newer, indent=2, sort_keys=True) + '\n') + + manifest = read_manifest(source, ModelManifest) + write_manifest(tmp_path / 'out.json', manifest) + + assert (tmp_path / 'out.json').read_bytes() == source.read_bytes() + assert manifest.extra == {'future': True} + assert manifest.fit.extra == {'gpu_seconds': 1.5} + + +def test_bundle_without_a_manifest_loads_and_predicts_identically( + fitted, tmp_path): + trainer, prediction_table, snapshot = fitted + result = trainer.result + old = tmp_path / 'old' + shutil.copytree(result.bundle_dir, old) + manifest = json.loads((old / MANIFEST_FILE).read_text()) + manifest['artifacts'] = [ + entry for entry in manifest['artifacts'] + if entry['role'] != 'provenance' + ] + (old / MANIFEST_FILE).write_text(json.dumps(manifest)) + (old / '_provenance.json').unlink() + restored = kumo.Trainer.load(old) + assert restored.result.id == result.bundle_id + assert restored.result.manifest is None + + def predict(model): + return model.predict( + prediction_table, snapshot=snapshot, + binary_classification_threshold=.5).frame.sort_values( + 'ENTITY').reset_index(drop=True) + + pd.testing.assert_frame_equal(predict(restored), predict(trainer)) + + +def test_each_fit_returns_a_new_result(fitted): + trainer, _, _ = fitted + first = trainer.result + second = trainer._finish_fit( + TrainingResult({}, dict(first.artifact_paths), first.bundle_dir), + predictive_query='PREDICT x', snapshot_id='snap', export_config=None, + training_table_generation=TrainingTableGenerationPlan()) + assert second is not first and trainer.result is second + + +def _finish_stub_fit(bundle, trainer, warm_start=None, settings=None): + (bundle / 'model').mkdir(parents=True) + (bundle / 'model' / 'config.json').write_text('{}') + (bundle / MANIFEST_FILE).write_text(json.dumps({'artifacts': []})) + add_bundle_file(bundle, 'trainer_config', 'model/config.json') + table = SimpleNamespace(tt_id='tt-raw', id='tt-1', manifest=None) + return trainer._finish_fit( + TrainingResult({}, {}, bundle), predictive_query='PREDICT x', + snapshot_id='snap', export_config=None, + training_table_generation=TrainingTableGenerationPlan(), + provenance=dict(train_table=table, warm_start=warm_start, + started_at=datetime.now(timezone.utc), fit_mode=None), + settings=settings) + + +def test_missing_parents_are_stubs_and_retrain_chains_stay_flat(tmp_path): + first = _finish_stub_fit(tmp_path / 'one', kumo.Trainer(ModelPlan())) + assert first.manifest.parents.training_table.to_dict() == { + 'kind': 'training_table', 'schema_version': 1, 'id': 'tt-1', + 'created_at': None, 'environment': None, 'reconstructed': True, + 'parents': None, 'query': None, 'generation_plan': None, + 'task_type': None, 'split_seed': None} + second = _finish_stub_fit(tmp_path / 'two', kumo.Trainer(ModelPlan()), + warm_start=first.manifest) + third = _finish_stub_fit(tmp_path / 'three', kumo.Trainer(ModelPlan()), + warm_start=second.manifest) + parent = third.manifest.parents.warm_start + assert parent == embedded_model(second.manifest) + assert parent.parents.warm_start == ModelManifest(id=first.id, + reconstructed=True) + assert parent.parents == ModelParents( + training_table=second.manifest.parents.training_table, + warm_start=ModelManifest(id=first.id, reconstructed=True)) + + +def test_a_mlops_block_stores_each_model_once(tmp_path): + cfg = kumo.MLOps(artifact_root=str(tmp_path / 'artifacts')) + result = _finish_stub_fit(tmp_path / 'bundle', kumo.Trainer(ModelPlan()), + settings=cfg) + stored = tmp_path / 'artifacts' / 'models' / result.id + assert kumo.Trainer.load(stored).result.id == result.id + (stored / 'model' / 'config.json').write_text('{"damaged": 1}') + again = _finish_stub_fit(tmp_path / 'again', kumo.Trainer(ModelPlan()), + settings=cfg) + assert again.id == result.id + assert (stored / 'model' / 'config.json').read_text() == '{}' + off = kumo.MLOps(artifact_root=str(tmp_path / 'off'), export={'snapshot'}) + _finish_stub_fit(tmp_path / 'off-bundle', kumo.Trainer(ModelPlan()), + settings=off) + assert not (tmp_path / 'off').exists() + + +def test_manifest_failure_warns_and_still_publishes(tmp_path, monkeypatch): + from kumo.engine.trainer import _provenance + + def fail(*args, **kwargs): + raise RuntimeError('unreadable manifest') + + monkeypatch.setattr(_provenance, 'write_manifest', fail) + trainer = kumo.Trainer(ModelPlan()) + with pytest.warns(UserWarning, match='unreadable manifest'): + result = _finish_stub_fit(tmp_path / 'bundle', trainer) + assert result.manifest is None + assert re.fullmatch('[0-9a-f]{64}', result.bundle_id) + assert trainer.result is result diff --git a/packages/kumo-engine/test/test_custom_sql.py b/packages/kumo-engine/test/test_custom_sql.py index b2f80c231..498b045ce 100644 --- a/packages/kumo-engine/test/test_custom_sql.py +++ b/packages/kumo-engine/test/test_custom_sql.py @@ -131,7 +131,7 @@ def spy(graph_config, spec): def test_custom_training_rows_stage_the_query_result(warehouse, tmp_path): - staged = stage_source_table(warehouse.sql(SQL), tmp_path / 'staged') + staged, _ = stage_source_table(warehouse.sql(SQL), tmp_path / 'staged') assert pq.read_table(staged).to_pydict() == { 'user_id': [1, 3], 'age': [30, 50] diff --git a/packages/kumo-engine/test/test_filestore_sink.py b/packages/kumo-engine/test/test_filestore_sink.py index 8716310f8..770b1ac2d 100644 --- a/packages/kumo-engine/test/test_filestore_sink.py +++ b/packages/kumo-engine/test/test_filestore_sink.py @@ -191,3 +191,28 @@ def test_parquet_writes_stay_plain_up_to_one_part(tmp_path, connector): assert files.write_files(connector, '@stage/out/', [str(part)]) assert pd.read_parquet('@stage/out/part-0.parquet', filesystem=PyFileSystem(store)).equals(frame) + + +def test_publishing_an_identical_bundle_again_reuses_the_copy( + tmp_path, connector, monkeypatch, caplog): + import kumo_core.object_store_sink as sink + + bundle = tmp_path / 'bundle' + bundle.mkdir() + (bundle / '_provenance.json').write_text('{"id": 1}') + (bundle / 'weights.bin').write_bytes(os.urandom(10000)) + files = FilesystemSink(connector.sink_backend) + files.publish_model(connector, str(bundle), path='@stage/models/m') + upload = sink.upload_bundle + monkeypatch.setattr(sink, 'upload_bundle', + lambda *a: pytest.fail('uploaded again')) + with caplog.at_level('INFO', logger=sink.__name__): + files.publish_model(connector, str(bundle), path='@stage/models/m') + assert 'reused it' in caplog.text + + (bundle / '_provenance.json').write_text('{"id": 2}') + uploads = [] + monkeypatch.setattr(sink, 'upload_bundle', + lambda *a: uploads.append(a) or upload(*a)) + files.publish_model(connector, str(bundle), path='@stage/models/m') + assert len(uploads) == 1 diff --git a/packages/kumo-engine/test/test_graph_mismatch.py b/packages/kumo-engine/test/test_graph_mismatch.py new file mode 100644 index 000000000..6714c5ba9 --- /dev/null +++ b/packages/kumo-engine/test/test_graph_mismatch.py @@ -0,0 +1,89 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import copy +import logging +import warnings +from dataclasses import replace + +import pytest + +import kumo.engine as kumo +from kumo.engine.graph import schema_changes +from kumo.engine.pquery.predictive_query import PredictiveQuery + + +class _Stop(Exception): + pass + + +def _other(snapshot, edit): + config = copy.deepcopy(snapshot.graph_config) + edit({t['table_name']: t for t in config['tables']}) + return replace(snapshot, id='f' * 64 + '@' + snapshot.id.partition('@')[2], + graph_config=config) + + +def _moved(tables): + for table in tables.values(): + table['connector_id'] = 'prod' + table['source_table_name'] = 'prod.' + table['source_table_name'] + + +def _predict(trainer, snapshot, monkeypatch): + def stop(*args, **kwargs): + raise _Stop + + monkeypatch.setattr(PredictiveQuery, '_from_snapshot', stop) + with pytest.raises(_Stop): + trainer.predict(snapshot, binary_classification_threshold=.5) + + +def test_predict_says_what_differs_from_the_training_graph( + fitted, monkeypatch, caplog, tmp_path): + trainer, _, snapshot = fitted + assert kumo.Trainer.load(trainer.result.bundle_dir)._graph_config == ( + snapshot.graph_config) + + with caplog.at_level(logging.INFO, logger='kumo.engine.graph'), \ + warnings.catch_warnings(): + warnings.filterwarnings('error', message='Snapshot .* was taken') + _predict(trainer, _other(snapshot, _moved), monkeypatch) + assert 'only table connectors or locations differ' in caplog.text + + def dropped(tables): + _moved(tables) + tables['events']['fkeys'] = [] + + with pytest.warns(UserWarning, match=r'trained on a snapshot of graph ' + r'\w+\. Schema changes: events edges\.$'): + _predict(trainer, _other(snapshot, dropped), monkeypatch) + + monkeypatch.setattr(trainer, '_graph_config', None) + with pytest.warns(UserWarning, match='does not record its training graph'): + _predict(trainer, _other(snapshot, _moved), monkeypatch) + + +def test_schema_changes_ignore_connectors_and_locations(): + def config(location, columns, tables=('users', )): + return { + 'name': location, + 'tables': [{ + 'table_name': name, + 'source_table_name': f'{location}.{name}', + 'connector_id': location, + 'pkey': {'name': 'id', 'dtype': 'int64', 'stype': 'ID'}, + 'cols': columns, + 'fkeys': [], + } for name in tables], + } + + age = {'name': 'age', 'dtype': 'int64', 'stype': 'numerical'} + label = {'name': 'label', 'dtype': 'bool', 'stype': 'categorical'} + assert schema_changes(config('dev', [age, label]), + config('prod', [label, age])) == [] + assert schema_changes( + config('dev', [age]), + config('dev', [{**age, 'dtype': 'float64'}], tables=('users', 'x')) + ) == ['table x added', 'users columns'] diff --git a/packages/kumo-engine/test/test_holdout_df.py b/packages/kumo-engine/test/test_holdout_df.py index 70bfe82c6..c70110813 100644 --- a/packages/kumo-engine/test/test_holdout_df.py +++ b/packages/kumo-engine/test/test_holdout_df.py @@ -123,14 +123,32 @@ def test_no_holdout(graph): result.holdout_df() -def test_loaded_trainer_has_no_holdout(graph): - graph, snapshot = graph - result = _fit((graph, snapshot), 'PREDICT COUNT(events.*, 0, 7, days) > 0 ' - 'FOR EACH ents.ID') - - loaded = kumo.Trainer.load(result.bundle_dir) +def test_a_loaded_or_fetched_model_keeps_its_holdout(graph, tmp_path, + monkeypatch): + monkeypatch.setattr('kumo_core.train.output._HOLDOUT_MAX_BYTES', 1000) + result = _fit(graph, QUERY) + expected = result.holdout_df() + registry = kumo.FileConnector(root_dir=str(tmp_path)) + registry.sink_backend.publish_model(registry, str(result.bundle_dir), + path='models/m') + + for loaded in (kumo.Trainer.load(result.bundle_dir).result, + kumo.Trainer.load_from_connector(registry, + path='models/m').result): + assert loaded.bundle_id == result.bundle_id + assert loaded.holdout_truncated is True + pd.testing.assert_frame_equal(loaded.holdout_df(), expected) + + +def test_a_bundle_without_a_holdout_says_so(graph, monkeypatch): + monkeypatch.setattr('kumo.engine.trainer.trainer._bundle_holdout', + lambda result: result) # as bundles before the copy + result = _fit(graph, QUERY) - assert loaded._training_result.holdout_path is None + loaded = kumo.Trainer.load(result.bundle_dir).result + assert loaded.holdout_path is None + with pytest.raises(FileNotFoundError, match='No holdout'): + loaded.holdout_df() def test_truncation_is_recorded(tmp_path): diff --git a/packages/kumo-engine/test/test_mlops.py b/packages/kumo-engine/test/test_mlops.py new file mode 100644 index 000000000..6384faa72 --- /dev/null +++ b/packages/kumo-engine/test/test_mlops.py @@ -0,0 +1,86 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import asyncio +import threading +from types import SimpleNamespace + +import pytest + +import kumo.engine as kumo +from kumo.engine._manifest import ( + SnapshotManifest, + SnapshotParents, + embedded, + stub, +) +from kumo.engine._mlops import resolve + + +def test_mlops_needs_an_artifact_root_it_can_write_files_to(tmp_path): + with pytest.raises(TypeError): + kumo.MLOps() + with pytest.raises(ValueError, match='artifact_root'): + kumo.MLOps(artifact_root='') + table_only = SimpleNamespace(sink_backend=SimpleNamespace( + holds_files=lambda c, p: False, filesystem=None)) + with pytest.raises(ValueError, match='file location'): + kumo.MLOps(artifact_root=(table_only, 'DB.SCHEMA')) + with pytest.raises(ValueError, match='Unknown export kinds'): + kumo.MLOps(artifact_root=str(tmp_path), export={'models'}) + + +def test_a_root_is_a_uri_or_a_connector_and_path(tmp_path): + uri = kumo.MLOps(artifact_root=str(tmp_path)) + assert uri.root.location('models/m') == 'models/m' + assert uri.root.uri('models/m') == f'{tmp_path}/models/m' + assert isinstance(uri.root.connector, kumo.FileConnector) + assert uri.exports == {'snapshot', 'training_table', 'model'} + + connector = kumo.FileConnector(root_dir=str(tmp_path)) + pair = kumo.MLOps(artifact_root=(connector, 'volume/kumo/'), + export={'model'}) + assert pair.root.connector is connector + assert pair.root.location('models/m') == 'volume/kumo/models/m' + assert pair.wants('model') and not pair.wants('snapshot') + assert not kumo.MLOps(artifact_root=str(tmp_path), export=False).exports + + +def test_a_call_argument_beats_the_nearest_block(tmp_path): + outer = kumo.MLOps(artifact_root=str(tmp_path / 'outer')) + inner = kumo.MLOps(artifact_root=str(tmp_path / 'inner')) + call = kumo.MLOps(artifact_root=str(tmp_path / 'call')) + assert resolve() is None + seen = [] + with kumo.mlops(outer): + with kumo.mlops(inner): + assert resolve() is inner + assert resolve(call) is call + assert resolve() is outer + thread = threading.Thread(target=lambda: seen.append(resolve())) + thread.start() + thread.join() + + async def task(): + return resolve() + + assert asyncio.run(task()) is outer + assert seen == [None] and resolve() is None + + +def test_stubs_stand_in_for_parents_without_a_manifest(): + snapshot = stub('snapshot', 'abc@0123') + assert isinstance(snapshot, SnapshotManifest) + assert snapshot.to_dict() == { + 'kind': 'snapshot', + 'schema_version': 1, + 'id': 'abc@0123', + 'created_at': None, + 'environment': None, + 'reconstructed': True, + 'parents': None, + } + full = SnapshotManifest(id='abc@0123', + parents=SnapshotParents(source_tables={})) + assert embedded(full) == full diff --git a/packages/kumo-engine/test/test_predictive_query.py b/packages/kumo-engine/test/test_predictive_query.py index 7f3794266..2f77622a3 100644 --- a/packages/kumo-engine/test/test_predictive_query.py +++ b/packages/kumo-engine/test/test_predictive_query.py @@ -465,7 +465,7 @@ def test_custom_training_table_materializes_weight_and_manual_split( ) monkeypatch.setattr('kumo.engine.pquery._custom_source.stage_source_table', - lambda source_table, destination: rows) + lambda source_table, destination: (rows, None)) monkeypatch.setattr( 'kumo.engine.pquery._custom_source.validate_training_source', lambda *args: None) @@ -512,7 +512,7 @@ def test_custom_training_table_uses_the_normal_inferred_split( staging = tmp_path / 'staging' captured = {} monkeypatch.setattr('kumo.engine.pquery._custom_source.stage_source_table', - lambda source_table, destination: rows) + lambda source_table, destination: (rows, None)) monkeypatch.setattr( 'kumo.engine.pquery._custom_source.validate_training_source', lambda *args: None) @@ -637,7 +637,7 @@ def test_custom_prediction_table_materializes_exact_source_rows( ) captured = {} monkeypatch.setattr('kumo.engine.pquery._custom_source.stage_source_table', - lambda source_table, destination: rows) + lambda source_table, destination: (rows, None)) monkeypatch.setattr( 'kumo.engine.pquery._custom_source.validate_prediction_source', lambda *args: None, diff --git a/packages/kumo-engine/test/test_query_from_snapshot.py b/packages/kumo-engine/test/test_query_from_snapshot.py new file mode 100644 index 000000000..c22709da8 --- /dev/null +++ b/packages/kumo-engine/test/test_query_from_snapshot.py @@ -0,0 +1,123 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from dataclasses import replace + +import pandas as pd +import pytest + +import kumo.engine as kumo +from kumo.engine._scratch import get_scratch_dir + + +def _scores(trainer, table, snapshot): + frame = trainer.predict(table, snapshot=snapshot, + binary_classification_threshold=0.5).frame + return frame.sort_values('ENTITY').reset_index(drop=True) + + +def test_predict_scores_a_snapshot_with_the_saved_query(fitted, tmp_path): + trainer, graph_table, snapshot = fitted + expected = _scores(trainer, graph_table, snapshot) + + def scores(**options): + return trainer.predict( + snapshot, binary_classification_threshold=0.5, + **options).frame.sort_values('ENTITY').reset_index(drop=True) + + pd.testing.assert_frame_equal(scores(), expected) + pd.DataFrame({ + 'ENTITY1': range(50), + 'TIMESTAMP': [pd.Timestamp(graph_table.timestamps[0])] * 50, + }).to_parquet(tmp_path / 'rows.parquet') + rows = kumo.FileConnector(root_dir=str(tmp_path))['rows.parquet'] + pd.testing.assert_frame_equal(scores(entities=rows), expected) + pd.testing.assert_frame_equal(scores(prediction_table=graph_table), expected) + legacy = trainer.predict(prediction_table=graph_table, snapshot=snapshot, + binary_classification_threshold=0.5).frame + pd.testing.assert_frame_equal( + legacy.sort_values('ENTITY').reset_index(drop=True), expected) + with pytest.raises(ValueError, match='not both'): + scores(entities=rows, prediction_time=graph_table.timestamps[0]) + # Same schema under another graph id: scored, with an info log only. + other = replace(snapshot, id='other@' + snapshot.id.split('@')[1]) + scored = trainer.predict(other, binary_classification_threshold=0.5).frame + pd.testing.assert_frame_equal( + scored.sort_values('ENTITY').reset_index(drop=True), expected) + + +def test_exported_snapshot_scores_in_a_fresh_scratch_without_a_graph( + fitted, tmp_path): + trainer, graph_table, snapshot = fitted + expected = _scores(trainer, graph_table, snapshot) + exports = kumo.FileConnector(root_dir=str(tmp_path / 'exports')) + location = snapshot.export(kumo.SnapshotOutputConfig(exports, 'snapshots')) + bundle = trainer._training_result.bundle_dir + previous = get_scratch_dir() + kumo.configure(scratch_dir=tmp_path / 'fresh') + try: + loaded = kumo.GraphSnapshot.load_from_connector(exports, path=location) + assert loaded.id == snapshot.id + assert loaded.index_path.is_relative_to(tmp_path / 'fresh') + model = kumo.Trainer.load(bundle) + scored = model.result.predict( + loaded, binary_classification_threshold=0.5).frame + pd.testing.assert_frame_equal( + scored.sort_values('ENTITY').reset_index(drop=True), expected, + atol=1e-5) + finally: + kumo.configure(scratch_dir=previous) + + +@pytest.mark.parametrize('scope', [{ + 'entities': object() +}, { + 'prediction_time': '2024-01-01' +}]) +def test_a_prediction_table_takes_no_entities_or_time(scope): + from kumoapi.model_plan import ModelPlan + + with pytest.raises(ValueError, match='not both'): + kumo.Trainer(ModelPlan()).predict(snapshot=object(), + prediction_table=object(), **scope) + + +@pytest.mark.parametrize('call', [ + lambda trainer, snapshot: trainer.predict(object(), snapshot=snapshot, + prediction_table=object()), + lambda trainer, snapshot: trainer.predict(snapshot, snapshot=snapshot), +]) +def test_an_argument_given_twice_is_refused(fitted, call): + trainer, _, snapshot = fitted + with pytest.raises(TypeError, match='two'): + call(trainer, snapshot) + + +def test_each_fit_result_scores_with_its_own_model(fitted): + from kumoapi.model_plan import ModelPlan + + from kumo.engine.pquery.predictive_query import PredictiveQuery + + fixture_trainer, prediction_table, snapshot = fitted + table = PredictiveQuery._from_snapshot( + snapshot, + fixture_trainer._training_query).generate_training_table(snapshot) + plan = ModelPlan() + plan.training_job.disable_explain = True + plan.optimization.max_epochs = 1 + plan.optimization.max_steps_per_epoch = 2 + trainer = kumo.Trainer(model_plan=plan) + first = trainer.fit(table, snapshot=snapshot) + second = trainer.fit(table, snapshot=snapshot) + assert first.bundle_id != second.bundle_id + + scores = [ + _scores(result, prediction_table, snapshot) + for result in (first, second) + ] + assert not scores[0].equals(scores[1]) + for result, score in zip((first, second), scores): + own = kumo.Trainer.load(result.bundle_dir) + pd.testing.assert_frame_equal( + score, _scores(own, prediction_table, snapshot)) diff --git a/packages/kumo-engine/test/test_semantic_carry_forward.py b/packages/kumo-engine/test/test_semantic_carry_forward.py index a1533b131..9481d14d1 100644 --- a/packages/kumo-engine/test/test_semantic_carry_forward.py +++ b/packages/kumo-engine/test/test_semantic_carry_forward.py @@ -874,10 +874,9 @@ def fake_train(training_input, **kwargs): both_workers.wait(timeout=5) time.sleep(.05) run_name = Path(training_input.home_dir).name - return SimpleNamespace( - metrics={'run_name': - run_name}, artifact_paths={}, explain_disabled=True, - best_trainer=_best_trainer(tmp_path)) + return SimpleNamespace(metrics={'run_name': run_name}, + artifact_paths={}, explain_disabled=True, + best_trainer=_best_trainer(tmp_path)) monkeypatch.setattr( 'kumo.engine.trainer.trainer.build_generated_training_task_dataset', diff --git a/packages/kumo-engine/test/test_snapshot_export.py b/packages/kumo-engine/test/test_snapshot_export.py new file mode 100644 index 000000000..e0f15057f --- /dev/null +++ b/packages/kumo-engine/test/test_snapshot_export.py @@ -0,0 +1,786 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import hashlib +import json +import os +import re +import shutil +import threading +import warnings +import zlib +from contextlib import contextmanager +from dataclasses import replace +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import Mock + +import pandas as pd +import pytest +from pyarrow.fs import LocalFileSystem, SubTreeFileSystem + +import kumo.engine as kumo +import kumo.engine._snapshot_io as snapshot_io +from kumo.engine import GraphSnapshot, SnapshotOutputConfig +from kumo.engine._scratch import get_scratch_dir + + +@pytest.fixture +def scratch(tmp_path): + previous = get_scratch_dir() + + def use(name): + kumo.configure(scratch_dir=tmp_path / name) + + use('cpu') + yield use + kumo.configure(scratch_dir=previous) + + +@pytest.fixture +def exports(tmp_path): + return kumo.FileConnector(root_dir=str(tmp_path / 'exports')) + + +def _index_files(path): + return { + str(file.relative_to(path)): file.read_bytes() + for file in Path(path).rglob('*') if file.is_file() + and file.name not in ('_SUCCESS', '_provenance.json', '.kumo_export') + and not str(file.relative_to(path)).startswith('train_tables/') + } + + +def _export(snapshot, exports, tmp_path): + location = snapshot.export(SnapshotOutputConfig(exports, 'one')) + return location, tmp_path / 'exports' / location + + +def _generation(exported): + marker = json.loads((exported / 'kumo_snapshot.json').read_text()) + return exported / marker['generation'] + + +def _edit_manifest(exported, edit): + path = exported / 'kumo_snapshot.json' + manifest = json.loads(path.read_text()) + edit(manifest) + path.write_text(json.dumps(manifest)) + + +def _edit_index_manifest(exported, edit): + path = _generation(exported) / 'index/_MANIFEST.json' + manifest = json.loads(path.read_text()) + edit(manifest) + path.write_text(json.dumps(manifest)) + _edit_manifest( + exported, lambda m: m.update( + index_manifest=hashlib.sha256(path.read_bytes()).hexdigest())) + + +def _load(exports, location, **kwargs): + return GraphSnapshot.load_from_connector(exports, path=location, **kwargs) + + +def test_export_load_round_trip_relocates_and_is_reused( + tmp_path, scratch, graph, exports, monkeypatch): + monkeypatch.setattr(snapshot_io, '_CHUNK', 1024) + monkeypatch.setattr(snapshot_io, '_RANGE_BYTES', 4096) + snapshot = graph.snapshot() + location = snapshot.export(SnapshotOutputConfig(exports, 'one/')) + graph_id, data_version = snapshot.id.split('@') + assert location == f'one/{graph_id}/{data_version}' + assert snapshot.export(SnapshotOutputConfig(exports, 'one')) == location + exported = tmp_path / 'exports' / location + recorded = (snapshot.index_path / '_provenance.json').read_bytes() + assert (_generation(exported) / '_provenance.json').read_bytes() == recorded + + shutil.rmtree(tmp_path / 'cpu') + scratch('gpu') + loaded = _load(exports, location, graph=graph) + assert (loaded.id, loaded.graph_config) == (snapshot.id, + snapshot.graph_config) + assert loaded.index_path.is_relative_to(tmp_path / 'gpu') + assert _index_files(loaded.index_path) == _index_files( + _generation(exported) / 'index') + assert (loaded.index_path / + '_provenance.json').read_bytes() == recorded + reexported = tmp_path / 'exports' / loaded.export( + SnapshotOutputConfig(exports, 'two')) + assert (_generation(reexported) / + '_provenance.json').read_bytes() == recorded + + monkeypatch.setattr( + 'kumo_core.diskgraph.ingest.execute_diskgraph_ingest', + Mock(side_effect=AssertionError('ingest must not run'))) + monkeypatch.setattr(snapshot_io, 'load_into', + Mock(side_effect=AssertionError('already loaded'))) + assert _load(exports, location) == loaded + assert graph.snapshot() == loaded + + +def test_damaged_index_is_reused_by_snapshot_but_reloaded_by_load( + tmp_path, scratch, graph, exports, monkeypatch): + snapshot = graph.snapshot() + location, _ = _export(snapshot, exports, tmp_path) + expected = _index_files(snapshot.index_path) + stats = snapshot.index_path / 'col_stats.json' + stats.write_bytes(stats.read_bytes()[:-1]) + with warnings.catch_warnings(): + warnings.simplefilter('error') + assert graph.snapshot() == snapshot + + with pytest.warns(UserWarning, match='Reloading the damaged'): + loaded = _load(exports, location) + assert _index_files(loaded.index_path) == expected + + +def test_an_interrupted_reexport_leaves_the_last_export_loadable( + tmp_path, scratch, graph, exports, monkeypatch): + @contextmanager + def four(found, worker_ceiling=None): + yield 4 + + monkeypatch.setattr('kumo_core.sources.object_store.copies', four) + snapshot = graph.snapshot() + location, exported = _export(snapshot, exports, tmp_path) + first = _generation(exported) + open_output, finished = snapshot_io._open_output, threading.Event() + + def interrupted(client, path, size): + if path.endswith('diskgraph_config.json'): + finished.wait(10) + raise OSError('connection reset') + if path.endswith('col_stats.json'): + finished.set() + return open_output(client, path, size) + + rebuilt = _rebuilt(snapshot, tmp_path, 'rebuilt') + with monkeypatch.context() as patch: + patch.setattr(snapshot_io, '_open_output', interrupted) + patch.setattr(snapshot_io, '_ATTEMPTS', 1) + with pytest.raises(OSError, match='connection reset'): + rebuilt.export(SnapshotOutputConfig(exports, 'one')) + assert _generation(exported) == first + scratch('gpu') + assert _index_files(_load(exports, location).index_path) == _index_files( + snapshot.index_path) + + assert rebuilt.export(SnapshotOutputConfig(exports, 'one')) == location + assert _generation(exported) != first and not first.exists() + + +def test_reexporting_an_identical_snapshot_reuses_the_export( + tmp_path, scratch, graph, exports, monkeypatch, caplog): + snapshot = graph.snapshot() + location, exported = _export(snapshot, exports, tmp_path) + first = _generation(exported) + monkeypatch.setattr(snapshot_io, '_transfer', + Mock(side_effect=AssertionError('uploaded again'))) + with caplog.at_level('INFO', logger=snapshot_io.__name__): + assert snapshot.export(SnapshotOutputConfig(exports, 'one')) == location + assert _generation(exported) == first and 'reused it' in caplog.text + + (first / '_DONE').unlink() + with pytest.raises(AssertionError, match='uploaded again'): + snapshot.export(SnapshotOutputConfig(exports, 'one')) + + +def test_concurrent_exports_of_one_snapshot_both_succeed( + tmp_path, scratch, graph, exports, monkeypatch): + snapshot = graph.snapshot() + transfer, raced = snapshot_io._transfer, [] + + def racing(*args): + if not raced: + raced.append(None) + raced[0] = snapshot.export(SnapshotOutputConfig(exports, 'one')) + return transfer(*args) + + monkeypatch.setattr(snapshot_io, '_transfer', racing) + assert raced == [snapshot.export(SnapshotOutputConfig(exports, 'one'))] + assert _load(exports, raced[0]).id == snapshot.id + + +def _rebuilt(snapshot, tmp_path, name): + r"""Another build of the same logical snapshot: same id and file sizes, + one digit different in ``col_stats.json``.""" + index = (tmp_path / name / snapshot.index_path.parent.name / + snapshot.index_path.name) + shutil.copytree(snapshot.index_path, index) + (index / '.kumo_export').unlink(missing_ok=True) + stats = index / 'col_stats.json' + data = stats.read_bytes() + position = re.search(rb'[0-9]([0-9])', data).start(1) + digit = b'1' if data[position:position + 1] != b'1' else b'2' + stats.write_bytes(data[:position] + digit + data[position + 1:]) + return replace(snapshot, index_path=index) + + +def test_interleaved_exports_of_different_builds_stay_consistent( + tmp_path, scratch, graph, exports, monkeypatch): + first = graph.snapshot() + second = _rebuilt(first, tmp_path, 'second') + commit, paused = snapshot_io._commit, threading.Event() + resumed = threading.Event() + + def held(*args): + if threading.current_thread() is not threading.main_thread(): + paused.set() + resumed.wait(30) + return commit(*args) + + monkeypatch.setattr(snapshot_io, '_commit', held) + errors = [] + + def export_first(): + try: + first.export(SnapshotOutputConfig(exports, 'one')) + except Exception as error: + errors.append(error) + + worker = threading.Thread(target=export_first) + worker.start() + assert paused.wait(30) + location = second.export(SnapshotOutputConfig(exports, 'one')) + resumed.set() + worker.join(30) + assert not errors + + scratch('gpu') + assert _index_files(_load(exports, location).index_path) == _index_files( + first.index_path) + generations = [path for path in (tmp_path / 'exports' / location).iterdir() + if path.is_dir()] + assert generations == [_generation(tmp_path / 'exports' / location)] + + +def test_a_load_follows_an_export_that_commits_while_it_downloads( + tmp_path, scratch, graph, exports, monkeypatch): + first = graph.snapshot() + second = _rebuilt(first, tmp_path, 'second') + location = first.export(SnapshotOutputConfig(exports, 'one')) + load_into, calls = snapshot_io.load_into, [] + + def overtaken(*args): + calls.append(args[2]) + if len(calls) == 1: + second.export(SnapshotOutputConfig(exports, 'one')) + return load_into(*args) + + monkeypatch.setattr(snapshot_io, 'load_into', overtaken) + scratch('gpu') + loaded = _load(exports, location) + assert len(calls) == 2 and calls[0] != calls[1] + assert _index_files(loaded.index_path) == _index_files(second.index_path) + + +def test_a_cached_load_is_reused_only_for_the_same_export( + tmp_path, scratch, graph, exports): + first = graph.snapshot() + second = _rebuilt(first, tmp_path, 'second') + one = first.export(SnapshotOutputConfig(exports, 'one')) + two = second.export(SnapshotOutputConfig(exports, 'two')) + assert _index_files(first.index_path) != _index_files(second.index_path) + + scratch('gpu') + for location, build in ((one, first), (two, second), (one, first)): + loaded = _load(exports, location) + assert _index_files(loaded.index_path) == _index_files( + build.index_path) + + +def test_overlapping_marker_commits_leave_one_whole_marker( + tmp_path, monkeypatch): + root = tmp_path / 'location' + markers = { + name * 32: json.dumps({'generation': name * 32, 'crc': crc}).encode() + for name, crc in (('a', 1), ('b', 1234567)) + } + for generation in markers: + (root / generation).mkdir(parents=True) + write, halfway = snapshot_io._write, threading.Barrier(2) + + def interleaved(filesystem, path, payload): + if not path.endswith('/kumo_snapshot.json'): + return write(filesystem, path, payload) + with filesystem.open_output_stream(path) as writer: + writer.write(payload[:5]) + halfway.wait(10) + writer.write(payload[5:]) + + monkeypatch.setattr(snapshot_io, '_write', interleaved) + threads = [ + threading.Thread(target=snapshot_io._commit, + args=(LocalFileSystem(), str(root), generation, + marker)) + for generation, marker in markers.items() + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join(30) + committed = (root / 'kumo_snapshot.json').read_bytes() + assert committed in markers.values() + + +def test_an_unreadable_marker_is_corrupt_on_load_and_reexported_by_a_store( + tmp_path, scratch, graph): + snapshot = graph.snapshot() + settings = kumo.MLOps(artifact_root=str(tmp_path / 'root')) + snapshot_io.store(snapshot, settings) + location = snapshot_io.export_location(snapshot, 'snapshots') + marker = tmp_path / 'root' / location / 'kumo_snapshot.json' + marker.write_bytes(b'{"generation": "a"}34567}') + connector = kumo.FileConnector(root_dir=str(tmp_path / 'root')) + with pytest.raises(ValueError, match='corrupt.*unreadable'): + _load(connector, location) + + snapshot_io.store(snapshot, settings) + assert _load(connector, location).id == snapshot.id + + +def test_loading_an_export_on_the_exporting_machine_skips_the_checksum_pass( + tmp_path, scratch, graph, exports, monkeypatch): + snapshot = graph.snapshot() + location, _ = _export(snapshot, exports, tmp_path) + monkeypatch.setattr( + 'diskgraph.snapshot_range_checksums', + Mock(side_effect=AssertionError('the local index was re-read'))) + assert _load(exports, location) == snapshot + + +def test_failed_upload_is_retried_with_a_fresh_client( + scratch, graph, exports, monkeypatch): + open_output, clients = snapshot_io._open_output, [] + + def expiring(client, path, size): + if path.endswith('col_stats.json'): + clients.append(client) + if len(clients) == 1: + raise OSError('ExpiredToken') + return open_output(client, path, size) + + monkeypatch.setattr(snapshot_io, '_open_output', expiring) + graph.snapshot().export(SnapshotOutputConfig(exports, 'one')) + assert len(clients) == 2 and clients[0] is not clients[1] + + +def test_reexport_after_redaction_publishes_the_redacted_manifest( + tmp_path, scratch, graph, exports): + full = graph.snapshot() + location, exported = _export(full, exports, tmp_path) + first = _generation(exported) + redacted = graph.snapshot(record_sources=False) + assert redacted.id == full.id + assert redacted.export(SnapshotOutputConfig(exports, 'one')) == location + committed = _generation(exported) + assert committed != first + assert (committed / '_provenance.json').read_bytes() == ( + redacted.index_path / '_provenance.json').read_bytes() + assert json.loads((committed / '_provenance.json').read_text())[ + 'record_sources'] is False + + +def test_reexport_replaces_an_export_missing_a_remote_file( + tmp_path, scratch, graph, exports): + snapshot = graph.snapshot() + location, exported = _export(snapshot, exports, tmp_path) + first = _generation(exported) + (first / 'index' / 'col_stats.json').unlink() + assert snapshot.export(SnapshotOutputConfig(exports, 'one')) == location + assert _generation(exported) != first + scratch('empty') + assert _index_files(_load(exports, location).index_path) == _index_files( + snapshot.index_path) + + +def test_reexport_supersedes_a_damaged_export(tmp_path, scratch, graph, + exports): + snapshot = graph.snapshot() + location, exported = _export(snapshot, exports, tmp_path) + damaged = _generation(exported) + stats = damaged / 'index/col_stats.json' + stats.write_bytes(stats.read_bytes()[:3]) + (damaged / 'index/index_metadata.json').unlink() + _edit_manifest(exported, lambda m: m.update(snapshot_id='other')) + assert snapshot.export(SnapshotOutputConfig(exports, 'one')) == location + assert not damaged.exists() + scratch('gpu') + assert _load(exports, location).id == snapshot.id + + +def test_failed_load_leaves_nothing_and_restarts_cleanly( + tmp_path, scratch, graph, exports): + snapshot = graph.snapshot() + location, exported = _export(snapshot, exports, tmp_path) + stats = _generation(exported) / 'index/col_stats.json' + original = stats.read_bytes() + stats.write_bytes(original[:-1] + bytes([original[-1] ^ 1])) + + scratch('gpu') + with pytest.raises(ValueError, match='corrupt.*checksum'): + _load(exports, location) + graph_dir = tmp_path / 'gpu/snapshots' / snapshot.index_path.parent.name + assert not [path for path in graph_dir.iterdir() if path.suffix != '.lock'] + + stats.write_bytes(original) + (graph_dir / f'{snapshot.index_path.name}.loading').mkdir() + loaded = _load(exports, location) + assert _index_files(loaded.index_path) == _index_files(snapshot.index_path) + assert [ + path.name for path in graph_dir.iterdir() if path.suffix != '.lock' + ] == [snapshot.index_path.name] + + +def test_load_refuses_before_downloading_when_disk_is_too_small( + tmp_path, scratch, graph, exports, monkeypatch): + location, _ = _export(graph.snapshot(), exports, tmp_path) + scratch('gpu') + monkeypatch.setattr(snapshot_io.os, 'statvfs', + lambda path: SimpleNamespace(f_frsize=1, f_bavail=0)) + monkeypatch.setattr(snapshot_io, '_download', + Mock(side_effect=AssertionError('downloaded'))) + with pytest.raises(OSError, match='free bytes'): + _load(exports, location) + + +@pytest.mark.parametrize('index, edit', [ + (False, lambda manifest: manifest.update(data_version='../../victim')), + (False, lambda manifest: manifest.pop('crc32')), + (True, lambda manifest: manifest['files'][0].update( + rel_path='../../victim/x')), + (True, lambda manifest: manifest['files'].append(manifest['files'][0])), +]) +def test_load_rejects_a_hostile_or_malformed_manifest(tmp_path, scratch, graph, + exports, index, edit): + location, exported = _export(graph.snapshot(), exports, tmp_path) + victim = tmp_path / 'victim' + victim.mkdir() + (victim / 'keep').write_text('x') + (_edit_index_manifest if index else _edit_manifest)(exported, edit) + scratch('gpu/a/b') + with pytest.raises(ValueError): + _load(exports, location) + assert (victim / 'keep').read_text() == 'x' + + +def test_load_rejects_unknown_formats_and_accepts_unknown_fields( + tmp_path, scratch, graph, exports): + snapshot = graph.snapshot() + location, exported = _export(snapshot, exports, tmp_path) + scratch('gpu') + + _edit_manifest(exported, lambda m: m.update(format=2)) + with pytest.raises(ValueError, match='format 2'): + _load(exports, location) + + _edit_manifest(exported, lambda m: m.update(format=1, added_later={'x': 1})) + other = kumo.Graph(tables={'users': graph.tables['users']}) + with pytest.warns(UserWarning, match=r'another graph\. Schema changes: ' + r'table events added\.$'): + loaded = _load(exports, location, graph=other) + assert loaded.id == snapshot.id + + +def test_load_only_logs_a_graph_that_moved_with_the_same_schema( + tmp_path, scratch, graph, exports, caplog): + location, _ = _export(graph.snapshot(), exports, tmp_path) + shutil.copytree(tmp_path / 'data', tmp_path / 'prod') + prod = kumo.FileConnector(root_dir=str(tmp_path / 'prod')) + moved = kumo.Graph( + tables={ + name: replace(table, connector=prod) + for name, table in graph.tables.items() + }, edges=graph.edges) + assert moved.id != graph.id + + with caplog.at_level('INFO', logger='kumo.engine.graph'), \ + warnings.catch_warnings(): + warnings.filterwarnings('error', message='.*another graph') + _load(exports, location, graph=moved) + assert 'only table connectors or locations differ' in caplog.text + + +@pytest.mark.parametrize('recorded', ['newer', 'missing']) +def test_snapshot_manifest_travels_unchanged_or_is_reconstructed( + tmp_path, scratch, graph, exports, recorded): + snapshot = graph.snapshot() + path = snapshot.index_path / '_provenance.json' + if recorded == 'newer': + path.write_bytes(b'{"schema_version": 99, "future": [1]}') + else: + path.unlink() + location, exported = _export(snapshot, exports, tmp_path) + scratch('gpu') + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + loaded = _load(exports, location) + shipped = (_generation(exported) / '_provenance.json').read_bytes() + assert (loaded.index_path / '_provenance.json').read_bytes() == shipped + if recorded == 'newer': + assert shipped == b'{"schema_version": 99, "future": [1]}' + else: + manifest = loaded.manifest + assert (manifest.id, manifest.reconstructed) == (snapshot.id, True) + assert manifest.parents.source_tables == {} + + +def test_warehouse_connectors_are_not_supported(scratch, graph): + connector = SimpleNamespace(sink_backend=object()) + with pytest.raises(NotImplementedError, match='FileConnector'): + graph.snapshot().export(SnapshotOutputConfig(connector, 'x')) + with pytest.raises(NotImplementedError, match='FileConnector'): + _load(connector, 'x') + + +def test_load_explains_an_index_the_installed_diskgraph_cannot_open( + tmp_path, scratch, graph, exports): + location, exported = _export(graph.snapshot(), exports, tmp_path) + metadata = _generation(exported) / 'index/index_metadata.json' + garbage = b'{' * metadata.stat().st_size + metadata.write_bytes(garbage) + + def unreadable(manifest): + manifest['crc32']['index_metadata.json'] = [zlib.crc32(garbage)] + manifest['environment']['kumo_engine'] = '0.0.1' + + _edit_manifest(exported, unreadable) + scratch('gpu') + with pytest.raises(RuntimeError, match='exported by kumo-engine 0.0.1'): + _load(exports, location) + + +@pytest.mark.parametrize('scheme', ['s3', 'gs']) +def test_object_store_round_trip(tmp_path, scratch, graph, monkeypatch, + scheme): + storage = tmp_path / 'storage' + storage.mkdir() + + class Bucket(SubTreeFileSystem): + def __init__(self, *args, **kwargs): + super().__init__(str(storage), LocalFileSystem()) + + monkeypatch.setattr('kumo_s3.filesystem', Bucket) + monkeypatch.setattr('kumo_gcs.backend.fs.GcsFileSystem', Bucket) + snapshot = graph.snapshot() + connector = kumo.S3Connector() if scheme == 's3' else kumo.GCSConnector() + location = snapshot.export( + SnapshotOutputConfig(connector, f'{scheme}://bucket/snapshots/')) + graph_id, data_version = snapshot.id.split('@') + assert location == f'{scheme}://bucket/snapshots/{graph_id}/{data_version}' + scratch('gpu') + loaded = _load(connector, location, graph=graph) + assert _index_files(loaded.index_path) == _index_files(snapshot.index_path) + + +def _set_mtime(path, seconds): + os.utime(path, (seconds, seconds)) + + +def _rewrite_events(data): + frame = pd.read_parquet(data / 'events.parquet') + frame.iloc[::-1].to_parquet(data / 'events.parquet') + _set_mtime(data / 'events.parquet', 1_700_009_123) + + +@pytest.fixture +def dated(tmp_path, graph): + data = tmp_path / 'data' + for name in ('users', 'events'): + _set_mtime(data / f'{name}.parquet', 1_700_000_123) + return data + + +@pytest.mark.parametrize('record_sources', [True, False]) +def test_source_changes_reads_stamps_without_staging( + scratch, graph, dated, monkeypatch, record_sources): + from kumo_core.sources.object_store import ObjectStoreSourceReader + + from kumo.engine.connector.file import _FileSourceBackend + + _set_mtime(dated / 'users.parquet', 1_700_002_800) + snapshot = graph.snapshot(record_sources=record_sources) + fail = Mock(side_effect=AssertionError('must not stage')) + monkeypatch.setattr('kumo_core.sources.prepare.read_table', fail) + monkeypatch.setattr(ObjectStoreSourceReader, 'prepare', fail) + monkeypatch.setattr(_FileSourceBackend, 'prepare_ingest', fail) + assert snapshot.source_changes(graph) == { + 'users': {'status': 'unchanged', 'recorded': '1700002800', + 'current': '1700002800'}, + 'events': {'status': 'unchanged', 'recorded': '1700000123', + 'current': '1700000123'}, + } + _rewrite_events(dated) + assert snapshot.source_changes(graph)['events'] == { + 'status': 'changed', 'recorded': '1700000123', 'current': '1700009123'} + + monkeypatch.setattr(_FileSourceBackend, 'get_data_version', + Mock(side_effect=OSError('listing failed'))) + assert snapshot.source_changes(graph)['events'] == { + 'status': 'unknown', 'recorded': '1700000123', 'current': None} + + +def test_source_changes_compares_directory_listings(tmp_path, scratch): + parts = tmp_path / 'dir' / 'users' + parts.mkdir(parents=True) + pd.DataFrame({'ID': range(5)}).to_parquet(parts / 'part-0.parquet') + connector = kumo.FileConnector(root_dir=str(tmp_path / 'dir')) + graph = kumo.Graph(tables={ + 'users': kumo.Table.from_source_table(connector['users'], + primary_key='ID') + }, edges=[]) + snapshot = graph.snapshot() + assert snapshot.source_changes(graph)['users']['status'] == 'unchanged' + pd.DataFrame({'ID': range(5, 9)}).to_parquet(parts / 'part-1.parquet') + assert snapshot.source_changes(graph)['users']['status'] == 'changed' + + +@pytest.mark.parametrize('record_sources', [True, False]) +def test_source_changes_notices_a_rebound_source(tmp_path, scratch, graph, + dated, record_sources): + import shutil + + from kumoapi.graph import Edge + from kumoapi.typing import Stype + + snapshot = graph.snapshot(record_sources=record_sources) + shutil.copytree(dated, tmp_path / 'other') + connector = kumo.FileConnector(root_dir=str(tmp_path / 'other')) + rebound = kumo.Graph( + tables={ + 'users': + kumo.Table.from_source_table(connector['users.parquet'], + primary_key='ID'), + 'events': + kumo.Table.from_source_table( + connector['events.parquet'], primary_key='EID', + time_column='TS').with_column('USER_ID', stype=Stype.ID), + }, edges=[Edge(src_table='events', fkey='USER_ID', dst_table='users')]) + + changes = snapshot.source_changes(rebound) + assert {change['status'] for change in changes.values()} == {'unknown'} + assert changes['events']['recorded'] == changes['events']['current'] + + +@pytest.mark.parametrize('edit', ['absent', 'reconstructed', 'hour_bucket']) +def test_source_changes_unknown_without_a_comparable_version( + scratch, graph, dated, edit): + snapshot = graph.snapshot() + path = snapshot.index_path / '_provenance.json' + manifest = json.loads(path.read_text()) + if edit == 'absent': + path.unlink() + elif edit == 'reconstructed': + path.write_text(json.dumps({**manifest, 'reconstructed': True})) + else: + for entry in manifest['parents']['source_tables'].values(): + entry['version_kind'] = 'hour_bucket' + path.write_text(json.dumps(manifest)) + assert {change['status'] for change in snapshot.source_changes( + graph).values()} == {'unknown'} + + +def test_load_warns_only_with_graph_and_only_for_changed_tables( + tmp_path, scratch, graph, dated, exports): + location, _ = _export(graph.snapshot(), exports, tmp_path) + scratch('gpu') + with warnings.catch_warnings(): + warnings.simplefilter('error') + _load(exports, location, graph=graph) + _rewrite_events(dated) + _load(exports, location) + with pytest.warns(UserWarning) as caught: + _load(exports, location, graph=graph) + [message] = [str(warning.message) for warning in caught] + assert 'source table(s) events;' in message + + +class _Bytes: + def __init__(self, data): + self.data = data + + def read_at(self, offset, length): + return self.data[offset:offset + length] + + def close(self): + pass + + +def test_a_sink_with_only_a_filesystem_holds_exports(tmp_path, scratch, + graph): + from kumo_core.chunked_fs import ChunkedFileHandler + from pyarrow.fs import FileSelector, FileType, PyFileSystem + + class Store(ChunkedFileHandler): + part_bytes = 1000 + + def __init__(self): + self.objects = {} + + def get_type_name(self): + return 'memory' + + def _size(self, path): + data = self.objects.get(path) + return None if data is None else len(data) + + def _list(self, directory): + return {p: len(d) for p, d in self.objects.items() + if p.startswith(directory + '/')} + + def _put(self, local, path): + self.objects[path] = Path(local).read_bytes() + + def _part(self, path, size): + return _Bytes(self.objects[path]) + + def _remove(self, path, *, tree): + for name in list(self.objects): + if name.startswith(path + '/') if tree else name == path: + del self.objects[name] + + store = Store() + sink = SimpleNamespace( + holds_files=lambda connector, path: path.startswith('@'), + filesystem=lambda connector, path: (PyFileSystem(store), + path.rstrip('/'))) + connector = SimpleNamespace(sink_backend=sink) + snapshot = graph.snapshot() + location = snapshot.export(SnapshotOutputConfig(connector, '@stage/s')) + assert any('.chunks/' in name for name in store.objects) + scratch('gpu') + loaded = _load(connector, location) + assert loaded.id == snapshot.id + filesystem, root = sink.filesystem(connector, location) + with filesystem.open_input_stream(f'{root}/kumo_snapshot.json') as reader: + generation = json.loads(reader.readall())['generation'] + filesystem, root = sink.filesystem(connector, + f'{location}/{generation}/index') + for info in filesystem.get_file_info(FileSelector(root, recursive=True)): + if info.type is FileType.File: + with filesystem.open_input_file(info.path) as reader: + assert (loaded.index_path / info.path[len(root) + 1:] + ).read_bytes() == reader.read() + with pytest.raises(NotImplementedError, match='file destination'): + snapshot.export(SnapshotOutputConfig(connector, 'DB.SCHEMA')) + + +def test_a_mlops_block_stores_each_snapshot_once(tmp_path, scratch, graph, + monkeypatch): + root = tmp_path / 'artifacts' + with kumo.mlops(kumo.MLOps(artifact_root=str(root))): + snapshot = graph.snapshot() + monkeypatch.setattr(snapshot_io, 'export', + Mock(side_effect=AssertionError('re-exported'))) + assert graph.snapshot() == snapshot + graph_id, data_version = snapshot.id.split('@') + stored = root / 'snapshots' / graph_id / data_version + assert (stored / 'kumo_snapshot.json').is_file() + assert (_generation(stored) / '_provenance.json').is_file() + assert _load(kumo.FileConnector(root_dir=str(root)), + f'snapshots/{graph_id}/{data_version}') == snapshot + graph.snapshot(mlops=kumo.MLOps(artifact_root=str(tmp_path / 'off'), + export={'model'})) + assert not (tmp_path / 'off').exists() diff --git a/packages/kumo-engine/test/test_snapshot_export_live.py b/packages/kumo-engine/test/test_snapshot_export_live.py new file mode 100644 index 000000000..c28b19712 --- /dev/null +++ b/packages/kumo-engine/test/test_snapshot_export_live.py @@ -0,0 +1,45 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import contextlib +import os +import uuid + +import pytest + +import kumo.engine as kumo +from kumo.engine import SnapshotOutputConfig +from kumo.engine._scratch import get_scratch_dir +from kumo.engine._snapshot_io import _location + +ROOTS = { + 's3': os.environ.get('KUMO_SNAPSHOT_LIVE_S3_ROOT'), + 'gs': os.environ.get('KUMO_SNAPSHOT_LIVE_GCS_ROOT'), +} + + +@pytest.mark.parametrize('scheme', ['s3', 'gs']) +def test_snapshot_round_trip_through_object_store(tmp_path, graph, scheme): + root = ROOTS[scheme] + if not root: + pytest.skip('set KUMO_SNAPSHOT_LIVE_S3_ROOT / _GCS_ROOT to run') + connector = kumo.S3Connector() if scheme == 's3' else kumo.GCSConnector() + path = f"{root.rstrip('/')}/engine-snapshot-{uuid.uuid4().hex[:12]}" + previous = get_scratch_dir() + try: + kumo.configure(scratch_dir=tmp_path / 'cpu') + snapshot = graph.snapshot() + location = snapshot.export(SnapshotOutputConfig(connector, path)) + kumo.configure(scratch_dir=tmp_path / 'gpu') + loaded = kumo.GraphSnapshot.load_from_connector( + connector, path=location, graph=graph) + finally: + kumo.configure(scratch_dir=previous) + filesystem, prefix = _location(connector, path) + filesystem.delete_dir_contents(prefix, missing_dir_ok=True) + with contextlib.suppress(OSError): + filesystem.delete_dir(prefix) + assert loaded.id == snapshot.id + assert (loaded.index_path / 'col_stats.json').read_bytes() == ( + snapshot.index_path / 'col_stats.json').read_bytes() diff --git a/packages/kumo-engine/test/test_snapshot_manifest.py b/packages/kumo-engine/test/test_snapshot_manifest.py new file mode 100644 index 000000000..58ad7eab0 --- /dev/null +++ b/packages/kumo-engine/test/test_snapshot_manifest.py @@ -0,0 +1,256 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import hashlib +import json +import os +from dataclasses import replace +from datetime import datetime +from types import SimpleNamespace + +import numpy as np +import pandas as pd +import pytest +from kumo_core.custom_sql import custom_sql_synthetic_name +from kumo_core.sources.identity import SourceVersion +from kumo_core.types import ConnectorType +from kumo_core.version import get_kumo_version +from kumoapi.graph import Edge +from kumoapi.typing import Stype + +import kumo.engine._provenance as provenance +from kumo.engine import ( + BigQueryConnector, + DatabricksConnector, + FileConnector, + Graph, + S3Connector, + SnowflakeConnector, + Table, +) +from kumo.engine._manifest import ( + SnapshotManifest, + SourceTableEntry, + read_manifest, + write_manifest, +) +from kumo.engine.connector.identity import ( + DatabricksIdentity, + SnowflakeIdentity, +) +from kumo.engine.graph import SNAPSHOT_MANIFEST_FILE, GraphSnapshot + + +@pytest.fixture +def make_snapshot(tmp_path, monkeypatch): + monkeypatch.setattr('kumo.engine._scratch._root', tmp_path / 'scratch') + data = tmp_path / 'data' + data.mkdir() + pd.DataFrame({'ID': np.arange(5)}).to_parquet(data / 'ents.parquet') + + def make(events: int | None = None, **options) -> GraphSnapshot: + if events is not None: + pd.DataFrame({ + 'EID': np.arange(events), + 'ENT_ID': np.arange(events) % 5, + }).to_parquet(data / 'events.parquet') + # Single-file versions are mtime seconds. + stamp = 1.7e9 + events + os.utime(data / 'events.parquet', (stamp, stamp)) + connector = FileConnector(root_dir=str(data)) + return Graph( + { + 'ents': + Table.from_source_table(connector['ents'], + primary_key='ID').infer_metadata(), + 'events': + Table.from_source_table( + connector['events.parquet'], + primary_key='EID').infer_metadata().with_column( + 'ENT_ID', stype=Stype.ID), + }, [Edge('events', 'ENT_ID', 'ents')]).snapshot(**options) + + return make + + +def test_snapshot_manifest_records_sources_and_diffs(make_snapshot, + tmp_path): + before = make_snapshot(30) + manifest = before.manifest + + assert manifest.kind == 'snapshot' + assert SNAPSHOT_MANIFEST_FILE == '_provenance.json' + assert manifest.id == before.id + data = (tmp_path / 'data').resolve() + assert {a: e.location + for a, e in manifest.parents.source_tables.items()} == { + 'ents': str(data / 'ents.parquet'), + 'events': str(data / 'events.parquet'), + } + assert {e.version_kind + for e in manifest.parents.source_tables.values() + } == {'timestamp'} + assert manifest.extra == {} + assert manifest.environment.kumo_engine == get_kumo_version() + assert manifest.environment.diskgraph + assert 'scheduler' not in manifest.to_dict() + assert datetime.fromisoformat(manifest.created_at).tzinfo is not None + assert str(tmp_path / 'scratch') not in json.dumps(manifest.to_dict()) + + after = make_snapshot(40) + sources = manifest.parents.source_tables + new_sources = after.manifest.parents.source_tables + assert new_sources['ents'] == sources['ents'] + assert new_sources['events'].version != sources['events'].version + + +def test_index_without_manifest_is_reused_and_backfilled( + make_snapshot, monkeypatch): + import kumo_core.diskgraph.ingest as ingest + + snapshot = make_snapshot(30) + original = snapshot.manifest + (snapshot.index_path / SNAPSHOT_MANIFEST_FILE).unlink() + monkeypatch.setattr(ingest, 'execute_diskgraph_ingest', + lambda *args: pytest.fail('cached index re-ingested')) + + reused = make_snapshot() + + assert reused == snapshot + assert reused.manifest == replace(original, created_at=None, + environment=None, + reconstructed=True) + + +def test_manifest_write_failure_leaves_snapshot_unchanged( + make_snapshot, monkeypatch): + def fail(path, manifest): + raise OSError('disk full') + + monkeypatch.setattr('kumo.engine.graph.write_manifest', fail) + with pytest.warns(UserWarning, match='disk full'): + snapshot = make_snapshot(30) + + assert snapshot.manifest is None + assert (snapshot.index_path / '_SUCCESS').is_file() + + +@pytest.mark.parametrize('content', [ + None, '{"schema_version": 2}', '{"schema_version": true}', '{"kind": ' +]) +def test_unsupported_manifest_is_ignored(tmp_path, content): + if content is not None: + (tmp_path / SNAPSHOT_MANIFEST_FILE).write_text(content) + with pytest.warns(UserWarning): + assert GraphSnapshot('s', tmp_path, {}).manifest is None + else: + assert GraphSnapshot('s', tmp_path, {}).manifest is None + + +@pytest.mark.parametrize('connector, name, sql, location', [ + (S3Connector(root_dir='s3://bucket-a/data/'), 'exports/tt/', None, + 's3://bucket-a/data/exports/tt'), + (S3Connector(root_dir='s3://bucket-a/data/'), 'q4#1 50%.parquet', None, + 's3://bucket-a/data/q4%231%2050%25.parquet'), + (S3Connector(), 's3://bucket-a/sales%23q4.parquet', None, + 's3://bucket-a/sales%23q4.parquet'), + (SnowflakeConnector(warehouse='wh', database='db', schema_name='sch', + identity=SnowflakeIdentity(account='org-acct')), + 'TT', None, 'org-acct/"DB"."SCH"."TT"'), + (SnowflakeConnector(warehouse='wh', database='db', schema_name='sch'), + 'TT', None, '"DB"."SCH"."TT"'), + (DatabricksConnector( + catalog='cat', schema_name='sch', warehouse_id='w', + identity=DatabricksIdentity(host='https://workspace.example.com', + client_id='c', client_secret='s')), + 'cat.sch.tt', None, '`cat`.`sch`.`tt`'), + (BigQueryConnector(project_id='data-project', dataset_id='analytics'), + 'tt', None, 'data-project.analytics.tt'), + (SimpleNamespace(source_type='OTHER', source_backend=None), 'TT', None, + 'TT'), + (None, 'synthetic', 'SELECT 1', 'sql:' + + hashlib.sha256(b'SELECT 1').hexdigest()), +]) +def test_source_location(connector, name, sql, location): + assert provenance.source_location(connector, name, sql) == location + + +def test_unrecordable_location_is_null(): + with pytest.warns(UserWarning, match='location'): + entry = provenance.source_table(None, 'users', None, 7) + assert entry == SourceTableEntry(connector=None, location=None, + version='7', version_kind='unknown') + + +@pytest.mark.parametrize('record_sources', [True, False]) +def test_custom_sql_records_text_unless_sources_are_off(record_sources): + version = SourceVersion('abc', 'hour_bucket', {'orders': '17'}) + entry = provenance.source_table( + SnowflakeConnector(warehouse='wh', database='db', schema_name='sch'), + 'q', 'SELECT * FROM orders', version, + record_sources=record_sources).to_dict() + expected = { + 'connector': 'SNOWFLAKE', + 'location': 'sql:' + hashlib.sha256(b'SELECT * FROM orders').hexdigest(), + 'version': 'abc', + 'version_kind': 'hour_bucket', + 'dependencies': {'orders': '17'}, + } + if record_sources: + expected.update(sql='SELECT * FROM orders', dialect='snowflake', + namespace=['db', 'sch']) + reformatted = provenance.source_table( + SnowflakeConnector(warehouse='wh', database='db', schema_name='sch'), + 'q', 'select *\n from ORDERS -- note', version) + assert entry.pop('normalized_sql_fingerprint') == ( + reformatted.normalized_sql_fingerprint) == custom_sql_synthetic_name( + 'SELECT * FROM orders', ConnectorType.SNOWFLAKE) + assert reformatted.location != entry['location'] + assert entry == expected + + +def test_record_sources_false_keeps_table_names_only(make_snapshot, + tmp_path): + names = {'ents': 'ents', 'events': 'events.parquet'} + redacted = make_snapshot(30, record_sources=False) + sources = redacted.manifest.parents.source_tables + assert {alias: e.location for alias, e in sources.items()} == names + assert str(tmp_path) not in json.dumps(redacted.manifest.to_dict()) + assert redacted.manifest.record_sources is False + assert redacted.record_sources is False + + full = make_snapshot(record_sources=True) + assert full == redacted + assert full.record_sources is True + assert full.manifest.parents.source_tables['ents'].location == str( + (tmp_path / 'data' / 'ents.parquet').resolve()) + assert full.manifest.created_at == redacted.manifest.created_at + + +def test_toggling_record_sources_keeps_unknown_parent_fields(make_snapshot): + written = make_snapshot(30).index_path / SNAPSHOT_MANIFEST_FILE + data = json.loads(written.read_text()) + data['parents']['future'] = {'kept': True} + written.write_text(json.dumps(data)) + + redacted = make_snapshot(record_sources=False) + assert redacted.manifest.record_sources is False + assert redacted.manifest.parents.extra == {'future': {'kept': True}} + + +def test_snapshot_manifest_round_trips_unchanged(make_snapshot, tmp_path): + written = make_snapshot(30).index_path / SNAPSHOT_MANIFEST_FILE + newer = json.loads(written.read_text()) + newer['future'] = {'kept': [1, 2]} + newer['environment']['cuda'] = '13.0' + newer['parents']['source_tables']['ents']['row_count'] = 5 + source = tmp_path / 'newer.json' + source.write_text(json.dumps(newer, indent=2, sort_keys=True) + '\n') + + manifest = read_manifest(source, SnapshotManifest) + write_manifest(tmp_path / 'out.json', manifest) + + assert (tmp_path / 'out.json').read_bytes() == source.read_bytes() + assert manifest.extra == {'future': {'kept': [1, 2]}} + assert manifest.environment.extra == {'cuda': '13.0'} diff --git a/packages/kumo-engine/test/test_trainer_registry.py b/packages/kumo-engine/test/test_trainer_registry.py index 931200303..7ea9ba950 100644 --- a/packages/kumo-engine/test/test_trainer_registry.py +++ b/packages/kumo-engine/test/test_trainer_registry.py @@ -141,11 +141,11 @@ def test_write_then_read_training_spec_round_trips(tmp_path: Path) -> None: def test_training_spec_manifest_update_is_atomic_and_retryable( tmp_path: Path, monkeypatch) -> None: - import kumo.engine.trainer._spec as spec_module + import kumo_core.io.filesystem as fs_module bundle = _fake_published_bundle(tmp_path) original_manifest = (bundle / MANIFEST_FILE).read_bytes() - real_replace = spec_module.os.replace + real_replace = fs_module.os.replace failed = False def fail_first_manifest_replace(source, destination): @@ -155,7 +155,7 @@ def fail_first_manifest_replace(source, destination): raise OSError('simulated interrupted manifest commit') return real_replace(source, destination) - monkeypatch.setattr(spec_module.os, 'replace', fail_first_manifest_replace) + monkeypatch.setattr(fs_module.os, 'replace', fail_first_manifest_replace) with pytest.raises(OSError, match='interrupted manifest commit'): write_training_spec(bundle, predictive_query='PREDICT COUNT(*)', model_plan=ModelPlan(), snapshot_id='snap-123', @@ -210,13 +210,13 @@ def test_finish_fit_persists_spec_without_export_before_exposing_state( trainer = Trainer(ModelPlan()) result = TrainingResult({'auc': .9}, {}, bundle) - assert trainer._finish_fit(result, predictive_query='PREDICT COUNT(*)', - snapshot_id='snap-123', - training_table_generation=_generation_plan(), - export_config=None) is result + finished = trainer._finish_fit( + result, predictive_query='PREDICT COUNT(*)', snapshot_id='snap-123', + training_table_generation=_generation_plan(), export_config=None) assert (bundle / SPEC_FILE).is_file() - assert trainer._training_result is result + assert finished.bundle_id == Trainer.load(bundle).result.id + assert trainer._training_result is finished assert trainer._snapshot_id == 'snap-123' assert trainer._training_query == 'PREDICT COUNT(*)' assert result.published_model is None diff --git a/packages/kumo-engine/test/test_training_table_export.py b/packages/kumo-engine/test/test_training_table_export.py index cedee0e8c..fd97a20f1 100644 --- a/packages/kumo-engine/test/test_training_table_export.py +++ b/packages/kumo-engine/test/test_training_table_export.py @@ -4,7 +4,9 @@ import json from collections import Counter +from dataclasses import replace from functools import partial +from types import SimpleNamespace import numpy as np import pandas as pd @@ -17,6 +19,13 @@ from kumoapi.typing import Stype import kumo.engine as kumo +from kumo.engine._manifest import ( + SourceTableEntry, + TrainingTableManifest, + TrainingTableParents, + read_manifest, + write_manifest, +) from kumo.engine._scratch import get_scratch_dir from kumo.engine.trainer.output_config import TrainingTableOutputConfig @@ -188,7 +197,9 @@ def test_forecasting_export_roundtrip(graph): assert _seeds(reimported)[1] == counts -def test_reexport_and_trailing_slash(graph): +def test_reexport_and_trailing_slash(graph, monkeypatch, caplog): + import kumo_core.artifact_export + _, _, connector, out = graph query = 'PREDICT ents.LABEL FOR EACH ents.ID' _, table, first = _export(graph, query, 'retry') @@ -200,8 +211,12 @@ def listing(): before = listing() config = TrainingTableOutputConfig(connector, 'retry') - table.export(config) - table.export(config) + monkeypatch.setattr(kumo_core.artifact_export, 'export_selected_files', + lambda *a, **k: pytest.fail('uploaded again')) + with caplog.at_level('INFO', logger='kumo_core.object_store_sink'): + table.export(config) + table.export(config) + assert 'reused it' in caplog.text assert config.table_name == 'retry/' assert listing() == before assert before[1] == first.num_rows == sum(_seeds(table)[1].values()) @@ -225,9 +240,9 @@ def test_custom_link_table_ingests_multiple_weighted_parts( stage_source_table = _custom_source.stage_source_table def stage_multipart_source(*args, **kwargs): - path = stage_source_table(*args, **kwargs) + path, version = stage_source_table(*args, **kwargs) assert len(list(path.glob('*.parquet'))) > 1 - return path + return path, version monkeypatch.setattr(_custom_source, 'stage_source_table', stage_multipart_source) @@ -292,3 +307,155 @@ def test_custom_training_table_declares_integer_timestamp_unit(graph, unit, for stage in seeds: columns = ['ENTITY1', 'TIMESTAMP', 'TARGET1'] assert _rows(seeds[stage], columns) == _rows(original_seeds[stage], columns) + + +def test_manifest_and_export_sidecar(graph, tmp_path): + from kumo_core.training.export import export_train_table + + _, snapshot, connector, out = graph + query = 'PREDICT ents.LABEL FOR EACH ents.ID' + pquery, table, _ = _export(graph, query, 'lineage') + manifest = table.manifest + assert manifest.id == table.id and table.id.startswith('tt-') + assert table.id != _spec(table)['tt_id'] + assert manifest.parents == TrainingTableParents( + snapshot=snapshot.manifest) + assert manifest.query == query + assert manifest.extra == {} + assert manifest.split_seed == 0 + again = pquery.generate_training_table(snapshot) + assert again.id == table.id + assert again.tt_id != table.tt_id + + directory = out / 'lineage' + baseline = tmp_path / 'baseline' + export_train_table(str(table.staging_dir), str(baseline)) + for part in baseline.glob('*.parquet'): + assert (directory / part.name).read_bytes() == part.read_bytes() + assert json.loads( + (directory / '_provenance.json').read_bytes()) == manifest.to_dict() + + reloaded = kumo.TrainingTable.load_from_connector( + connector, path='lineage', pq=pquery, snapshot=snapshot) + manifest = reloaded.manifest + assert reloaded.id != table.id + assert manifest.parents.source_table == SourceTableEntry( + connector='LOCAL_FILE', + location=str(directory.resolve()), + version=connector.source_backend.get_data_version( + connector, 'lineage'), + version_kind='listing', + ) + assert _seeds(reloaded)[1] == _seeds(table)[1] + + +def test_training_table_manifest_round_trips_unchanged(graph, tmp_path): + _, table, _ = _export(graph, 'PREDICT ents.LABEL FOR EACH ents.ID', + 'round_trip') + newer = json.loads( + (table.staging_dir / '_provenance.json').read_text()) + newer['future'] = 1 + newer['parents']['label_source'] = 'x' + source = tmp_path / 'newer.json' + source.write_text(json.dumps(newer, indent=2, sort_keys=True) + '\n') + + manifest = read_manifest(source, TrainingTableManifest) + write_manifest(tmp_path / 'out.json', manifest) + + assert (tmp_path / 'out.json').read_bytes() == source.read_bytes() + assert manifest.extra == {'future': 1} + assert manifest.parents.extra == {'label_source': 'x'} + + +def test_custom_table_follows_the_snapshots_record_sources( + graph, monkeypatch): + graph_, snapshot, connector, _ = graph + pquery, _, _ = _export(graph, 'PREDICT ents.LABEL FOR EACH ents.ID', + 'unlocated') + + def location(snapshot): + table = kumo.TrainingTable.from_source_table( + connector['unlocated'], TrainingTableSpec(split_col='SPLIT'), + pq=pquery, snapshot=snapshot) + return table.manifest.parents.source_table.location + + try: + assert location(graph_.snapshot(record_sources=False)) == 'unlocated' + finally: + graph_.snapshot() + assert location(snapshot) != 'unlocated' + monkeypatch.setattr(kumo.GraphSnapshot, 'manifest', property(lambda _: None)) + assert location(replace(snapshot, record_sources=None)) == 'unlocated' + + +def test_manifest_write_failure_keeps_generation_and_export( + graph, monkeypatch): + from kumo.engine.pquery import training_table + + _, _, connector, out = graph + + def fail(*args, **kwargs): + raise OSError('read-only scratch') + + monkeypatch.setattr(training_table, 'write_manifest', fail) + with pytest.warns(UserWarning, match='read-only scratch'): + _, table, exported = _export( + graph, 'PREDICT ents.LABEL FOR EACH ents.ID', 'no_manifest') + assert table.manifest is None + assert not (out / 'no_manifest' / '_provenance.json').exists() + assert exported.num_rows == sum(_seeds(table)[1].values()) + + +def test_unreadable_manifest_still_exports(graph): + graph_, snapshot, connector, out = graph + table = kumo.PredictiveQuery( + graph=graph_, query='PREDICT ents.LABEL FOR EACH ents.ID' + ).generate_training_table(snapshot) + manifest = table.staging_dir / '_provenance.json' + manifest.chmod(0) + try: + with pytest.warns(UserWarning, match='_provenance.json'): + table.export(TrainingTableOutputConfig(connector, 'unreadable')) + finally: + manifest.chmod(0o644) + assert list((out / 'unreadable').glob('*.parquet')) + assert not (out / 'unreadable' / '_provenance.json').exists() + + +def test_no_sidecar_for_warehouse_or_csv_exports(graph): + graph_, snapshot, connector, out = graph + table = kumo.PredictiveQuery( + graph=graph_, query='PREDICT ents.LABEL FOR EACH ents.ID' + ).generate_training_table(snapshot) + calls = [] + + def write_files(connector, table_name, files, write_mode, **kwargs): + calls.append(kwargs) + + warehouse = SimpleNamespace( + source_type='SNOWFLAKE', + sink_backend=SimpleNamespace(write_files=write_files)) + table.export(TrainingTableOutputConfig(warehouse, 'DB.SCHEMA.TT')) + assert calls == [{'export_id': f'{table.tt_id}:tt-export-v1'}] + table.export(TrainingTableOutputConfig(connector, 'rows.csv/')) + assert not (out / 'rows.csv' / '_provenance.json').exists() + + +def test_a_mlops_block_stores_each_table_once_under_the_root(graph, tmp_path): + graph_, snapshot, _, _ = graph + pquery = kumo.PredictiveQuery(graph=graph_, + query='PREDICT ents.LABEL FOR EACH ents.ID') + root = tmp_path / 'artifacts' + with kumo.mlops(kumo.MLOps(artifact_root=str(root))): + table = pquery.generate_training_table(snapshot) + stored = root / 'training_tables' / table.id + parts = {p.name: p.stat().st_mtime_ns for p in stored.iterdir()} + pquery.generate_training_table(snapshot) + assert '_provenance.json' in parts and '_SUCCESS' in parts + assert {p.name: p.stat().st_mtime_ns for p in stored.iterdir()} == parts + assert json.loads((stored / '_provenance.json').read_text())['id'] == ( + table.id) + off = pquery.generate_training_table( + snapshot, mlops=kumo.MLOps(artifact_root=str(tmp_path / 'other'), + export={'model'})) + assert off.id == table.id and not (tmp_path / 'other').exists() diff --git a/packages/kumo-engine/test/test_warm_start.py b/packages/kumo-engine/test/test_warm_start.py index e79130045..86d274770 100644 --- a/packages/kumo-engine/test/test_warm_start.py +++ b/packages/kumo-engine/test/test_warm_start.py @@ -182,6 +182,8 @@ def configure(self): patch.setattr(BaseLightningModule, 'configure_optimizers', configure) second = destination.fit(table, snapshot=snapshot, warm_start=source) assert restored_states and configured_states + assert second.manifest.parents.warm_start.id == source.result.id + assert second.manifest.parents.warm_start.parents.warm_start is None assert destination.model_plan == configured_plan assert _hash_bundle(durable) == original_hashes assert _hash_bundle(source_bundle) == original_hashes diff --git a/packages/kumo-snowflake/src/kumo_snowflake/identity.py b/packages/kumo-snowflake/src/kumo_snowflake/identity.py index 7b17c9a78..e62270e42 100644 --- a/packages/kumo-snowflake/src/kumo_snowflake/identity.py +++ b/packages/kumo-snowflake/src/kumo_snowflake/identity.py @@ -7,7 +7,11 @@ from typing import Any, Iterable, Mapping from kumo_core.sources.errors import PermanentSourceError, SourceNotFoundError -from kumo_core.sources.identity import epoch_second, hour_version +from kumo_core.sources.identity import ( + SourceVersion, + epoch_second, + hour_version, +) from kumo_core.sources.sql import resolve_snowflake_identifier from kumo_snowflake.sql import placeholder, quote_identifier @@ -95,7 +99,8 @@ def data_version( altered = {str(row[0]): row[1] for row in cursor.fetchall()} if table not in altered: return hour_version() - return epoch_second(pd.to_datetime(altered[table])) + return SourceVersion(epoch_second(pd.to_datetime(altered[table])), + "timestamp") class MissingColumnTypesError(PermanentSourceError): diff --git a/packages/kumo-snowflake/src/kumo_snowflake/session.py b/packages/kumo-snowflake/src/kumo_snowflake/session.py index 0e237c868..08b8c4ebf 100644 --- a/packages/kumo-snowflake/src/kumo_snowflake/session.py +++ b/packages/kumo-snowflake/src/kumo_snowflake/session.py @@ -461,6 +461,16 @@ def _principal(self, connector: SourceConnector) -> Principal: f'accepted, as {self._default_hint}.'), ) + def table_location(self, connector: SourceConnector, + table_name: str) -> str: + """``account/"DB"."SCHEMA"."TABLE"``; the account when it resolves.""" + try: + account = _snowflake_config(self._principal(connector)).account + except Exception: + account = None + qualified = self._qualified(connector, table_name) + return f'{account}/{qualified}' if account else qualified + @staticmethod def _qualified(connector: SourceConnector, table_name: str) -> str: return (f'{_quoted(resolved(_target_config(connector).database))}.' diff --git a/packages/kumo-snowflake/tests/test_identity.py b/packages/kumo-snowflake/tests/test_identity.py index f4ba9fd8d..eade78ae6 100644 --- a/packages/kumo-snowflake/tests/test_identity.py +++ b/packages/kumo-snowflake/tests/test_identity.py @@ -30,6 +30,7 @@ def stamp(recorder=None, table=TABLE, **overrides): def test_a_table_is_stamped_with_the_second_it_was_last_altered(): """A version carrying `.0` is a different path segment and a different key.""" assert stamp() == ALTERED_VERSION + assert stamp().kind == "timestamp" assert stamp(last_altered=ALTERED.replace( microsecond=900000)) == (ALTERED_VERSION) @@ -46,6 +47,7 @@ def test_a_differently_cased_table_is_not_substituted(): def test_a_view_of_either_kind_falls_back_to_the_hour_window(table_type): recorder = cursor(table_type=table_type) assert stamp(recorder) == hour_version() + assert stamp(cursor(table_type=table_type)).kind == "hour_bucket" assert recorder.matching("LAST_ALTERED") == [] From 1ee61f71e3ffd2255116c5657009cbfea9a27c50 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 28 Oct 2026 20:39:15 +0000 Subject: [PATCH 005/278] Read artifacts back through the MLOps root and rename export= to store= kumo.MLOps stands in for a connector in load_from_connector: snapshots by id, training tables and models by their path under the artifact root. MLOps(store=...) replaces export=. TrainingResult.snapshot_id names the training snapshot. DatabricksConnector without an identity is ambient. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/docs/mlops.md | 65 +++++++++++++++---- packages/kumo-engine/kumo/engine/_mlops.py | 34 +++++++--- .../kumo/engine/connector/databricks.py | 9 +-- packages/kumo-engine/kumo/engine/graph.py | 16 ++++- .../kumo/engine/pquery/training_table.py | 6 +- .../kumo/engine/trainer/trainer.py | 16 +++++ .../test/test_bundle_provenance.py | 5 +- .../test/test_databricks_connector.py | 17 ++--- packages/kumo-engine/test/test_mlops.py | 10 +-- .../kumo-engine/test/test_snapshot_export.py | 9 ++- .../test/test_training_table_export.py | 8 ++- 11 files changed, 146 insertions(+), 49 deletions(-) diff --git a/packages/kumo-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index 77bc7f609..71cc3a81c 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -1,30 +1,69 @@ # MLOps with Kumo Engine -## Turning MLOps on +## Quickstart -`kumo.MLOps` says where artifacts go. Use it as a block, which notebooks and scheduled scripts alike can do, or pass it as `mlops=` on a single call: +With a `graph` and a `pquery` (see the README), turn MLOps on with an artifact root and every snapshot, training table and model is stored under it: ```python import kumo.engine as kumo -cfg = kumo.MLOps(artifact_root="s3://bucket/kumo") # or (connector, path), e.g. a volume or stage +cfg = kumo.MLOps("artifacts") # or s3://..., gs://..., /Volumes/..., (connector, "@DB.SCHEMA.STAGE/kumo") with kumo.mlops(cfg): snapshot = graph.snapshot() + result = kumo.Trainer(pquery.suggest_model_plan(snapshot)).fit( + pquery.generate_training_table(snapshot), snapshot=snapshot) +# Later, on any machine that reaches the root: +snapshot = kumo.GraphSnapshot.load_from_connector(cfg, id=result.snapshot_id) +trainer = kumo.Trainer.load_from_connector(cfg, path=f"models/{result.id}") ``` -- **`artifact_root`** is required. It is a URI (`s3://`, `gs://`, `file://`, or a local path, including `/Volumes/...` on a Databricks cluster) or a `(connector, path)` pair for a store a URI cannot reach. -- **Volumes and stages:** `(DatabricksConnector(...), "/Volumes////kumo")` and `(SnowflakeConnector(...), "@DB.SCHEMA.STAGE/kumo")` work as roots once the installed kumo-databricks or kumo-snowflake sink holds files at those paths; with an older plugin, building `kumo.MLOps` raises. -- **`export`** names the artifact kinds stored under the root when they are produced: `snapshot`, `training_table` and `model` (default: all). +## How it fits together + +| artifact | id | stored at (MLOps on) | +|---|---|---| +| snapshot | `snapshot.id` | `/snapshots///` | +| training table | `table.id` | `/training_tables//` | +| model | `result.id` (the bundle id) | `/models//` | + +Each carries a `_provenance.json` manifest that embeds its parents' manifests (see [Lineage](#lineage)). + +## Turning MLOps on + +`kumo.MLOps` says where artifacts go. Use it as a block, or pass it as `mlops=` on a single call: + +```python +cfg = kumo.MLOps("s3://bucket/kumo") +with kumo.mlops(cfg): + snapshot = graph.snapshot() +trainer.fit(table, snapshot=snapshot, mlops=cfg) +``` + +- **`artifact_root`** is required: a URI (`s3://`, `gs://`, `file://`, or a local path, including `/Volumes/...` on a Databricks cluster) or a `(connector, path)` pair for a store a URI cannot reach, such as `(DatabricksConnector(...), "/Volumes////kumo")` or `(SnowflakeConnector(...), "@DB.SCHEMA.STAGE/kumo")`. A connector whose sink cannot hold files there raises when `kumo.MLOps` is built. +- **`store`** names the artifact kinds stored when they are produced: `snapshot`, `training_table` and `model` (default: all). - **Precedence:** a call's `mlops=` replaces the enclosing block, and an inner block replaces an outer one. A block covers the calling thread or asyncio task. -- **Off by default.** Without a block or `mlops=`, every call behaves as it always has. Plain exports (`export(config)`, `export_config=`, `output_config=`) never write to the artifact root. A path export that finds a complete copy of the same artifact at its destination reuses it instead of uploading again (logged at info): a snapshot with the same bytes, a training table or prediction directory with `_SUCCESS` and the same `_KUMO_EXPORT` identity, or a model bundle with the same files and `_provenance.json`. +- **Off by default.** Without a block or `mlops=`, every call behaves as it always has, and plain exports (`export(config)`, `export_config=`) never write to the artifact root. + +## Exporting and loading artifacts + +Every artifact has one way out, `.export(OutputConfig(...))`, and one way back, `.load_from_connector(...)`. Pass the `kumo.MLOps` in place of the connector to read from the artifact root, with paths relative to it: + +| Artifact | Export | Load | +|---|---|---| +| Snapshot | `snapshot.export(SnapshotOutputConfig(connector, path))`, returns the location | `GraphSnapshot.load_from_connector(connector, path=location)`, or `(cfg, id=snapshot_id)` | +| Training table | `table.export(TrainingTableOutputConfig(connector, name))` | `TrainingTable.load_from_connector(connector, path=name, pq=pq, snapshot=snapshot)`, or `(cfg, path=f"training_tables/{id}", ...)` | +| Model | `result.export(ModelOutputConfig(connector, path=... or location=...))` | `Trainer.load_from_connector(connector, path=... or location=...)`, or `(cfg, path=f"models/{id}")` | + +`fit(export_config=...)` is shorthand for fitting, then exporting. Loading needs a file destination (S3, GCS, local files, a volume or stage). Exporting again to a destination that already holds a complete copy of the same artifact reuses it (logged at info) instead of uploading again. + +## Reference -## Lineage +### Lineage Every artifact carries a `_provenance.json` manifest, metadata only. The leading underscore keeps Spark and Hive Parquet readers from treating it as data. A graph snapshot's manifest, next to its index, records each source table's connector, location and version, plus the engine build that wrote it. `graph.snapshot(record_sources=False)` keeps only connector and table names, plus a SHA-256 of each full location so a source moved elsewhere is still noticed, and keeps Custom SQL only as a hash. -### Training tables +#### Training tables A generated or caller-supplied training table carries `_provenance.json` in its staging directory: - the query, the effective generation plan and the task type; @@ -36,7 +75,7 @@ A generated or caller-supplied training table carries `_provenance.json` in its With MLOps on, each table is stored once under `/training_tables//`. -### Models +#### Models Every fitted bundle carries `_provenance.json` (`ModelManifest`): - its `id`, the bundle id (a content hash that leaves the manifest file out); @@ -73,11 +112,11 @@ With MLOps on (`graph.snapshot(mlops=...)` or a `kumo.mlops` block), each snapsh ### Scoring without the graph -A scoring job does not need the training job's `Graph` code. It loads the snapshot the training job exported and the model, then scores: +A scoring job does not need the training job's `Graph` code. It loads a snapshot and the model from the artifact root, then scores: ```python -snapshot = GraphSnapshot.load_from_connector(store, path=location) -trainer = Trainer.load_from_connector(store, path="models/") +snapshot = GraphSnapshot.load_from_connector(cfg, id=snapshot_id) +trainer = Trainer.load_from_connector(cfg, path=f"models/{bundle_id}") result = trainer.predict(snapshot, binary_classification_threshold=0.5) ``` diff --git a/packages/kumo-engine/kumo/engine/_mlops.py b/packages/kumo-engine/kumo/engine/_mlops.py index 3c7b16463..527508eb9 100644 --- a/packages/kumo-engine/kumo/engine/_mlops.py +++ b/packages/kumo-engine/kumo/engine/_mlops.py @@ -37,15 +37,18 @@ class MLOps: (``/Volumes/...`` on a Databricks cluster), or a ``(connector, path)`` pair for a store a URI cannot reach, such as a volume through a ``DatabricksConnector`` or a Snowflake stage. - export: the artifact kinds stored under the root when produced + store: the artifact kinds stored under the root when produced (``snapshot``, ``training_table``, ``model``); ``True`` for all, ``False`` for none. + + In ``load_from_connector`` calls, an ``MLOps`` stands in for the + connector, with paths relative to its artifact root. """ artifact_root: Union[str, Tuple[Any, str]] - export: Union[bool, Iterable[str]] = field(default=True, kw_only=True) + store: Union[bool, Iterable[str]] = field(default=True, kw_only=True) root: _Root = field(init=False, repr=False, compare=False) - exports: FrozenSet[str] = field(init=False, repr=False, compare=False) + kinds: FrozenSet[str] = field(init=False, repr=False, compare=False) def __post_init__(self) -> None: from kumo.engine.connector import ( @@ -65,19 +68,18 @@ def __post_init__(self) -> None: f'{type(connector).__name__} cannot hold files at ' f'{path!r}; artifact_root needs a file location.') root = _Root(connector, path, path) - kinds = (frozenset(KINDS) if self.export is True else - frozenset() if self.export is False else frozenset( - self.export)) + kinds = (frozenset(KINDS) if self.store is True else + frozenset() if self.store is False else frozenset(self.store)) unknown = sorted(kinds - set(KINDS)) if unknown: - raise ValueError(f'Unknown export kinds {unknown}; choose from ' + raise ValueError(f'Unknown store kinds {unknown}; choose from ' f'{list(KINDS)}.') object.__setattr__(self, 'root', root) - object.__setattr__(self, 'exports', kinds) + object.__setattr__(self, 'kinds', kinds) def wants(self, kind: str) -> bool: r"""Whether artifacts of ``kind`` are stored under the root.""" - return kind in self.exports + return kind in self.kinds _CURRENT: contextvars.ContextVar[Optional[MLOps]] = contextvars.ContextVar( @@ -103,3 +105,17 @@ def resolve(settings: Optional[MLOps] = None) -> Optional[MLOps]: r"""The settings a call runs with: its ``mlops=`` argument, else the enclosing block's, else None (off).""" return settings if settings is not None else _CURRENT.get() + + +def on_root(source: Any, path: str) -> Tuple[Any, str]: + r"""The connector and path a ``load_from_connector(source, path=...)`` + reads: an `MLOps` ``source`` reads ``path`` under its artifact root.""" + if isinstance(source, MLOps): + return source.root.connector, source.root.location(path.lstrip('/')) + return source, path + + +def snapshot_path(snapshot: Any) -> str: + r"""Where a snapshot, or a snapshot id, is stored under a root.""" + snapshot_id = snapshot if isinstance(snapshot, str) else snapshot.id + return 'snapshots/' + snapshot_id.replace('@', '/') diff --git a/packages/kumo-engine/kumo/engine/connector/databricks.py b/packages/kumo-engine/kumo/engine/connector/databricks.py index b7f52f892..4bc380a66 100644 --- a/packages/kumo-engine/kumo/engine/connector/databricks.py +++ b/packages/kumo-engine/kumo/engine/connector/databricks.py @@ -36,8 +36,9 @@ class DatabricksConnector(Connector): Args: catalog: Databricks Unity Catalog name. schema_name: Schema within `catalog`. - identity: The `DatabricksIdentity` used to authenticate. See - `from_notebook` for the credential-free notebook case. + identity: The `DatabricksIdentity` used to authenticate. `None` + uses the Databricks SDK's ambient auth (a notebook, a job, or + `databricks auth login`). warehouse_id: Databricks SQL warehouse ID. Mutually exclusive with `cluster_id`. cluster_id: Databricks all-purpose cluster ID. Mutually exclusive @@ -53,7 +54,7 @@ def __init__( *, catalog: str, schema_name: str, - identity: DatabricksIdentity, + identity: Optional[DatabricksIdentity] = None, warehouse_id: Optional[str] = None, cluster_id: Optional[str] = None, staging_volume_path: Optional[str] = None, @@ -61,7 +62,7 @@ def __init__( ) -> None: if require_staged and not staging_volume_path: raise ValueError("require_staged needs a staging_volume_path") - super().__init__(identity=identity) + super().__init__(identity=identity or DatabricksIdentity()) self._config = DatabricksConnectorResourceConfig( catalog=catalog, schema_name=schema_name, diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index 04df470a2..e6c991133 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -406,19 +406,33 @@ def export(self, output_config: 'SnapshotOutputConfig') -> str: return export(self, output_config.connector, output_config.path) @classmethod - def load_from_connector(cls, connector: Connector, *, path: str, + def load_from_connector(cls, connector: Connector, *, + path: Optional[str] = None, + id: Optional[str] = None, graph: Optional['Graph'] = None) -> 'GraphSnapshot': r"""Copy the snapshot exported at `path` into the local store. Args: + connector: a connector, or a `kumo.MLOps` for its artifact root. path: the location `export` returned. + id: instead of `path`, the id of a snapshot stored under the + artifact root of the `kumo.MLOps` passed as `connector`. graph: if given, warn when its source tables changed since the snapshot was built (see `source_changes`). """ from kumo_core.diskgraph.snapshot import DiskgraphSnapshotHandle + from kumo.engine._mlops import MLOps, on_root, snapshot_path from kumo.engine._snapshot_io import holds_export, load_into, read + if (path is None) == (id is None): + raise ValueError('Pass exactly one of path or id.') + if id is not None: + if not isinstance(connector, MLOps): + raise TypeError('id= loads from a kumo.MLOps artifact root; ' + 'pass path= with a connector.') + path = snapshot_path(id) + connector, path = on_root(connector, path) root, manifest = read(connector, path) graph_key = manifest['graph_key'] if graph is not None and graph.id != graph_key: diff --git a/packages/kumo-engine/kumo/engine/pquery/training_table.py b/packages/kumo-engine/kumo/engine/pquery/training_table.py index d8e88689e..417295022 100644 --- a/packages/kumo-engine/kumo/engine/pquery/training_table.py +++ b/packages/kumo-engine/kumo/engine/pquery/training_table.py @@ -125,7 +125,11 @@ def load_from_connector(cls, connector: Any, *, path: str, snapshot: GraphSnapshot) -> "TrainingTable": r"""Materialize a table written by :meth:`export` (a directory with a ``SPLIT`` column) against ``snapshot``, as :meth:`from_source_table` - with ``split_col='SPLIT'``.""" + with ``split_col='SPLIT'``. A `kumo.MLOps` ``connector`` reads + ``path`` under its artifact root (``training_tables/``).""" + from kumo.engine._mlops import on_root + + connector, path = on_root(connector, path) return cls.from_source_table(connector[path], TrainingTableSpec(split_col='SPLIT'), pq=pq, snapshot=snapshot) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 45be9d35f..d59169e92 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -65,6 +65,7 @@ from kumo.engine.trainer._spec import ( HOLDOUT_META_ROLE, HOLDOUT_ROLE, + SPEC_FILE, listed_file, read_bundle_artifact_paths, read_training_spec, @@ -116,6 +117,13 @@ def manifest(self) -> Optional[ModelManifest]: return read_provenance(self.bundle_dir) + @property + def snapshot_id(self) -> Optional[str]: + """The id of the snapshot the model was trained on.""" + if not (self.bundle_dir / SPEC_FILE).is_file(): + return None + return read_training_spec(self.bundle_dir).snapshot_id + def predict(self, *args: Any, **kwargs: Any) -> 'PredictionResult': """Score with this model, as :meth:`Trainer.predict`.""" return self._model('predict').predict(*args, **kwargs) @@ -1043,7 +1051,15 @@ def load_from_connector( Hashes detect corruption, not publisher authentication. Load only models from trusted publishers; checkpoint deserialization happens at predict. ``accelerator`` selects the device for this instance (default: ``auto``). + A `kumo.MLOps` ``model_registry`` reads ``path`` under its artifact + root (``models/``). """ + from kumo.engine._mlops import MLOps, on_root + + if isinstance(model_registry, MLOps) and path is None: + raise TypeError('Loading from a kumo.MLOps takes path=.') + if path is not None: + model_registry, path = on_root(model_registry, path) address = ModelOutputConfig(model_registry, path=path, location=location) bundle_dir = trainer_scratch_dir(new_run_id()) / 'fetched-model' diff --git a/packages/kumo-engine/test/test_bundle_provenance.py b/packages/kumo-engine/test/test_bundle_provenance.py index dab92e4e7..854e6e0dc 100644 --- a/packages/kumo-engine/test/test_bundle_provenance.py +++ b/packages/kumo-engine/test/test_bundle_provenance.py @@ -138,12 +138,15 @@ def test_a_mlops_block_stores_each_model_once(tmp_path): settings=cfg) stored = tmp_path / 'artifacts' / 'models' / result.id assert kumo.Trainer.load(stored).result.id == result.id + loaded = kumo.Trainer.load_from_connector(cfg, path=f'models/{result.id}') + assert loaded.result.id == result.id + assert loaded.result.snapshot_id == result.snapshot_id (stored / 'model' / 'config.json').write_text('{"damaged": 1}') again = _finish_stub_fit(tmp_path / 'again', kumo.Trainer(ModelPlan()), settings=cfg) assert again.id == result.id assert (stored / 'model' / 'config.json').read_text() == '{}' - off = kumo.MLOps(artifact_root=str(tmp_path / 'off'), export={'snapshot'}) + off = kumo.MLOps(artifact_root=str(tmp_path / 'off'), store={'snapshot'}) _finish_stub_fit(tmp_path / 'off-bundle', kumo.Trainer(ModelPlan()), settings=off) assert not (tmp_path / 'off').exists() diff --git a/packages/kumo-engine/test/test_databricks_connector.py b/packages/kumo-engine/test/test_databricks_connector.py index 9a7b08fea..2abbd539b 100644 --- a/packages/kumo-engine/test/test_databricks_connector.py +++ b/packages/kumo-engine/test/test_databricks_connector.py @@ -58,17 +58,12 @@ def test_strict_staging_requires_a_volume(staging_volume_path): ) -def test_identity_is_required(): - with pytest.raises(TypeError, match='identity'): - DatabricksConnector( - catalog='main', - schema_name='analytics', - warehouse_id='warehouse-id', - ) - - -def test_from_notebook_builds_an_ambient_identity(): - connector = DatabricksConnector.from_notebook( +@pytest.mark.parametrize('build', [ + DatabricksConnector.from_notebook, + DatabricksConnector, +]) +def test_without_an_identity_the_connector_is_ambient(build): + connector = build( catalog='main', schema_name='analytics', warehouse_id='warehouse-id', diff --git a/packages/kumo-engine/test/test_mlops.py b/packages/kumo-engine/test/test_mlops.py index 6384faa72..22e3d80eb 100644 --- a/packages/kumo-engine/test/test_mlops.py +++ b/packages/kumo-engine/test/test_mlops.py @@ -27,8 +27,8 @@ def test_mlops_needs_an_artifact_root_it_can_write_files_to(tmp_path): holds_files=lambda c, p: False, filesystem=None)) with pytest.raises(ValueError, match='file location'): kumo.MLOps(artifact_root=(table_only, 'DB.SCHEMA')) - with pytest.raises(ValueError, match='Unknown export kinds'): - kumo.MLOps(artifact_root=str(tmp_path), export={'models'}) + with pytest.raises(ValueError, match='Unknown store kinds'): + kumo.MLOps(artifact_root=str(tmp_path), store={'models'}) def test_a_root_is_a_uri_or_a_connector_and_path(tmp_path): @@ -36,15 +36,15 @@ def test_a_root_is_a_uri_or_a_connector_and_path(tmp_path): assert uri.root.location('models/m') == 'models/m' assert uri.root.uri('models/m') == f'{tmp_path}/models/m' assert isinstance(uri.root.connector, kumo.FileConnector) - assert uri.exports == {'snapshot', 'training_table', 'model'} + assert uri.kinds == {'snapshot', 'training_table', 'model'} connector = kumo.FileConnector(root_dir=str(tmp_path)) pair = kumo.MLOps(artifact_root=(connector, 'volume/kumo/'), - export={'model'}) + store={'model'}) assert pair.root.connector is connector assert pair.root.location('models/m') == 'volume/kumo/models/m' assert pair.wants('model') and not pair.wants('snapshot') - assert not kumo.MLOps(artifact_root=str(tmp_path), export=False).exports + assert not kumo.MLOps(artifact_root=str(tmp_path), store=False).kinds def test_a_call_argument_beats_the_nearest_block(tmp_path): diff --git a/packages/kumo-engine/test/test_snapshot_export.py b/packages/kumo-engine/test/test_snapshot_export.py index e0f15057f..41a1556af 100644 --- a/packages/kumo-engine/test/test_snapshot_export.py +++ b/packages/kumo-engine/test/test_snapshot_export.py @@ -770,7 +770,8 @@ def _remove(self, path, *, tree): def test_a_mlops_block_stores_each_snapshot_once(tmp_path, scratch, graph, monkeypatch): root = tmp_path / 'artifacts' - with kumo.mlops(kumo.MLOps(artifact_root=str(root))): + cfg = kumo.MLOps(artifact_root=str(root)) + with kumo.mlops(cfg): snapshot = graph.snapshot() monkeypatch.setattr(snapshot_io, 'export', Mock(side_effect=AssertionError('re-exported'))) @@ -781,6 +782,10 @@ def test_a_mlops_block_stores_each_snapshot_once(tmp_path, scratch, graph, assert (_generation(stored) / '_provenance.json').is_file() assert _load(kumo.FileConnector(root_dir=str(root)), f'snapshots/{graph_id}/{data_version}') == snapshot + assert GraphSnapshot.load_from_connector(cfg, id=snapshot.id) == snapshot + with pytest.raises(TypeError, match='kumo.MLOps'): + GraphSnapshot.load_from_connector(kumo.FileConnector(root_dir='.'), + id=snapshot.id) graph.snapshot(mlops=kumo.MLOps(artifact_root=str(tmp_path / 'off'), - export={'model'})) + store={'model'})) assert not (tmp_path / 'off').exists() diff --git a/packages/kumo-engine/test/test_training_table_export.py b/packages/kumo-engine/test/test_training_table_export.py index fd97a20f1..13f40ebd4 100644 --- a/packages/kumo-engine/test/test_training_table_export.py +++ b/packages/kumo-engine/test/test_training_table_export.py @@ -446,7 +446,8 @@ def test_a_mlops_block_stores_each_table_once_under_the_root(graph, tmp_path): pquery = kumo.PredictiveQuery(graph=graph_, query='PREDICT ents.LABEL FOR EACH ents.ID') root = tmp_path / 'artifacts' - with kumo.mlops(kumo.MLOps(artifact_root=str(root))): + cfg = kumo.MLOps(artifact_root=str(root)) + with kumo.mlops(cfg): table = pquery.generate_training_table(snapshot) stored = root / 'training_tables' / table.id parts = {p.name: p.stat().st_mtime_ns for p in stored.iterdir()} @@ -455,7 +456,10 @@ def test_a_mlops_block_stores_each_table_once_under_the_root(graph, tmp_path): assert {p.name: p.stat().st_mtime_ns for p in stored.iterdir()} == parts assert json.loads((stored / '_provenance.json').read_text())['id'] == ( table.id) + reloaded = kumo.TrainingTable.load_from_connector( + cfg, path=f'training_tables/{table.id}', pq=pquery, snapshot=snapshot) + assert _seeds(reloaded)[1] == _seeds(table)[1] off = pquery.generate_training_table( snapshot, mlops=kumo.MLOps(artifact_root=str(tmp_path / 'other'), - export={'model'})) + store={'model'})) assert off.id == table.id and not (tmp_path / 'other').exists() From 5766eb9144d5a5b2195d3e8157576d879ea1f854 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 28 Oct 2026 21:36:05 +0000 Subject: [PATCH 006/278] Fail a fit whose test split is empty, and a stage root given without its connector A query window that leaves no test rows raised an opaque evaluation error after training; fit now refuses before training and says to shorten the window or add data. kumo.MLOps('@DB.SCHEMA.STAGE/...') was read as a local directory; it now asks for the (SnowflakeConnector, path) pair. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_mlops.py | 4 ++++ packages/kumo-engine/kumo/engine/trainer/trainer.py | 4 ++++ packages/kumo-engine/test/test_holdout_df.py | 10 ++++++++++ packages/kumo-engine/test/test_mlops.py | 2 ++ 4 files changed, 20 insertions(+) diff --git a/packages/kumo-engine/kumo/engine/_mlops.py b/packages/kumo-engine/kumo/engine/_mlops.py index 527508eb9..5f8b57b60 100644 --- a/packages/kumo-engine/kumo/engine/_mlops.py +++ b/packages/kumo-engine/kumo/engine/_mlops.py @@ -59,6 +59,10 @@ def __post_init__(self) -> None: if not self.artifact_root: raise ValueError('MLOps needs an artifact_root.') if isinstance(self.artifact_root, str): + if self.artifact_root.startswith('@'): + raise ValueError( + 'A stage artifact_root needs its connector: pass ' + f'(SnowflakeConnector(...), {self.artifact_root!r}).') root = _Root(connector_for_root(self.artifact_root), '', self.artifact_root) else: diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index d59169e92..4e0708b08 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -601,6 +601,10 @@ def fit( tt_id=ingest_tt_id, ) del task_dataset + if (materialized.task.split.stage_counts.get('test') == 0 + and not materialized.task.task_type.is_link_pred): + raise ValueError('The test split is empty: shorten the query ' + 'window or add data.') training_input = ResolvedTrainingInput( task=materialized.task, model_plan=materialized.model_plan, diff --git a/packages/kumo-engine/test/test_holdout_df.py b/packages/kumo-engine/test/test_holdout_df.py index c70110813..290cce9cb 100644 --- a/packages/kumo-engine/test/test_holdout_df.py +++ b/packages/kumo-engine/test/test_holdout_df.py @@ -194,3 +194,13 @@ def test_unreadable_size_record_does_not_fail_fit(tmp_path, caplog, content): assert truncated is None assert len(pd.read_parquet(path)) == 2 assert 'truncation is unknown' in caplog.text + + +def test_an_empty_test_split_fails_the_fit_with_a_clear_message(graph): + graph, snapshot = graph + pquery = kumo.PredictiveQuery( + graph=graph, + query='PREDICT COUNT(events.*, 0, 30, days) > 15 FOR EACH ents.ID') + with pytest.raises(ValueError, match='test split is empty'): + kumo.Trainer(pquery.suggest_model_plan(snapshot)).fit( + pquery.generate_training_table(snapshot), snapshot=snapshot) diff --git a/packages/kumo-engine/test/test_mlops.py b/packages/kumo-engine/test/test_mlops.py index 22e3d80eb..179606015 100644 --- a/packages/kumo-engine/test/test_mlops.py +++ b/packages/kumo-engine/test/test_mlops.py @@ -27,6 +27,8 @@ def test_mlops_needs_an_artifact_root_it_can_write_files_to(tmp_path): holds_files=lambda c, p: False, filesystem=None)) with pytest.raises(ValueError, match='file location'): kumo.MLOps(artifact_root=(table_only, 'DB.SCHEMA')) + with pytest.raises(ValueError, match='needs its connector'): + kumo.MLOps(artifact_root='@DB.SCHEMA.STAGE/kumo') with pytest.raises(ValueError, match='Unknown store kinds'): kumo.MLOps(artifact_root=str(tmp_path), store={'models'}) From a5dd3bf3e589d5397215cb2f4eb205b504e7d178 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 28 Oct 2026 22:48:15 +0000 Subject: [PATCH 007/278] Give the fit test fakes a test split The empty-test-split check reads the materialized task's stage counts; the fakes now carry one. Co-Authored-By: Claude Opus 5.5 --- .../kumo-engine/test/test_semantic_carry_forward.py | 12 +++++++----- .../kumo-engine/test/test_training_explanations.py | 4 +++- 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/packages/kumo-engine/test/test_semantic_carry_forward.py b/packages/kumo-engine/test/test_semantic_carry_forward.py index 9481d14d1..fc5788a9b 100644 --- a/packages/kumo-engine/test/test_semantic_carry_forward.py +++ b/packages/kumo-engine/test/test_semantic_carry_forward.py @@ -26,6 +26,8 @@ ) from kumo.engine.trainer.trainer import Trainer, _stage_model_bundle +_SPLIT = SimpleNamespace(stage_counts={'test': 1}) + @dataclass(frozen=True) class _TaskDataset: @@ -395,7 +397,7 @@ def fake_build(**kwargs): def fake_materialize(request, task_dataset, generated, **kwargs): calls.append((request, task_dataset, generated, kwargs.get('tt_id'))) - return SimpleNamespace(task=Mock(spec=ResolvedNumericalTask), + return SimpleNamespace(task=Mock(spec=ResolvedNumericalTask, split=_SPLIT), model_plan=CoreModelPlan(), bundle_dir=str(tmp_path)) @@ -528,7 +530,7 @@ def stats(index_path, table_name, mode): def fake_materialize(request, task_dataset, generated, **kwargs): calls.append((request, task_dataset, generated)) - return SimpleNamespace(task=Mock(spec=ResolvedNumericalTask), + return SimpleNamespace(task=Mock(spec=ResolvedNumericalTask, split=_SPLIT), model_plan=CoreModelPlan(), bundle_dir=str(tmp_path)) @@ -604,7 +606,7 @@ def test_fit_uses_selected_graph_task_semantics_without_regenerating_rows( def fake_materialize(request, task_dataset, generated, **kwargs): captured.update(request=request, task_dataset=task_dataset, generated=generated) - return SimpleNamespace(task=Mock(spec=ResolvedNumericalTask), + return SimpleNamespace(task=Mock(spec=ResolvedNumericalTask, split=_SPLIT), model_plan=CoreModelPlan(), bundle_dir=str(tmp_path)) @@ -780,7 +782,7 @@ def fake_build(**kwargs): def fake_materialize(request, task_dataset, generated, **kwargs): tt_ids.append(kwargs.get('tt_id')) task_datasets.append(task_dataset) - return SimpleNamespace(task=Mock(spec=ResolvedNumericalTask), + return SimpleNamespace(task=Mock(spec=ResolvedNumericalTask, split=_SPLIT), model_plan=CoreModelPlan(), bundle_dir=str(tmp_path)) @@ -864,7 +866,7 @@ def fake_build(**kwargs): def fake_materialize(request, task_dataset, generated, **kwargs): with records_lock: records.append((kwargs['tt_id'], task_dataset, request.run_id)) - return SimpleNamespace(task=Mock(spec=ResolvedNumericalTask), + return SimpleNamespace(task=Mock(spec=ResolvedNumericalTask, split=_SPLIT), model_plan=CoreModelPlan(), bundle_dir=str(tmp_path)) diff --git a/packages/kumo-engine/test/test_training_explanations.py b/packages/kumo-engine/test/test_training_explanations.py index 399dd1eed..7b6ef408d 100644 --- a/packages/kumo-engine/test/test_training_explanations.py +++ b/packages/kumo-engine/test/test_training_explanations.py @@ -14,6 +14,8 @@ from kumo.engine import Trainer, TrainingTable +_SPLIT = SimpleNamespace(stage_counts={'test': 1}) + @pytest.mark.parametrize('disabled,explain_fails,aggregate_fails', [ (False, False, False), @@ -84,7 +86,7 @@ def train(training_input, **kwargs): monkeypatch.setattr( 'kumo_core.training.workflow.materialize_generated_training_task', lambda *args, **kwargs: SimpleNamespace( - task=Mock(spec=ResolvedNumericalTask), model_plan=core_plan, + task=Mock(spec=ResolvedNumericalTask, split=_SPLIT), model_plan=core_plan, bundle_dir=str(tmp_path))) monkeypatch.setattr( 'kumo_core.training.orchestration.run_callable_training', train) From 30eb44896de2cb1cfd305a301cf01bff57ecb0eb Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 07:39:29 +0000 Subject: [PATCH 008/278] Warn when the scored graph differs only in table connectors or locations Scoring tables read from another location (a dev graph scored on prod tables, or a renamed copy) changes the graph id, but predict and GraphSnapshot.load_from_connector only logged that at info level, so with default logging nobody saw it. It is now a UserWarning like the other graph-mismatch cases, saying the schema is the same. Verified: test_graph_mismatch, test_snapshot_export. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/docs/mlops.md | 4 ++-- packages/kumo-engine/kumo/engine/graph.py | 7 +++---- packages/kumo-engine/test/test_graph_mismatch.py | 10 +++------- packages/kumo-engine/test/test_snapshot_export.py | 10 ++++------ 4 files changed, 12 insertions(+), 19 deletions(-) diff --git a/packages/kumo-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index 71cc3a81c..a22ce00c6 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -103,7 +103,7 @@ snapshot = GraphSnapshot.load_from_connector(S3Connector(), path=location) - **Layout:** `export` writes `///`, the two halves of `snapshot.id`, and returns that location. It holds `kumo_snapshot.json` and one directory per export generation (`index/` and `_provenance.json`). - **Transfer:** files move in parallel with a CRC32 per range. Each export uploads into a new generation and commits by writing `kumo_snapshot.json` last, naming that generation. Until then the previous export stays loadable, and concurrent exports never mix files: the last commit wins whole. Superseded generations, and abandoned ones over a day old, are removed best-effort. A load either completes or leaves nothing; if a newer export commits while it downloads, it switches to that one. - **Reuse:** loading again reuses the local copy only if it holds this export's bytes (the same range checksums) and passes DiskGraph's integrity check; otherwise it is downloaded again. Exporting again skips the upload when the destination already holds a complete export (its `kumo_snapshot.json` and `_DONE`) with the local index's bytes, and logs that it reused it. -- **Staleness:** `snapshot.source_changes(graph)` reports, per table of `graph`, whether its source is `unchanged`, `changed` or `unknown` since the snapshot was built. `load_from_connector(..., graph=graph)` warns when the snapshot was taken of another graph, naming the schema changes (tables, columns, types, keys, edges), or when any of its tables changed. A graph id also covers each table's connector and location, so a graph that differs only there is logged at info level, not warned about. +- **Staleness:** `snapshot.source_changes(graph)` reports, per table of `graph`, whether its source is `unchanged`, `changed` or `unknown` since the snapshot was built. `load_from_connector(..., graph=graph)` warns when the snapshot was taken of another graph, naming the schema changes (tables, columns, types, keys, edges), or when any of its tables changed. A graph id also covers each table's connector and location, so a graph that differs only there also warns, saying the schema is the same. - **Manifest:** the snapshot's `_provenance.json` travels with the export. - **Trust:** load only from paths whose writers you trust. Checksums detect corruption, not tampering. - **Destinations:** any destination whose sink holds files at the path works: S3, GCS, local files, or a volume or stage once those sinks are installed. A warehouse connector raises. @@ -122,6 +122,6 @@ result = trainer.predict(snapshot, binary_classification_threshold=0.5) - `predict(snapshot)` reuses the trained query and table-generation plan, and scores its population at `prediction_time`. - `entities=` scores given rows instead. -- It warns when the snapshot was taken of a different graph than the one the model was trained on, naming the schema changes, for example `Schema changes: events edges.` A graph that differs only in table connectors or locations (a dev graph and a prod graph of the same schema) is logged at info level instead. The model records its training graph from this release on; for an older model the warning cannot say what differs. +- It warns when the snapshot was taken of a different graph than the one the model was trained on, naming the schema changes, for example `Schema changes: events edges.` A graph that differs only in table connectors or locations (a dev graph and a prod graph of the same schema) warns that the schema is the same. The model records its training graph from this release on; for an older model the warning cannot say what differs. - `predict(prediction_table, snapshot=...)` scores a prediction table generated beforehand, as before. - `fit` results score too: `result.predict(snapshot, ...)`. diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index e6c991133..ea5c66bed 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -60,7 +60,6 @@ """ import json -import logging import threading import warnings from dataclasses import dataclass, field, replace @@ -98,7 +97,6 @@ _warned_scratch: set[Path] = set() _warned_backfill: set[Path] = set() -logger = logging.getLogger(__name__) def _warn_if_network_scratch(root: Path) -> None: @@ -265,8 +263,9 @@ def report_other_graph(message: str, before: Optional[Mapping[str, Any]], warnings.warn(f"{message} Schema changes: {'; '.join(changes)}.", UserWarning, stacklevel=stacklevel + 1) else: - logger.info('%s The schema is the same; only table connectors or ' - 'locations differ.', message) + warnings.warn(f"{message} The schema is the same; only table " + "connectors or locations differ.", UserWarning, + stacklevel=stacklevel + 1) def _store() -> Any: diff --git a/packages/kumo-engine/test/test_graph_mismatch.py b/packages/kumo-engine/test/test_graph_mismatch.py index 6714c5ba9..d36e4a3db 100644 --- a/packages/kumo-engine/test/test_graph_mismatch.py +++ b/packages/kumo-engine/test/test_graph_mismatch.py @@ -3,8 +3,6 @@ # SPDX-License-Identifier: Apache-2.0 import copy -import logging -import warnings from dataclasses import replace import pytest @@ -41,16 +39,14 @@ def stop(*args, **kwargs): def test_predict_says_what_differs_from_the_training_graph( - fitted, monkeypatch, caplog, tmp_path): + fitted, monkeypatch, tmp_path): trainer, _, snapshot = fitted assert kumo.Trainer.load(trainer.result.bundle_dir)._graph_config == ( snapshot.graph_config) - with caplog.at_level(logging.INFO, logger='kumo.engine.graph'), \ - warnings.catch_warnings(): - warnings.filterwarnings('error', message='Snapshot .* was taken') + with pytest.warns(UserWarning, match='The schema is the same; only ' + 'table connectors or locations differ'): _predict(trainer, _other(snapshot, _moved), monkeypatch) - assert 'only table connectors or locations differ' in caplog.text def dropped(tables): _moved(tables) diff --git a/packages/kumo-engine/test/test_snapshot_export.py b/packages/kumo-engine/test/test_snapshot_export.py index 41a1556af..4452bc06d 100644 --- a/packages/kumo-engine/test/test_snapshot_export.py +++ b/packages/kumo-engine/test/test_snapshot_export.py @@ -484,8 +484,8 @@ def test_load_rejects_unknown_formats_and_accepts_unknown_fields( assert loaded.id == snapshot.id -def test_load_only_logs_a_graph_that_moved_with_the_same_schema( - tmp_path, scratch, graph, exports, caplog): +def test_load_warns_about_a_graph_that_moved_with_the_same_schema( + tmp_path, scratch, graph, exports): location, _ = _export(graph.snapshot(), exports, tmp_path) shutil.copytree(tmp_path / 'data', tmp_path / 'prod') prod = kumo.FileConnector(root_dir=str(tmp_path / 'prod')) @@ -496,11 +496,9 @@ def test_load_only_logs_a_graph_that_moved_with_the_same_schema( }, edges=graph.edges) assert moved.id != graph.id - with caplog.at_level('INFO', logger='kumo.engine.graph'), \ - warnings.catch_warnings(): - warnings.filterwarnings('error', message='.*another graph') + with pytest.warns(UserWarning, match=r'another graph\. The schema is ' + r'the same; only table connectors or locations differ'): _load(exports, location, graph=moved) - assert 'only table connectors or locations differ' in caplog.text @pytest.mark.parametrize('recorded', ['newer', 'missing']) From 32020c5e5101a96ba13b154d88b80eed44ae7b86 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 09:58:25 +0000 Subject: [PATCH 009/278] Refuse an artifact root outside its connector's root_dir when built MLOps((S3Connector(root_dir='s3://b/data'), 's3://b/artifacts')) was accepted, then every store failed with "Output path is outside the selected connector target". Under on_error='warn' those were warnings, so the fit returned with nothing stored and no model registered. kumo.MLOps now checks a (connector, path) root against the connector's root_dir when it is built and raises, naming the root_dir and the connector to pass instead. Verified: test_mlops. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_mlops.py | 21 +++++++++++++++++++++ packages/kumo-engine/test/test_mlops.py | 9 +++++++++ 2 files changed, 30 insertions(+) diff --git a/packages/kumo-engine/kumo/engine/_mlops.py b/packages/kumo-engine/kumo/engine/_mlops.py index 5f8b57b60..b4848c15e 100644 --- a/packages/kumo-engine/kumo/engine/_mlops.py +++ b/packages/kumo-engine/kumo/engine/_mlops.py @@ -71,6 +71,7 @@ def __post_init__(self) -> None: raise ValueError( f'{type(connector).__name__} cannot hold files at ' f'{path!r}; artifact_root needs a file location.') + _require_within_root(connector, path) root = _Root(connector, path, path) kinds = (frozenset(KINDS) if self.store is True else frozenset() if self.store is False else frozenset(self.store)) @@ -86,6 +87,26 @@ def wants(self, kind: str) -> bool: return kind in self.kinds +def _require_within_root(connector: Any, path: str) -> None: + from kumo_core.sources.object_store_sessions import ObjectStoreSessions + + sink = connector.sink_backend + if not isinstance(sink, ObjectStoreSessions): + return + try: + sink._resolved(connector, path) + except ValueError as error: + target = sink._root_for(connector) + within = (f'root_dir {sink.scheme}://{target.path}' + if target is not None else 'target') + name = type(connector).__name__ + raise ValueError( + f'artifact_root {path!r} is outside this {name}\'s {within} ' + f'({error}), so nothing could be stored there. Give the root ' + f'a connector that contains it: ({name}(root_dir={path!r}, ' + f'identity=...), {path!r}).') from error + + _CURRENT: contextvars.ContextVar[Optional[MLOps]] = contextvars.ContextVar( 'kumo_mlops', default=None) diff --git a/packages/kumo-engine/test/test_mlops.py b/packages/kumo-engine/test/test_mlops.py index 179606015..e289fc5ee 100644 --- a/packages/kumo-engine/test/test_mlops.py +++ b/packages/kumo-engine/test/test_mlops.py @@ -49,6 +49,15 @@ def test_a_root_is_a_uri_or_a_connector_and_path(tmp_path): assert not kumo.MLOps(artifact_root=str(tmp_path), store=False).kinds +def test_a_connector_root_outside_its_root_dir_fails_when_built(): + data = kumo.S3Connector(root_dir='s3://bucket/data') + assert kumo.MLOps(artifact_root=(data, 's3://bucket/data/kumo')).root + with pytest.raises(ValueError, match=r"'s3://bucket/kumo' is outside " + r"this S3Connector's root_dir s3://bucket/data .*" + r"S3Connector\(root_dir='s3://bucket/kumo'"): + kumo.MLOps(artifact_root=(data, 's3://bucket/kumo')) + + def test_a_call_argument_beats_the_nearest_block(tmp_path): outer = kumo.MLOps(artifact_root=str(tmp_path / 'outer')) inner = kumo.MLOps(artifact_root=str(tmp_path / 'inner')) From f8f10686c992efa9f1cb225c9cb0559d184854dd Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 16:21:34 +0000 Subject: [PATCH 010/278] Name TrainingTable.load_from_connector's query argument pquery Every other call that takes a PredictiveQuery names it pquery; this one said pq=. It is new in this PR, so it needs no alias. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/docs/mlops.md | 4 ++-- packages/kumo-engine/kumo/engine/pquery/training_table.py | 4 ++-- packages/kumo-engine/test/test_training_table_export.py | 4 ++-- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/packages/kumo-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index a22ce00c6..338c2eb73 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -50,7 +50,7 @@ Every artifact has one way out, `.export(OutputConfig(...))`, an | Artifact | Export | Load | |---|---|---| | Snapshot | `snapshot.export(SnapshotOutputConfig(connector, path))`, returns the location | `GraphSnapshot.load_from_connector(connector, path=location)`, or `(cfg, id=snapshot_id)` | -| Training table | `table.export(TrainingTableOutputConfig(connector, name))` | `TrainingTable.load_from_connector(connector, path=name, pq=pq, snapshot=snapshot)`, or `(cfg, path=f"training_tables/{id}", ...)` | +| Training table | `table.export(TrainingTableOutputConfig(connector, name))` | `TrainingTable.load_from_connector(connector, path=name, pquery=pquery, snapshot=snapshot)`, or `(cfg, path=f"training_tables/{id}", ...)` | | Model | `result.export(ModelOutputConfig(connector, path=... or location=...))` | `Trainer.load_from_connector(connector, path=... or location=...)`, or `(cfg, path=f"models/{id}")` | `fit(export_config=...)` is shorthand for fitting, then exporting. Loading needs a file destination (S3, GCS, local files, a volume or stage). Exporting again to a destination that already holds a complete copy of the same artifact reuses it (logged at info) instead of uploading again. @@ -71,7 +71,7 @@ A generated or caller-supplied training table carries `_provenance.json` in its - the snapshot it was generated from (that snapshot's manifest, embedded); - for `TrainingTable.from_source_table`, the source table's connector, location and version. -`TrainingTable.id` is deterministic: the same snapshot, query, plan, split seed and engine build give the same id. A caller-supplied source without a version gets a random id. An export to files (`TrainingTable.export`) writes the manifest next to the Parquet parts. `TrainingTable.load_from_connector(connector, path=..., pq=..., snapshot=...)` materializes an export again. +`TrainingTable.id` is deterministic: the same snapshot, query, plan, split seed and engine build give the same id. A caller-supplied source without a version gets a random id. An export to files (`TrainingTable.export`) writes the manifest next to the Parquet parts. `TrainingTable.load_from_connector(connector, path=..., pquery=..., snapshot=...)` materializes an export again. With MLOps on, each table is stored once under `/training_tables//`. diff --git a/packages/kumo-engine/kumo/engine/pquery/training_table.py b/packages/kumo-engine/kumo/engine/pquery/training_table.py index 417295022..db3888293 100644 --- a/packages/kumo-engine/kumo/engine/pquery/training_table.py +++ b/packages/kumo-engine/kumo/engine/pquery/training_table.py @@ -121,7 +121,7 @@ def _export_sidecars(self, connector, output_config) -> Dict[str, Any]: @classmethod def load_from_connector(cls, connector: Any, *, path: str, - pq: "PredictiveQuery", # noqa: F821 + pquery: "PredictiveQuery", # noqa: F821 snapshot: GraphSnapshot) -> "TrainingTable": r"""Materialize a table written by :meth:`export` (a directory with a ``SPLIT`` column) against ``snapshot``, as :meth:`from_source_table` @@ -132,7 +132,7 @@ def load_from_connector(cls, connector: Any, *, path: str, connector, path = on_root(connector, path) return cls.from_source_table(connector[path], TrainingTableSpec(split_col='SPLIT'), - pq=pq, snapshot=snapshot) + pq=pquery, snapshot=snapshot) def _store(self, settings: Any) -> None: from pyarrow.fs import FileType diff --git a/packages/kumo-engine/test/test_training_table_export.py b/packages/kumo-engine/test/test_training_table_export.py index 13f40ebd4..484f56481 100644 --- a/packages/kumo-engine/test/test_training_table_export.py +++ b/packages/kumo-engine/test/test_training_table_export.py @@ -336,7 +336,7 @@ def test_manifest_and_export_sidecar(graph, tmp_path): (directory / '_provenance.json').read_bytes()) == manifest.to_dict() reloaded = kumo.TrainingTable.load_from_connector( - connector, path='lineage', pq=pquery, snapshot=snapshot) + connector, path='lineage', pquery=pquery, snapshot=snapshot) manifest = reloaded.manifest assert reloaded.id != table.id assert manifest.parents.source_table == SourceTableEntry( @@ -457,7 +457,7 @@ def test_a_mlops_block_stores_each_table_once_under_the_root(graph, tmp_path): assert json.loads((stored / '_provenance.json').read_text())['id'] == ( table.id) reloaded = kumo.TrainingTable.load_from_connector( - cfg, path=f'training_tables/{table.id}', pq=pquery, snapshot=snapshot) + cfg, path=f'training_tables/{table.id}', pquery=pquery, snapshot=snapshot) assert _seeds(reloaded)[1] == _seeds(table)[1] off = pquery.generate_training_table( snapshot, mlops=kumo.MLOps(artifact_root=str(tmp_path / 'other'), From f17c1a286cc2ed17be37a1ee99f024af78e7d103 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 16:33:47 +0000 Subject: [PATCH 011/278] Version a single S3 or GCS object by its provider identity A single-object source was versioned by its modification time in seconds, so two uploads in the same second shared a data version (and a snapshot id). When the reader has the provider's object identity, the version is now its VersionId, else its entity tag, recorded as exact. Directories keep the listing hash. Verified: kumo-core tests/sources (454 passed). Co-Authored-By: Claude Opus 5.5 --- .../kumo-core/src/kumo_core/sources/object_store.py | 13 ++++++++++--- .../tests/sources/parquet/test_object_store.py | 10 ++++++++++ 2 files changed, 20 insertions(+), 3 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/sources/object_store.py b/packages/kumo-core/src/kumo_core/sources/object_store.py index bf1058d07..5f7f91d6b 100644 --- a/packages/kumo-core/src/kumo_core/sources/object_store.py +++ b/packages/kumo-core/src/kumo_core/sources/object_store.py @@ -98,8 +98,14 @@ _io_wanted: list[int] = [] -def _source_version(path: str, entries: Sequence[Any]) -> str: +def _source_version(path: str, entries: Sequence[Any], + identity: ObjectIdentity | None = None) -> str: if len(entries) == 1 and entries[0].path.rstrip("/") == path.rstrip("/"): + if identity is not None: + # The provider's version or entity tag; a second-resolution + # modification time cannot tell same-second rewrites apart. + token, version, *_ = (*identity(entries[0].path), "") + return SourceVersion(version or token.strip('"'), "exact") latest = entries[0].mtime return (SourceVersion(epoch_second(latest), "timestamp") if latest is not None else SourceVersion("-1")) @@ -161,7 +167,7 @@ def data_version(self, table: str, *, sql: str | None = None, exact=self.file_format is not None or not table) entries = ([info] if info.type is FileType.File else self._objects(table, info.path, include_hidden=True)) - return _source_version(info.path, entries) + return _source_version(info.path, entries, self.object_identity) def resolve(self, table: str) -> tuple[str, str]: """An explicit format addresses an exact path; discovery infers both.""" @@ -327,7 +333,8 @@ def prepare( include_hidden=not transform) # Stamp the same metadata listing the read consumes, before fetching # any files; later additions must not be stamped onto this output. - stamped = version or _source_version(path, found) + stamped = version or _source_version(path, found, + self.object_identity) source_bytes = (sum(entry.size for entry in found) if all( entry.size is not None and entry.size >= 0 for entry in found) else None) diff --git a/packages/kumo-core/tests/sources/parquet/test_object_store.py b/packages/kumo-core/tests/sources/parquet/test_object_store.py index 5989a09f5..b4fbf4c4e 100644 --- a/packages/kumo-core/tests/sources/parquet/test_object_store.py +++ b/packages/kumo-core/tests/sources/parquet/test_object_store.py @@ -149,6 +149,16 @@ def get_file_info(self, target): assert int(version) > 0 and version.kind == "timestamp" +def test_a_single_object_is_versioned_by_its_provider_identity(root): + reader = ObjectStoreSourceReader(LocalFileSystem(), root=str(root)) + reader.object_identity = lambda path: ('"etag-1"', '', '10') + version = reader.data_version("single") + assert (version, version.kind) == ("etag-1", "exact") + reader.object_identity = lambda path: ('"etag-1"', 'v2', '10') + assert reader.data_version("single") == "v2" + assert reader.data_version("events").kind == "listing" + + def test_version_of_exact_root_does_not_probe_other_suffixes(root): reader = reader_over(root / "single") with pytest.raises(SourceNotFoundError): From ed6c24f714a3a39899b8cf2cc540d86a92996594 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 19:00:29 +0000 Subject: [PATCH 012/278] Publish and load models through a filesystem-only sink's file paths A sink holding files through filesystem() but without model methods had its file paths rejected by fit/export and load_from_connector, which checked the raw sink. Both now take FilesystemSink over the sink for such paths; registry locations stay on the raw sink. Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/connector/__init__.py | 16 +++++++++++++++ .../kumo/engine/trainer/trainer.py | 9 ++++++--- .../kumo-engine/test/test_filestore_sink.py | 20 +++++++++++++++++++ 3 files changed, 42 insertions(+), 3 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/connector/__init__.py b/packages/kumo-engine/kumo/engine/connector/__init__.py index 64c869506..93265951d 100644 --- a/packages/kumo-engine/kumo/engine/connector/__init__.py +++ b/packages/kumo-engine/kumo/engine/connector/__init__.py @@ -61,6 +61,22 @@ def file_store(connector: Any, path: str) -> Optional[Tuple[Any, str]]: return None if sink is None else sink.filesystem(connector, path) +def model_sink(connector: Any, path: Optional[str]) -> Any: + r"""The sink that publishes or fetches a model at ``path``: the file + writers over ``connector``'s sink (`FilesystemSink`) for a file path of + a sink without its own, else the sink.""" + from kumo_core.object_store_sink import ( + FilesystemSink, + ObjectStoreSinkBackend, + ) + + sink = connector.sink_backend + if (path is None or isinstance(sink, ObjectStoreSinkBackend) + or _file_sink(connector, path) is None): + return sink + return FilesystemSink(sink) + + __all__ = [ "Connector", "FileConnector", diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 28c3156a3..7c9768bfc 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -53,6 +53,7 @@ require_trained_embedding_models, ) from kumo.engine._scratch import new_run_id, trainer_scratch_dir +from kumo.engine.connector import model_sink from kumo.engine.connector.base import Connector from kumo.engine.graph import GraphSnapshot from kumo.engine.pquery.prediction_table import PredictionTable @@ -358,7 +359,8 @@ def fit( from kumo.engine.trainer._runtime import local_training_bindings if export_config is not None and not callable( - getattr(export_config.connector.sink_backend, 'publish_model', + getattr(model_sink(export_config.connector, + export_config.path), 'publish_model', None)): raise NotImplementedError( f'{type(export_config.connector).__name__} cannot publish ' @@ -547,7 +549,8 @@ def _finish_fit( ) if export_config is not None: spec = read_training_spec(training_result.bundle_dir) - published = export_config.connector.sink_backend.publish_model( + sink = model_sink(export_config.connector, export_config.path) + published = sink.publish_model( export_config.connector, str(training_result.bundle_dir), path=export_config.path, @@ -819,7 +822,7 @@ def load_from_connector( address = ModelOutputConfig(model_registry, path=path, location=location) bundle_dir = trainer_scratch_dir(new_run_id()) / 'fetched-model' - sink = getattr(model_registry, 'sink_backend') + sink = model_sink(model_registry, address.path) if not isinstance(sink, ModelFetchingSink): raise TypeError( 'model_registry.sink_backend must implement ModelFetchingSink') diff --git a/packages/kumo-engine/test/test_filestore_sink.py b/packages/kumo-engine/test/test_filestore_sink.py index 21a59362d..40a54ac4a 100644 --- a/packages/kumo-engine/test/test_filestore_sink.py +++ b/packages/kumo-engine/test/test_filestore_sink.py @@ -135,6 +135,26 @@ def test_bundles_round_trip_through_a_filesystem_only_sink( bundle / name).read_bytes() +def test_a_trainer_publishes_and_loads_through_a_filesystem_only_sink( + fitted): + trainer, *_ = fitted + connector = SimpleNamespace(source_type='SNOWFLAKE', + sink_backend=StageSink()) + result = trainer._finish_fit( + trainer._training_result, predictive_query=trainer._training_query, + snapshot_id=trainer._snapshot_id, + training_table_generation=trainer.training_table_generation_plan, + export_config=kumo.ModelOutputConfig(connector, + path='@stage/models/m')) + assert result.published_model.uri == '@stage/models/m' + loaded = kumo.Trainer.load_from_connector(connector, + path='@stage/models/m') + assert loaded._snapshot_id == trainer._snapshot_id + with pytest.raises(NotImplementedError, match='cannot publish'): + trainer.fit(None, snapshot=None, + export_config=kumo.ModelOutputConfig(connector, path='T')) + + def test_parquet_writes_stay_plain_up_to_one_part(tmp_path, connector): files = FilesystemSink(connector.sink_backend) frame = pd.DataFrame({'x': range(10)}) From b808c9aa59b56b958dc0a7ecaa655592a25f7615 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 16:33:47 +0000 Subject: [PATCH 013/278] Check that the artifact root is writable when kumo.MLOps is built A root the identity could not write was found only at the first store, after a fit or predict had run. Building kumo.MLOps (when it stores anything) now writes and deletes a probe object under an object-store, volume or stage root, and checks a local root's nearest existing directory, raising PermissionError with the root and the cause. Verified: test_mlops, test_bundle_provenance, test_snapshot_export, test_training_table_export (71 passed). Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_mlops.py | 36 ++++++++++++++++++++++ packages/kumo-engine/test/test_mlops.py | 23 ++++++++++++-- 2 files changed, 57 insertions(+), 2 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_mlops.py b/packages/kumo-engine/kumo/engine/_mlops.py index b4848c15e..120b71792 100644 --- a/packages/kumo-engine/kumo/engine/_mlops.py +++ b/packages/kumo-engine/kumo/engine/_mlops.py @@ -5,6 +5,8 @@ import contextlib import contextvars +import os +import posixpath from dataclasses import dataclass, field from typing import Any, FrozenSet, Iterable, Iterator, Optional, Tuple, Union @@ -79,6 +81,8 @@ def __post_init__(self) -> None: if unknown: raise ValueError(f'Unknown store kinds {unknown}; choose from ' f'{list(KINDS)}.') + if kinds: + _require_writable(root) object.__setattr__(self, 'root', root) object.__setattr__(self, 'kinds', kinds) @@ -107,6 +111,38 @@ def _require_within_root(connector: Any, path: str) -> None: f'identity=...), {path!r}).') from error +def _require_writable(root: _Root) -> None: + from uuid import uuid4 + + from kumo.engine.connector import file_store + + store = file_store(root.connector, + root.location(f'.kumo-write-check-{uuid4().hex}')) + if store is None: + return + filesystem, probe = store + try: + if filesystem.type_name == 'local': + existing = posixpath.dirname(probe) + while not os.path.exists(existing): + existing = posixpath.dirname(existing) + if not os.access(existing, os.W_OK | os.X_OK): + raise PermissionError(f'{existing} is not writable') + return + try: + filesystem.open_output_stream(probe).close() + except FileNotFoundError: # a store with directories, before first use + filesystem.create_dir(posixpath.dirname(probe), recursive=True) + filesystem.open_output_stream(probe).close() + filesystem.delete_file(probe) + except Exception as error: + raise PermissionError( + f'kumo.MLOps cannot write under artifact_root ' + f'{root.uri("")!r}: {type(error).__name__}: {error}. Grant this ' + f'identity write and delete rights on it, or pick another ' + f'root.') from error + + _CURRENT: contextvars.ContextVar[Optional[MLOps]] = contextvars.ContextVar( 'kumo_mlops', default=None) diff --git a/packages/kumo-engine/test/test_mlops.py b/packages/kumo-engine/test/test_mlops.py index e289fc5ee..3d3af18d9 100644 --- a/packages/kumo-engine/test/test_mlops.py +++ b/packages/kumo-engine/test/test_mlops.py @@ -3,6 +3,7 @@ # SPDX-License-Identifier: Apache-2.0 import asyncio +import os import threading from types import SimpleNamespace @@ -51,11 +52,29 @@ def test_a_root_is_a_uri_or_a_connector_and_path(tmp_path): def test_a_connector_root_outside_its_root_dir_fails_when_built(): data = kumo.S3Connector(root_dir='s3://bucket/data') - assert kumo.MLOps(artifact_root=(data, 's3://bucket/data/kumo')).root + assert kumo.MLOps(artifact_root=(data, 's3://bucket/data/kumo'), + store=False).root with pytest.raises(ValueError, match=r"'s3://bucket/kumo' is outside " r"this S3Connector's root_dir s3://bucket/data .*" r"S3Connector\(root_dir='s3://bucket/kumo'"): - kumo.MLOps(artifact_root=(data, 's3://bucket/kumo')) + kumo.MLOps(artifact_root=(data, 's3://bucket/kumo'), store=False) + + +def test_a_root_that_cannot_be_written_fails_when_built(tmp_path): + if os.geteuid() == 0: + pytest.skip('root writes through a read-only directory') + read_only = tmp_path / 'read-only' + read_only.mkdir() + read_only.chmod(0o500) + try: + with pytest.raises(PermissionError, match='cannot write under ' + 'artifact_root'): + kumo.MLOps(artifact_root=str(read_only / 'kumo')) + assert kumo.MLOps(artifact_root=str(read_only), store=False) + finally: + read_only.chmod(0o700) + kumo.MLOps(artifact_root=str(tmp_path / 'kumo')) + assert not (tmp_path / 'kumo').exists() def test_a_call_argument_beats_the_nearest_block(tmp_path): From 312085b9cbe43eeffaff5cf5e4e5874a29307988 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 19:53:57 +0000 Subject: [PATCH 014/278] Keep the previous file until a chunked replacement commits A chunked write now stages its parts in its own .chunks// directory and removes the old plain object and older generations only after its _SUCCESS is uploaded. Readers take the newest committed generation, so a failed or interrupted write leaves the previous file readable. A writer whose write raised commits nothing on close and the original error surfaces. Co-Authored-By: Claude Opus 5.5 --- .../kumo-core/src/kumo_core/chunked_fs.py | 65 ++++++++++++----- packages/kumo-core/tests/test_chunked_fs.py | 71 ++++++++++++++++--- .../kumo-engine/test/test_filestore_sink.py | 8 ++- 3 files changed, 112 insertions(+), 32 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/chunked_fs.py b/packages/kumo-core/src/kumo_core/chunked_fs.py index 8972badd8..24a582f2b 100644 --- a/packages/kumo-core/src/kumo_core/chunked_fs.py +++ b/packages/kumo-core/src/kumo_core/chunked_fs.py @@ -4,9 +4,12 @@ r"""A pyarrow filesystem over stores that only put, list, read and remove. A file larger than :attr:`ChunkedFileHandler.part_bytes` is stored as -``.chunks/00000 ... NNNNN`` followed by an empty ``_SUCCESS``; anything -smaller is one plain object. Readers see one logical file either way. A write -removes the other form; if a crash leaves both, the plain object wins. A writer +``.chunks//00000 ... NNNNN`` followed by an empty +``_SUCCESS``; anything smaller is one plain object. Readers see one logical +file either way. A write removes the other form and older generations only +once its own is committed, and a write that failed commits nothing; if a +crash leaves both forms, the plain object wins, and of several committed +generations the newest. A writer spools one part at a time on local disk, under :func:`spool_root`. A handler that can stream an upload implements ``_stream``: a writer told the size (:data:`SIZE_KEY` in the stream metadata) streams a file of up to @@ -14,11 +17,14 @@ disk; without a size, up to ``plain_bytes`` is spooled and stays plain. """ +import contextlib import io import os import posixpath import shutil import tempfile +import time +import uuid from typing import ( BinaryIO, Callable, @@ -209,10 +215,17 @@ def open_output_stream(self, path, metadata): def _logical_parts(objects: Dict[str, int], directory: str) -> Optional[List[Tuple[str, int]]]: - if f'{directory}/{DONE}' not in objects: + committed = [ + posixpath.dirname(path) for path in objects + if posixpath.basename(path) == DONE + and posixpath.dirname(posixpath.dirname(path)) == directory + ] + if not committed: return None + generation = max(committed) return sorted((path, size) for path, size in objects.items() - if posixpath.basename(path) != DONE) + if posixpath.dirname(path) == generation + and posixpath.basename(path) != DONE) def _logical_files(objects: Dict[str, int]) -> Dict[str, int]: @@ -220,7 +233,7 @@ def _logical_files(objects: Dict[str, int]) -> Dict[str, int]: files: Dict[str, int] = {} chunked: Dict[str, Dict[str, int]] = {} for path, size in objects.items(): - directory = posixpath.dirname(path) + directory = posixpath.dirname(posixpath.dirname(path)) if directory.endswith(CHUNKS): chunked.setdefault(directory, {})[path] = size else: @@ -322,6 +335,8 @@ def __init__(self, handler: ChunkedFileHandler, path: str, size: Optional[int] = None) -> None: self._handler = handler self._path = path + self._chunks = (f'{path}{CHUNKS}/{time.time_ns():020d}-' + f'{uuid.uuid4().hex[:8]}') self._dir = tempfile.mkdtemp(prefix='kumo-chunked-', dir=spool_root(handler)) self._spool = open(os.path.join(self._dir, 'spool'), 'wb') @@ -333,11 +348,11 @@ def __init__(self, handler: ChunkedFileHandler, path: str, self._limit = (min(handler.plain_bytes, handler.part_bytes) if self._streams else handler.part_bytes) self._plain: Optional[BinaryIO] = None + self._failed = False if self._streams and size is not None: if size <= handler.part_bytes: self._plain = handler._stream(path) else: - self._clear() self._part = self._open_part() def writable(self): @@ -347,7 +362,13 @@ def tell(self): return self._written def write(self, data): - view = memoryview(data).cast('B') + try: + return self._write(memoryview(data).cast('B')) + except BaseException: + self._failed = True + raise + + def _write(self, view: memoryview) -> int: total = len(view) if self._plain is not None: self._plain.write(view) @@ -383,23 +404,25 @@ def _upload(self, target: str) -> None: self._spool = open(os.path.join(self._dir, 'spool'), 'wb') self._filled = 0 - def _clear(self) -> None: - self._handler._remove(self._path + CHUNKS, tree=True) + def _commit_parts(self) -> None: + self._upload(f'{self._chunks}/{DONE}') if self._handler._size(self._path) is not None: self._handler._remove(self._path, tree=False) + stale = { + posixpath.dirname(path) + for path in self._handler._list(self._path + CHUNKS) + } - {self._chunks} + for generation in stale: + self._handler._remove(generation, tree=True) def _flush_part(self) -> None: - if not self._parts: - self._clear() - self._upload(f'{self._path}{CHUNKS}/{self._parts:05d}') + self._upload(f'{self._chunks}/{self._parts:05d}') self._parts += 1 def _open_part(self) -> BinaryIO: - return self._handler._stream( - f'{self._path}{CHUNKS}/{self._parts:05d}') + return self._handler._stream(f'{self._chunks}/{self._parts:05d}') def _start_stream(self) -> None: - self._clear() self._spool.close() spooled = os.path.join(self._dir, 'spool') self._part = self._open_part() @@ -418,7 +441,11 @@ def close(self): if self.closed: return try: - if self._plain is not None: + if self._failed: + with contextlib.suppress(Exception): + if self._part is not None: + self._part.close() + elif self._plain is not None: self._plain.close() self._plain = None if self._handler._list(self._path + CHUNKS): @@ -426,10 +453,10 @@ def close(self): elif self._part is not None: self._part.close() self._part = None - self._upload(f'{self._path}{CHUNKS}/{DONE}') + self._commit_parts() elif self._parts: self._flush_part() - self._upload(f'{self._path}{CHUNKS}/{DONE}') + self._commit_parts() else: self._upload(self._path) if self._handler._list(self._path + CHUNKS): diff --git a/packages/kumo-core/tests/test_chunked_fs.py b/packages/kumo-core/tests/test_chunked_fs.py index f93b19e60..771014fd8 100644 --- a/packages/kumo-core/tests/test_chunked_fs.py +++ b/packages/kumo-core/tests/test_chunked_fs.py @@ -77,6 +77,14 @@ def _read(fs, path): return reader.readall() +def _generations(handler): + parts = {} + for name in sorted(handler.objects): + directory, _, part = name.rpartition('/') + parts.setdefault(directory, []).append(part) + return parts + + @pytest.mark.parametrize('size', [0, 3, PART]) def test_files_up_to_one_part_are_one_plain_object(store, size): handler, fs = store @@ -93,10 +101,9 @@ def test_larger_files_are_parts_committed_last_and_read_back_exactly( handler, fs = store data = os.urandom(size) _write(fs, 'root/a.bin', data) - assert sorted(handler.objects) == [ - 'root/a.bin.chunks/00000', 'root/a.bin.chunks/00001', - 'root/a.bin.chunks/00002', 'root/a.bin.chunks/_SUCCESS' - ] + (generation, parts), = _generations(handler).items() + assert generation.startswith('root/a.bin.chunks/') + assert parts == ['00000', '00001', '00002', '_SUCCESS'] info = fs.get_file_info('root/a.bin') assert (info.type, info.size) == (FileType.File, len(data)) assert _read(fs, 'root/a.bin') == data @@ -110,7 +117,7 @@ def test_larger_files_are_parts_committed_last_and_read_back_exactly( def test_an_uncommitted_chunked_write_is_not_a_file(store): handler, fs = store - handler.objects['root/a.bin.chunks/00000'] = b'x' * PART + handler.objects['root/a.bin.chunks/1/00000'] = b'x' * PART assert fs.get_file_info('root/a.bin').type is not FileType.File with pytest.raises(FileNotFoundError): fs.open_input_file('root/a.bin') @@ -127,6 +134,49 @@ def test_a_write_removes_the_other_form(store, first, second): assert any('.chunks/' in name for name in handler.objects) != plain +@pytest.mark.parametrize('sized', [False, True]) +@pytest.mark.parametrize('streams', [False, True]) +@pytest.mark.parametrize('old', [3, 2 * PART]) +def test_a_failed_chunked_write_leaves_the_old_file(monkeypatch, streams, + sized, old): + from kumo_core.object_store_sink import open_sized + + handler = StreamingHandler() if streams else MemoryHandler() + fs = PyFileSystem(handler) + _write(fs, 'root/a.bin', b'o' * old) + upload = '_stream' if streams else '_put' + original = getattr(handler, upload) + + def second_part_fails(*args): + if args[-1].endswith('/00001'): + raise OSError('upload failed') + return original(*args) + + monkeypatch.setattr(handler, upload, second_part_fails) + with pytest.raises(OSError, match='upload failed'): + with (open_sized(fs, 'root/a.bin', 3 * PART) + if sized else fs.open_output_stream('root/a.bin')) as writer: + writer.write(b'n' * (3 * PART)) + assert _read(PyFileSystem(handler), 'root/a.bin') == b'o' * old + + monkeypatch.setattr(handler, upload, original) + _write(fs, 'root/a.bin', b'n' * (3 * PART)) + assert _read(PyFileSystem(handler), 'root/a.bin') == b'n' * (3 * PART) + assert list(_generations(handler).values()) == [ + ['00000', '00001', '00002', '_SUCCESS'] + ] + + +def test_the_newest_committed_generation_is_read(store): + handler, fs = store + _write(fs, 'root/a.bin', b'1' * (2 * PART)) + old = dict(handler.objects) + _write(fs, 'root/a.bin', b'2' * (2 * PART)) + handler.objects.update(old) + handler._forget('root/a.bin') + assert _read(fs, 'root/a.bin') == b'2' * (2 * PART) + + def test_plain_wins_when_a_crash_left_both_forms(store): handler, fs = store _write(fs, 'root/a.bin', b'c' * (2 * PART)) @@ -223,12 +273,11 @@ def test_a_streaming_handler_streams_parts_and_spools_only_small_files(size): if size <= 6: assert handler.objects == {'root/a.bin': data} return - parts = -(-size // PART) - assert sorted(handler.objects) == [ - f'root/a.bin.chunks/{i:05d}' for i in range(parts) - ] + ['root/a.bin.chunks/_SUCCESS'] - assert handler.streamed[-parts:] == [ - f'root/a.bin.chunks/{i:05d}' for i in range(parts) + parts = [f'{i:05d}' for i in range(-(-size // PART))] + (generation, stored), = _generations(handler).items() + assert stored == parts + ['_SUCCESS'] + assert handler.streamed[-len(parts):] == [ + f'{generation}/{part}' for part in parts ] assert len(handler.spooled) == 2 with fs.open_input_file('root/a.bin') as reader: diff --git a/packages/kumo-engine/test/test_filestore_sink.py b/packages/kumo-engine/test/test_filestore_sink.py index 40a54ac4a..0a769cad5 100644 --- a/packages/kumo-engine/test/test_filestore_sink.py +++ b/packages/kumo-engine/test/test_filestore_sink.py @@ -127,7 +127,9 @@ def test_bundles_round_trip_through_a_filesystem_only_sink( assert published.uri == '@stage/models/m' store = connector.sink_backend.store assert list(store.objects)[-1] == '@stage/models/m/_provenance.json' - assert '@stage/models/m/sub/weights.bin.chunks/_SUCCESS' in store.objects + assert any( + name.startswith('@stage/models/m/sub/weights.bin.chunks/') + and name.endswith('/_SUCCESS') for name in store.objects) files.fetch_model(connector, str(tmp_path / 'fetched'), path='@stage/models/m') for name in ('_provenance.json', 'sub/weights.bin'): @@ -163,7 +165,9 @@ def test_parquet_writes_stay_plain_up_to_one_part(tmp_path, connector): pd.DataFrame({'x': [os.urandom(50).hex() for _ in range(100)]})) store = connector.sink_backend.store assert '@stage/small.parquet' in store.objects - assert '@stage/big.parquet.chunks/_SUCCESS' in store.objects + assert any( + name.startswith('@stage/big.parquet.chunks/') + and name.endswith('/_SUCCESS') for name in store.objects) assert pd.read_parquet( '@stage/big.parquet', filesystem=PyFileSystem(store)).shape == (100, 1) From f86cdf5843a02330fd98ac3cd4a883f32b1b39a2 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 18:30:08 +0000 Subject: [PATCH 015/278] Name the full path when no snapshot export is found Loading a snapshot from a path with no kumo_snapshot.json said 'No complete snapshot export at ; re-run export', which hid where it looked and pointed at the wrong fix when the path was wrong. It now prints the resolved location and says to check the path or export there. Verified: test_snapshot_export (new: the message has the full path). Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_snapshot_io.py | 6 ++++-- packages/kumo-engine/test/test_snapshot_export.py | 7 +++++++ 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_snapshot_io.py b/packages/kumo-engine/kumo/engine/_snapshot_io.py index 7f83970c2..5b8060244 100644 --- a/packages/kumo-engine/kumo/engine/_snapshot_io.py +++ b/packages/kumo-engine/kumo/engine/_snapshot_io.py @@ -372,8 +372,10 @@ def read(connector: Any, path: str) -> Tuple[str, dict]: except ValueError as error: raise _corrupt(path, f'unreadable {_MANIFEST}: {error}') from error if manifest is None: - raise FileNotFoundError(f'No complete snapshot export at {path}; ' - 're-run export.') + raise FileNotFoundError( + f'No complete snapshot export at {root}: it has no {_MANIFEST}. ' + f'Check the path, or export the snapshot there with ' + f'GraphSnapshot.export.') if not isinstance(manifest, dict): raise _corrupt(path, 'manifest is not a JSON object') if manifest.get('format') != _FORMAT: diff --git a/packages/kumo-engine/test/test_snapshot_export.py b/packages/kumo-engine/test/test_snapshot_export.py index 4452bc06d..0e67c3286 100644 --- a/packages/kumo-engine/test/test_snapshot_export.py +++ b/packages/kumo-engine/test/test_snapshot_export.py @@ -787,3 +787,10 @@ def test_a_mlops_block_stores_each_snapshot_once(tmp_path, scratch, graph, graph.snapshot(mlops=kumo.MLOps(artifact_root=str(tmp_path / 'off'), store={'model'})) assert not (tmp_path / 'off').exists() + + +def test_loading_a_missing_export_names_its_full_path(tmp_path, exports): + missing = tmp_path / 'exports' / 'snapshots' / 'graph' / 'version' + with pytest.raises(FileNotFoundError, + match=f'No complete snapshot export at {missing}: '): + _load(exports, 'snapshots/graph/version') From dbcadb480f67eff8990400fe89deb998918c84cb Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 20:47:24 +0000 Subject: [PATCH 016/278] Export predictions through a filesystem-only sink's file writers Prediction export called write_files on the raw sink, so a sink that only holds files through filesystem() was refused. Export now takes the same file writers as model publishing for file destinations. Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/trainer/output_config.py | 10 ++++++- .../kumo/engine/trainer/trainer.py | 29 +++++++------------ .../kumo-engine/test/test_filestore_sink.py | 20 +++++++++++-- 3 files changed, 36 insertions(+), 23 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/output_config.py b/packages/kumo-engine/kumo/engine/trainer/output_config.py index 43da33eac..d947fe775 100644 --- a/packages/kumo-engine/kumo/engine/trainer/output_config.py +++ b/packages/kumo-engine/kumo/engine/trainer/output_config.py @@ -11,7 +11,7 @@ from kumoapi.plugin import ModelRegistryLocation, SourceConnector from kumoapi.prediction_outputs import resolve_prediction_output_paths -from kumo.engine.connector import is_file_destination +from kumo.engine.connector import is_file_destination, model_sink from kumo.engine.connector.base import Connector @@ -64,6 +64,14 @@ def destinations(self) -> dict[str, str]: assert self.artifact_type is not None and self.table_name is not None return {self.artifact_type.value.lower(): self.table_name} + def writer(self, destination: str): + sink = model_sink(self.connector, destination) + if not callable(getattr(sink, 'write_files', None)): + raise NotImplementedError( + 'Prediction export requires FileArtifactSink.write_files; ' + 'this sink only supports in-memory table output') + return sink + @dataclass(frozen=True) class TrainingTableOutputConfig: diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 7c9768bfc..3e87b954e 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -157,20 +157,18 @@ def export(self, output_config: PredictionOutputConfig) -> None: self._export(output_config) def _export(self, output_config: PredictionOutputConfig) -> None: - write_files = getattr(output_config.connector.sink_backend, - 'write_files', None) - if not callable(write_files): - raise NotImplementedError( - 'Prediction export requires FileArtifactSink.write_files') + for destination in output_config.destinations().values(): + output_config.writer(destination) outputs = [] for name, destination in output_config.destinations().items(): kind = PredictionArtifactType(name.upper()) files = PredictionResult(self.selection, kind).files outputs.append((kind, destination, files)) for kind, destination, files in outputs: - write_files(output_config.connector, destination, list(files), - output_config.write_mode, - export_id=f'{self.selection.job_id}-{kind.value}') + output_config.writer(destination).write_files( + output_config.connector, destination, list(files), + output_config.write_mode, + export_id=f'{self.selection.job_id}-{kind.value}') def _find_holdout(log_dir: str) -> tuple[Optional[Path], Optional[bool]]: @@ -654,18 +652,11 @@ def predict( has_predictions=PredictionArtifactType.PREDICTIONS in artifacts, ) - write_files = None if output_config is not None: - write_files = getattr(output_config.connector.sink_backend, - 'write_files', None) - if not callable(write_files): - raise NotImplementedError( - 'Prediction export requires FileArtifactSink.write_files; ' - 'this sink only supports in-memory table output') - preflight = getattr(output_config.connector.sink_backend, - 'preflight_write_files', None) - if callable(preflight): - for destination in destinations.values(): + for destination in destinations.values(): + preflight = getattr(output_config.writer(destination), + 'preflight_write_files', None) + if callable(preflight): preflight(output_config.connector, destination, write_mode=output_config.write_mode) diff --git a/packages/kumo-engine/test/test_filestore_sink.py b/packages/kumo-engine/test/test_filestore_sink.py index 0a769cad5..8716310f8 100644 --- a/packages/kumo-engine/test/test_filestore_sink.py +++ b/packages/kumo-engine/test/test_filestore_sink.py @@ -80,9 +80,6 @@ def filesystem(self, connector, path): assert self.holds_files(connector, path) return PyFileSystem(self.store), path.rstrip('/') - def write_files(self, *args, **kwargs): - return FilesystemSink(self).write_files(*args, **kwargs) - @pytest.fixture def connector(tmp_path): @@ -157,6 +154,23 @@ def test_a_trainer_publishes_and_loads_through_a_filesystem_only_sink( export_config=kumo.ModelOutputConfig(connector, path='T')) +def test_predictions_export_through_a_filesystem_only_sink(fitted, connector): + trainer, table, snapshot = fitted + output = PredictionOutputConfig( + connector, output_paths={'predictions': '@stage/preds/'}) + result = trainer.predict(table, snapshot=snapshot, output_config=output, + binary_classification_threshold=0.5) + store = connector.sink_backend.store + assert any(name.startswith('@stage/preds/') for name in store.objects) + result.export( + PredictionOutputConfig(connector, + output_paths={'predictions': '@stage/again/'})) + assert any(name.startswith('@stage/again/') for name in store.objects) + with pytest.raises(NotImplementedError, match='requires'): + PredictionOutputConfig(connector, PredictionArtifactType.PREDICTIONS, + 'DB.S.T').writer('DB.S.T') + + def test_parquet_writes_stay_plain_up_to_one_part(tmp_path, connector): files = FilesystemSink(connector.sink_backend) frame = pd.DataFrame({'x': range(10)}) From fe8b8ed12f970db4626305fc9d1adba726ac12a6 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 27 Oct 2026 09:50:58 +0000 Subject: [PATCH 017/278] Hold MLOps artifacts on Unity Catalog volumes through the Files API Co-Authored-By: Claude Opus 5.5 --- packages/kumo-databricks/README.md | 30 +++ .../src/kumo_databricks/registry.py | 3 +- .../src/kumo_databricks/sink.py | 90 ++++++- .../src/kumo_databricks/volume.py | 240 ++++++++++++++++++ packages/kumo-databricks/tests/fakes.py | 81 +++++- packages/kumo-databricks/tests/test_volume.py | 180 +++++++++++++ tests/public/Dockerfile.dockerignore | 4 + 7 files changed, 622 insertions(+), 6 deletions(-) create mode 100644 packages/kumo-databricks/src/kumo_databricks/volume.py create mode 100644 packages/kumo-databricks/tests/test_volume.py diff --git a/packages/kumo-databricks/README.md b/packages/kumo-databricks/README.md index 6769362fc..865ae85a1 100644 --- a/packages/kumo-databricks/README.md +++ b/packages/kumo-databricks/README.md @@ -67,6 +67,36 @@ events = Table.from_source_table( snapshot = Graph(tables=[events]).snapshot() ``` +## Artifacts on a volume + +A `/Volumes////...` path is a file destination of a +`DatabricksConnector` from any machine, not only on Databricks compute: the +sink reads and writes it through the Files API with the connector's identity. + +```python +from kumo.engine import ModelOutputConfig, PredictionOutputConfig, Trainer + +root = "/Volumes/main/kumo/artifacts" +trainer.fit(training_table, snapshot=snapshot, export_config=ModelOutputConfig( + databricks, path=f"{root}/models/churn")) +trainer = Trainer.load_from_connector(databricks, path=f"{root}/models/churn") +trainer.predict(prediction_table, snapshot=snapshot, output_config=PredictionOutputConfig( + databricks, output_paths={"predictions": f"{root}/scores/"})) +``` + +For an MLOps artifact root on a volume, see `packages/kumo-engine/docs/mlops.md`. + +A file whose size is known when it is written (bundle files, exported +Parquet parts) is stored as one object up to 4 GiB, so Spark +and other tools read it directly. A larger one streams into +`.chunks/00000`, `00001`, ... (4 GiB each, under the Files API's 5 GiB +upload cap) with an empty `_SUCCESS` written last, and Kumo reads the parts +back as one file. A file written without a known size stays one object up to +256 MiB, staged on local disk under the engine scratch directory, and streams +into parts above that. The identity needs `USE CATALOG`, `USE SCHEMA`, and +`READ VOLUME` (loads) and `WRITE VOLUME` (exports) on the volume; no +warehouse is used. + ## Model registry The sink publishes Kumo bundles directly to Unity Catalog using the connector's diff --git a/packages/kumo-databricks/src/kumo_databricks/registry.py b/packages/kumo-databricks/src/kumo_databricks/registry.py index d2d88de09..1f7d2e352 100644 --- a/packages/kumo-databricks/src/kumo_databricks/registry.py +++ b/packages/kumo-databricks/src/kumo_databricks/registry.py @@ -34,7 +34,8 @@ def _address(connector: SourceConnector, path: Optional[str], location: Optional[ModelRegistryLocation]) -> str: if path is not None: raise ValueError('Databricks models take a ModelRegistryLocation, ' - 'rather than a path.') + 'or a /Volumes////... ' + 'path.') if location is None: raise ValueError('Pass a ModelRegistryLocation naming the model.') config = _target_config(connector) diff --git a/packages/kumo-databricks/src/kumo_databricks/sink.py b/packages/kumo-databricks/src/kumo_databricks/sink.py index b258c70a5..56d7717fb 100644 --- a/packages/kumo-databricks/src/kumo_databricks/sink.py +++ b/packages/kumo-databricks/src/kumo_databricks/sink.py @@ -22,16 +22,23 @@ import os import tempfile from pathlib import Path -from typing import Optional, Sequence +from typing import Mapping, Optional, Sequence import pandas as pd from databricks.sql.client import Connection +from kumo_core.object_store_sink import FilesystemSink from kumo_databricks.registry import DatabricksModelRegistryMixin from kumo_databricks.session import _target_config +from kumo_databricks.volume import is_volume_path, volume_filesystem from kumo_databricks.write import RunVolume from kumoapi.data_source import DataSourceType from kumoapi.jobs import WriteMode -from kumoapi.plugin import SourceConnector +from kumoapi.plugin import ( + ModelRegistryLocation, + PublishedModel, + SourceConnector, +) +from pyarrow.fs import FileSystem class DatabricksSinkBackend(DatabricksModelRegistryMixin): @@ -125,6 +132,9 @@ def write_table( names a schema other than the connector's, or if the staging volume is neither reachable nor creatable. """ + if self.holds_files(connector, table_name): + return FilesystemSink(self).write_table( + connector, table_name, frame, write_mode, export_id=export_id) del project_id, export_id from kumo_databricks import output, sql, write @@ -152,8 +162,13 @@ def write_table( def write_files(self, connector, table_name, files, write_mode=WriteMode.OVERWRITE, project_id=None, *, - export_id=None, extra_columns=None): - """Pass selected Parquet directly to the existing hosted COPY owner.""" + export_id=None, extra_columns=None, **options): + """Copy Parquet to a volume path, or COPY it into a table.""" + if self.holds_files(connector, table_name): + return FilesystemSink(self).write_files( + connector, table_name, files, write_mode, + export_id=export_id, extra_columns=extra_columns, + **options) del project_id, export_id from functools import partial @@ -187,6 +202,73 @@ def write_files(self, connector, table_name, files, prepare_file=prepare) return sql.qualified_name(config.catalog, schema, table) + def publish_model( + self, + connector: SourceConnector, + bundle_dir: str, + *, + path: Optional[str] = None, + location: Optional[ModelRegistryLocation] = None, + metadata: Mapping[str, str] = {}, + metrics: Mapping[str, float] = {}, + project_id: Optional[str] = None, + ) -> PublishedModel: + """Register in Unity Catalog, or copy to a volume ``path``.""" + if location is None and path and self.holds_files(connector, path): + return FilesystemSink(self).publish_model( + connector, bundle_dir, path=path, metadata=metadata) + return super().publish_model(connector, bundle_dir, path=path, + location=location, metadata=metadata, + metrics=metrics, project_id=project_id) + + def fetch_model( + self, + connector: SourceConnector, + local_dir: str, + *, + path: Optional[str] = None, + location: Optional[ModelRegistryLocation] = None, + project_id: Optional[str] = None, + ) -> PublishedModel: + """Download from Unity Catalog, or from a volume ``path``.""" + if location is None and path and self.holds_files(connector, path): + return FilesystemSink(self).fetch_model(connector, local_dir, + path=path) + return super().fetch_model(connector, local_dir, path=path, + location=location, project_id=project_id) + + def preflight_write_files(self, connector: SourceConnector, + table_name: str, + write_mode: Optional[WriteMode] = None) -> None: + """Refuse an occupied volume directory before scoring, unless + ``write_mode`` is ``OVERWRITE``; tables are not checked.""" + from kumo_core.artifact_export import ArtifactDestinationConflict + from pyarrow.fs import FileSelector + + if (write_mode is WriteMode.OVERWRITE or not table_name.endswith('/') + or not self.holds_files(connector, table_name)): + return + filesystem, path = self.filesystem(connector, table_name) + if filesystem.get_file_info(FileSelector(path, allow_not_found=True)): + raise ArtifactDestinationConflict( + f"Output directory '{path}/' already contains data from " + f"another export. Choose an empty destination path, or " + f"replace it with write_mode=WriteMode.OVERWRITE.") + + def holds_files(self, connector: SourceConnector, path: str) -> bool: + """``/Volumes////...``; anything else is a + table.""" + return is_volume_path(path) + + def filesystem(self, connector: SourceConnector, + path: str) -> tuple[FileSystem, str]: + """A volume path as a filesystem, as the connector's identity.""" + if not is_volume_path(path): + raise ValueError( + f"{path!r} is not a Unity Catalog volume path; Databricks " + f"artifacts go under /Volumes////.") + return volume_filesystem(self._client(connector)), path.rstrip('/') + def __repr__(self) -> str: return f'DatabricksSinkBackend(identity={self._identity!r})' diff --git a/packages/kumo-databricks/src/kumo_databricks/volume.py b/packages/kumo-databricks/src/kumo_databricks/volume.py new file mode 100644 index 000000000..3a7b39e61 --- /dev/null +++ b/packages/kumo-databricks/src/kumo_databricks/volume.py @@ -0,0 +1,240 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Unity Catalog volumes as a pyarrow filesystem, over the Files API.""" + +from __future__ import annotations + +import io +import posixpath +import queue +import threading +from contextlib import contextmanager +from typing import BinaryIO, Dict, Iterator, Optional +from urllib.parse import quote + +from databricks.sdk import WorkspaceClient +from databricks.sdk.errors import DatabricksError, NotFound, PermissionDenied +from kumo_core.chunked_fs import ChunkedFileHandler +from pyarrow.fs import FileSystem, PyFileSystem + +_FILES = '/api/2.0/fs/files' +_READ = 8 << 20 + + +def is_volume_path(path: str) -> bool: + parts = path.rstrip('/').split('/') + return (path.startswith('/Volumes/') and len(parts) >= 5 + and all(parts[1:]) and posixpath.normpath(path) == path.rstrip('/')) + + +def volume_filesystem(client: WorkspaceClient) -> FileSystem: + return PyFileSystem(VolumeHandler(client)) + + +@contextmanager +def _translated(path: str) -> Iterator[None]: + """Nothing is chained, and only SDK errors keep their text: transport + errors from presigned uploads carry the signed URL.""" + from kumo_core.sources.diagnostics import safe_summary + + try: + yield + except NotFound as error: + raise FileNotFoundError(f'{path}: {error}') from None + except PermissionDenied as error: + raise PermissionError(f'{path}: {error}') from None + except DatabricksError as error: + raise OSError(f'{path}: {error}') from None + except Exception as error: + raise OSError(f'{path}: {safe_summary(error)}') from None + + +class VolumeHandler(ChunkedFileHandler): + """``/Volumes////...`` through ``client.files``. + + Parts stay under the Files API's 5 GiB single-upload cap. + """ + part_bytes = 4 << 30 + plain_bytes = 256 << 20 + + def __init__(self, client: WorkspaceClient) -> None: + self._client = client + + def __eq__(self, other: object) -> bool: + return (isinstance(other, VolumeHandler) + and other._client is self._client) + + def __ne__(self, other: object) -> bool: + return not self == other + + def get_type_name(self) -> str: + return 'databricks-volume' + + def _size(self, path: str) -> Optional[int]: + path = path.rstrip('/') + with _translated(path): + try: + return int(self._client.files.get_metadata(path).content_length) + except NotFound: + return None + + def _list(self, directory: str) -> Dict[str, int]: + directory = directory.rstrip('/') + with _translated(directory): + try: + entries = list( + self._client.files.list_directory_contents(directory)) + except NotFound: + return {} + found: Dict[str, int] = {} + for entry in entries: + path = entry.path.rstrip('/') + if entry.is_directory: + found.update(self._list(path)) + else: + found[path] = int(entry.file_size) + return found + + def _put(self, local: str, path: str) -> None: + with _translated(path), open(local, 'rb') as contents: + self._client.files.upload(path, contents, overwrite=True) + + def _stream(self, path: str) -> BinaryIO: + return _Upload(self._client, path) + + def _part(self, path: str, size: int) -> _Part: + return _Part(self._client, path) + + def _remove(self, path: str, *, tree: bool) -> None: + from kumo_databricks.staging import clear_volume + + path = path.rstrip('/') + with _translated(path): + if tree: + clear_volume(self._client, path) + return + try: + self._client.files.delete(path) + except NotFound: + pass + + +class _Part: + """Serves consecutive reads from one open-ended ranged download.""" + def __init__(self, client: WorkspaceClient, path: str) -> None: + self._client, self._path = client, path + self._stream = None + self._at = 0 + + def read_at(self, offset: int, length: int) -> bytes: + if length <= 0: + return b'' + if self._stream is None or offset != self._at: + self._open(offset) + with _translated(self._path): + pieces, got = [], 0 + while got < length: + piece = self._stream.read(min(_READ, length - got)) + if not piece: + break + pieces.append(piece) + got += len(piece) + self._at += got + return pieces[0] if len(pieces) == 1 else b''.join(pieces) + + def _open(self, offset: int) -> None: + self.close() + with _translated(self._path): + response = self._client.api_client.do( + 'GET', f'{_FILES}{quote(self._path)}', raw=True, + response_headers=['content-range'], headers={ + 'Accept': 'application/octet-stream', + 'Range': f'bytes={offset}-' + }) + self._stream = response['contents'] + served = response.get('content-range') or '' + if not served.startswith(f'bytes {offset}-'): + self.close() + raise OSError(f'{self._path}: asked for bytes from {offset}, ' + f'got {served or "the whole file"}') + self._stream.set_chunk_size(_READ) + self._at = offset + + def close(self) -> None: + if self._stream is not None: + self._stream.close() + self._stream = None + + +class _Feed(io.RawIOBase): + def __init__(self) -> None: + self.pieces: queue.Queue[Optional[bytes]] = queue.Queue(maxsize=4) + self._pending = b'' + self._done = False + + def readable(self) -> bool: + return True + + def readinto(self, buffer) -> int: + while not self._pending and not self._done: + piece = self.pieces.get() + if piece is None: + self._done = True + else: + self._pending = piece + size = min(len(buffer), len(self._pending)) + buffer[:size] = self._pending[:size] + self._pending = self._pending[size:] + return size + + +class _Upload(io.RawIOBase): + """One ``files.upload`` fed from writes; ``close`` commits it.""" + def __init__(self, client: WorkspaceClient, path: str) -> None: + self._path = path + self._feed = _Feed() + self._error: Optional[BaseException] = None + source = io.BufferedReader(self._feed, 8 << 20) + + def upload() -> None: + try: + with _translated(path): + client.files.upload(path, source, overwrite=True) + except BaseException as error: + self._error = error + + self._thread = threading.Thread(target=upload, daemon=True) + self._thread.start() + + def writable(self) -> bool: + return True + + def write(self, data) -> int: + piece = bytes(data) + self._put(piece) + return len(piece) + + def _put(self, piece: Optional[bytes]) -> None: + while self._thread.is_alive(): + try: + self._feed.pieces.put(piece, timeout=1) + return + except queue.Full: + continue + self._thread.join() + if self._error is not None: + raise self._error + if piece is not None: + raise OSError(f'{self._path}: the upload stopped reading') + + def close(self) -> None: + if self.closed: + return + try: + self._put(None) + self._thread.join() + if self._error is not None: + raise self._error + finally: + super().close() diff --git a/packages/kumo-databricks/tests/fakes.py b/packages/kumo-databricks/tests/fakes.py index a3f941b0a..b629e404d 100644 --- a/packages/kumo-databricks/tests/fakes.py +++ b/packages/kumo-databricks/tests/fakes.py @@ -8,9 +8,10 @@ import io from dataclasses import dataclass from typing import Any, Callable, Mapping +from urllib.parse import unquote import pyarrow as pa -from databricks.sdk.errors import NotFound +from databricks.sdk.errors import BadRequest, NotFound from databricks.sql.exc import DataError from sources.recording import RecordingConnection, RecordingCursor @@ -197,7 +198,85 @@ def delete_directory(self, path: str) -> None: return None +class FakeVolume(FakeFiles): + """A volume with real directories, uploads and ranged downloads.""" + def __init__(self, objects: Mapping[str, bytes] | None = None) -> None: + super().__init__(objects) + self.directories: set[str] = set() + self.ignore_range = False + + def _is_directory(self, path: str) -> bool: + return path in self.directories or any( + name.startswith(path + "/") for name in self.objects) + + def list_directory_contents(self, path: str) -> list[FakeEntry]: + path = path.rstrip("/") + if not self._is_directory(path): + raise NotFound(f"No such directory: {path}") + children: dict[str, FakeEntry] = {} + for name in sorted(self.objects) + sorted(self.directories): + if not name.startswith(path + "/"): + continue + child = name[len(path) + 1:].split("/")[0] + full = f"{path}/{child}" + children[full] = (FakeEntry(full, len(self.objects[full])) + if full in self.objects else FakeEntry( + full + "/", 0, is_directory=True)) + return list(children.values()) + + def upload(self, path: str, contents: Any, *, overwrite: bool) -> None: + assert overwrite + self.objects[path] = contents.read() + + def get_metadata(self, path: str) -> Any: + if path.endswith("/"): + raise BadRequest("Bad Request") + if path not in self.objects: + raise NotFound(f"No such file: {path}") + return FakeMetadata(len(self.objects[path])) + + def get_directory_metadata(self, path: str) -> None: + if not self._is_directory(path): + raise NotFound(f"No such directory: {path}") + + def create_directory(self, path: str) -> None: + self.directories.add(path) + + def delete_directory(self, path: str) -> None: + self.directories.discard(path.rstrip("/")) + + def do(self, method: str, path: str, *, raw: bool, + response_headers: list[str], headers: Mapping[str, str]) -> dict: + assert method == "GET" and raw + remote = unquote(path.removeprefix("/api/2.0/fs/files")) + if remote not in self.objects: + raise NotFound(f"No such file: {remote}") + if self.ignore_range: + return {"contents": FakeStreamingResponse(self.objects[remote])} + first, last = headers["Range"].removeprefix("bytes=").split("-") + size = len(self.objects[remote]) + last = last or str(size - 1) + stream = FakeStreamingResponse( + self.objects[remote][int(first):int(last) + 1]) + self.streams.append(stream) + return { + "contents": stream, + "content-range": f"bytes {first}-{last}/{size}" + } + + +@dataclass +class FakeMetadata: + content_length: int + + +class FakeStreamingResponse(TrackedStream): + def set_chunk_size(self, size: int) -> None: + pass + + class FakeWorkspaceClient: """Exposes only the files API the reader uses.""" def __init__(self, files: FakeFiles) -> None: self.files = files + self.api_client = files diff --git a/packages/kumo-databricks/tests/test_volume.py b/packages/kumo-databricks/tests/test_volume.py new file mode 100644 index 000000000..db60e65fd --- /dev/null +++ b/packages/kumo-databricks/tests/test_volume.py @@ -0,0 +1,180 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace + +import pyarrow as pa +import pyarrow.parquet as pq +import pytest +from databricks.sdk.errors import PermissionDenied +from fakes import FakeVolume, FakeWorkspaceClient +from kumo_databricks.sink import DatabricksSinkBackend +from kumo_databricks.volume import VolumeHandler +from pyarrow.fs import FileType, PyFileSystem + +ROOT = '/Volumes/main/analytics/artifacts' +BODY = bytes(range(256)) * 10 + + +def filesystem(objects=None, part_bytes=None): + volume = FakeVolume(objects) + handler = VolumeHandler(FakeWorkspaceClient(volume)) + if part_bytes is not None: + handler.part_bytes = part_bytes + return PyFileSystem(handler), volume + + +def write(fs, path, body, piece=700): + with fs.open_output_stream(path) as writer: + for start in range(0, len(body), piece): + writer.write(body[start:start + piece]) + + +def test_reads_cross_part_boundaries(): + fs, _ = filesystem(part_bytes=1000) + write(fs, f'{ROOT}/run/big.bin', BODY) + with fs.open_input_file(f'{ROOT}/run/big.bin') as reader: + assert reader.read_at(300, 900) == BODY[900:1200] + assert reader.read_at(10, 1200) == BODY[1200:1210] + + +def test_sized_writes_stay_plain_up_to_a_part(): + from kumo_core.object_store_sink import open_sized + + fs, volume = filesystem(part_bytes=1000) + fs.handler.plain_bytes = 100 + for name, body in (('plain', BODY[:1000]), ('parts', BODY[:1500])): + with open_sized(fs, f'{ROOT}/{name}', len(body)) as writer: + writer.write(body) + with fs.open_input_stream(f'{ROOT}/{name}') as reader: + assert reader.readall() == body + assert volume.objects[f'{ROOT}/plain'] == BODY[:1000] + assert f'{ROOT}/parts' not in volume.objects + assert f'{ROOT}/parts.chunks/_SUCCESS' in volume.objects + + +def test_a_failed_upload_raises_without_the_signed_url(): + from kumo_core.object_store_sink import open_sized + + fs, volume = filesystem() + signed = 'https://bucket.example/part?X-Amz-Signature=secret' + + def drop(path, contents, *, overwrite): + contents.read(100) + raise ConnectionError(f'Max retries exceeded with url: {signed}') + + volume.upload = drop + with pytest.raises(OSError) as raised: + with open_sized(fs, f'{ROOT}/x', len(BODY)) as writer: + for start in range(0, len(BODY), 64): + writer.write(BODY[start:start + 64]) + assert 'Signature' not in str(raised.value) + assert raised.value.__cause__ is None + + +def test_a_download_that_ignores_the_range_is_refused(): + fs, volume = filesystem({f'{ROOT}/a': BODY}) + volume.ignore_range = True + with fs.open_input_file(f'{ROOT}/a') as reader: + with pytest.raises(OSError, match='asked for bytes from 100'): + reader.read_at(10, 100) + + +def test_directories_and_missing_paths(): + fs, volume = filesystem({ + f'{ROOT}/s/kumo_snapshot.json': b'{}', + f'{ROOT}/s/index/sub/b': b'de', + }) + kinds = [info.type for info in fs.get_file_info( + [f'{ROOT}/s/kumo_snapshot.json', f'{ROOT}/s/index/', f'{ROOT}/s/no'])] + assert kinds == [FileType.File, FileType.Directory, FileType.NotFound] + with pytest.raises(FileNotFoundError): + fs.open_input_stream(f'{ROOT}/s/no') + fs.delete_dir_contents(f'{ROOT}/s/') + assert volume.objects == {} + + +def test_refused_upload_raises_permission_error(): + fs, volume = filesystem() + + def refuse(path, contents, *, overwrite): + raise PermissionDenied('no WRITE VOLUME') + + volume.upload = refuse + with pytest.raises(PermissionError): + write(fs, f'{ROOT}/x', b'data') + + +def test_sink_filesystem_only_takes_volume_paths(): + sink = DatabricksSinkBackend() + volume = FakeVolume() + sink._client = lambda connector: FakeWorkspaceClient(volume) + connector = SimpleNamespace() + fs, path = sink.filesystem(connector, f'{ROOT}/exports/') + assert path == f'{ROOT}/exports' + write(fs, f'{path}/x', b'1') + assert volume.objects == {f'{path}/x': b'1'} + assert sink.holds_files(connector, f'{ROOT}/exports/') + for bad in ('s3://bucket/x', '/Volumes/main/analytics', 'users_scores', + f'{ROOT}/../other/x', 'dbfs:/Volumes/main/analytics/v'): + assert not sink.holds_files(connector, bad) + with pytest.raises(ValueError, match='volume path'): + sink.filesystem(connector, bad) + + +def test_bundle_publish_and_fetch_through_a_volume_path(tmp_path): + sink = DatabricksSinkBackend() + volume = FakeVolume() + sink._client = lambda connector: FakeWorkspaceClient(volume) + bundle = tmp_path / 'bundle' + (bundle / 'model').mkdir(parents=True) + (bundle / 'training-spec.json').write_text('{}') + (bundle / 'model' / 'weights.pt').write_bytes(BODY) + published = sink.publish_model(SimpleNamespace(), str(bundle), + path=f'{ROOT}/models/b1') + assert published.uri == f'{ROOT}/models/b1' + assert volume.objects[f'{ROOT}/models/b1/model/weights.pt'] == BODY + fetched = tmp_path / 'fetched' + sink.fetch_model(SimpleNamespace(), str(fetched), path=f'{ROOT}/models/b1') + assert (fetched / 'model' / 'weights.pt').read_bytes() == BODY + with pytest.raises(ValueError, match='Nothing is published'): + sink.fetch_model(SimpleNamespace(), str(fetched), + path=f'{ROOT}/models/none') + + +def test_predictions_and_frames_written_to_a_volume_path(tmp_path): + sink = DatabricksSinkBackend() + volume = FakeVolume() + sink._client = lambda connector: FakeWorkspaceClient(volume) + part = tmp_path / 'part-0.parquet' + pq.write_table(pa.table({'users': [1, 2], 'score': [0.1, 0.9]}), part) + written = sink.write_files(SimpleNamespace(), f'{ROOT}/predictions/', + [str(part)], export_id='run-1') + assert written == f'{ROOT}/predictions/' + names = {path.rsplit('/', 1)[-1] for path in volume.objects} + assert {'part-0.parquet', '_SUCCESS'} <= names + frame = pa.table({'users': [3]}).to_pandas() + assert sink.write_table(SimpleNamespace(), f'{ROOT}/one', + frame) == f'{ROOT}/one.parquet' + + +def test_an_occupied_volume_directory_is_refused_before_scoring(tmp_path): + from kumo_core.artifact_export import ArtifactDestinationConflict + from kumoapi.jobs import WriteMode + + sink = DatabricksSinkBackend() + volume = FakeVolume() + sink._client = lambda connector: FakeWorkspaceClient(volume) + target = f'{ROOT}/predictions/' + sink.preflight_write_files(SimpleNamespace(), target) + part = tmp_path / 'part-0.parquet' + pq.write_table(pa.table({'users': [1]}), part) + sink.write_files(SimpleNamespace(), target, [str(part)], export_id='r1') + + with pytest.raises(ArtifactDestinationConflict, match='/predictions/'): + sink.preflight_write_files(SimpleNamespace(), target) + sink.preflight_write_files(SimpleNamespace(), target, + write_mode=WriteMode.OVERWRITE) + sink.preflight_write_files(SimpleNamespace(), f'{ROOT}/other/') + sink.preflight_write_files(SimpleNamespace(), 'main.ml.scores') diff --git a/tests/public/Dockerfile.dockerignore b/tests/public/Dockerfile.dockerignore index 4bedfe32e..b5e0341fb 100644 --- a/tests/public/Dockerfile.dockerignore +++ b/tests/public/Dockerfile.dockerignore @@ -14,6 +14,10 @@ !packages/kumo-core/src/kumo_core/artifact_staging.py # Model registries validate bundles through the portable artifact contract. !packages/kumo-core/src/kumo_core/artifacts.py +# Volume and stage sinks hold files through the shared chunked filesystem. +!packages/kumo-core/src/kumo_core/chunked_fs.py +!packages/kumo-core/src/kumo_core/object_store_sink.py +!packages/kumo-core/src/kumo_core/artifact_export.py !packages/kumo-core/src/kumo_core/config.py !packages/kumo-core/src/kumo_core/manifest.py !packages/kumo-core/src/kumo_core/io/__init__.py From 5bcdf0d4add40b14f6d98e77214843cf5342d3f4 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 28 Oct 2026 16:02:54 +0000 Subject: [PATCH 018/278] Record fits in experiment tracking and publish models to registries Co-Authored-By: Claude Opus 5.5 --- examples/mlops/smoke.py | 138 +++ packages/kumo-api/kumoapi/plugin.py | 70 +- .../kumo-api/test/test_plugin_registry.py | 33 + packages/kumo-core/src/kumo_core/artifacts.py | 32 + .../src/kumo_core/automl/auto_trainer.py | 37 +- .../src/kumo_core/automl/optuna_searcher.py | 1 + .../kumo-core/src/kumo_core/automl/runtime.py | 5 +- .../src/kumo_core/automl/searcher.py | 14 +- .../callbacks/train_progress_logger.py | 25 +- .../src/kumo_core/diskgraph/lightning_data.py | 11 +- .../src/kumo_core/instrumentation.py | 10 + .../src/kumo_core/train/formatting.py | 34 +- .../src/kumo_core/training/runtime.py | 4 +- packages/kumo-core/tests/test_artifacts.py | 8 + .../tests/test_prediction_output_columns.py | 23 + .../kumo-core/tests/test_searched_params.py | 38 + .../kumo-core/tests/test_worker_handles.py | 46 + .../src/kumo_databricks/registry.py | 103 ++- .../kumo-databricks/tests/test_registry.py | 96 +- .../kumo-databricks/tests/test_uc_registry.py | 2 +- packages/kumo-engine/docs/mlops.md | 356 ++++++-- packages/kumo-engine/kumo/engine/__init__.py | 2 + packages/kumo-engine/kumo/engine/_events.py | 86 ++ packages/kumo-engine/kumo/engine/_manifest.py | 9 + packages/kumo-engine/kumo/engine/_mlops.py | 93 +- packages/kumo-engine/kumo/engine/_scratch.py | 19 +- packages/kumo-engine/kumo/engine/_tracking.py | 482 ++++++++++ .../kumo/engine/connector/snowflake.py | 8 +- .../kumo/engine/integrations/__init__.py | 3 + .../kumo/engine/integrations/_mlflow_model.py | 217 +++++ .../engine/integrations/_mlflow_model_code.py | 7 + .../engine/integrations/_mlflow_target.py | 230 +++++ .../kumo/engine/integrations/mlflow.py | 859 ++++++++++++++++++ .../kumo/engine/integrations/snowflake.py | 281 ++++++ packages/kumo-engine/kumo/engine/tracking.py | 13 + .../kumo/engine/trainer/__init__.py | 2 + .../kumo/engine/trainer/_provenance.py | 32 +- .../kumo/engine/trainer/_runtime.py | 6 +- .../kumo/engine/trainer/_scoring.py | 81 ++ .../kumo-engine/kumo/engine/trainer/_spec.py | 4 +- .../kumo/engine/trainer/output_config.py | 20 +- .../kumo/engine/trainer/trainer.py | 335 ++++++- packages/kumo-engine/pyproject.toml | 11 + .../test/_snowflake_experiment_fake.py | 192 ++++ packages/kumo-engine/test/test_fit_events.py | 154 ++++ .../kumo-engine/test/test_mlflow_model.py | 516 +++++++++++ .../kumo-engine/test/test_mlflow_target.py | 290 ++++++ .../kumo-engine/test/test_mlflow_tracking.py | 404 ++++++++ packages/kumo-engine/test/test_mlops.py | 24 + .../kumo-engine/test/test_optional_imports.py | 17 + .../test/test_prediction_completion.py | 4 +- .../test/test_prediction_workers.py | 4 +- .../kumo-engine/test/test_scoring_defaults.py | 139 +++ .../test/test_semantic_carry_forward.py | 15 +- .../test/test_snowflake_connector.py | 15 + .../test/test_snowflake_tracking.py | 250 +++++ .../test/test_tracking_contract.py | 400 ++++++++ .../kumo-engine/test/test_trainer_registry.py | 187 +++- packages/kumo-snowflake/pyproject.toml | 8 +- .../src/kumo_snowflake/registry.py | 178 +++- .../kumo-snowflake/tests/test_registry.py | 191 +++- setup.py | 3 + uv.lock | 639 ++++++++++--- 63 files changed, 7187 insertions(+), 329 deletions(-) create mode 100644 examples/mlops/smoke.py create mode 100644 packages/kumo-api/test/test_plugin_registry.py create mode 100644 packages/kumo-core/tests/test_prediction_output_columns.py create mode 100644 packages/kumo-core/tests/test_searched_params.py create mode 100644 packages/kumo-core/tests/test_worker_handles.py create mode 100644 packages/kumo-engine/kumo/engine/_events.py create mode 100644 packages/kumo-engine/kumo/engine/_tracking.py create mode 100644 packages/kumo-engine/kumo/engine/integrations/__init__.py create mode 100644 packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py create mode 100644 packages/kumo-engine/kumo/engine/integrations/_mlflow_model_code.py create mode 100644 packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py create mode 100644 packages/kumo-engine/kumo/engine/integrations/mlflow.py create mode 100644 packages/kumo-engine/kumo/engine/integrations/snowflake.py create mode 100644 packages/kumo-engine/kumo/engine/tracking.py create mode 100644 packages/kumo-engine/kumo/engine/trainer/_scoring.py create mode 100644 packages/kumo-engine/test/_snowflake_experiment_fake.py create mode 100644 packages/kumo-engine/test/test_fit_events.py create mode 100644 packages/kumo-engine/test/test_mlflow_model.py create mode 100644 packages/kumo-engine/test/test_mlflow_target.py create mode 100644 packages/kumo-engine/test/test_mlflow_tracking.py create mode 100644 packages/kumo-engine/test/test_scoring_defaults.py create mode 100644 packages/kumo-engine/test/test_snowflake_tracking.py create mode 100644 packages/kumo-engine/test/test_tracking_contract.py diff --git a/examples/mlops/smoke.py b/examples/mlops/smoke.py new file mode 100644 index 000000000..47e152564 --- /dev/null +++ b/examples/mlops/smoke.py @@ -0,0 +1,138 @@ +# 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:// --client-id \ + --catalog main --schema ml --warehouse --experiment /Users//kumo-smoke + python smoke.py sagemaker --arn --role-arn + python smoke.py snowflake --database ML --schema RUNS --warehouse \ + --account --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=connector, + experiment_name='KUMO_SMOKE'), None + + +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')) + cfg = kumo.MLOps( + args.artifact_root or str(work / 'root'), tracking=tracker, + registry=None if model_name is None else ModelRegistryLocation( + model_name, alias='candidate'), 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 {tracker.run_id}, bundle {result.id}, ' + f'registered {result.registered_model}') + if model_name is not None: + from kumo.engine.integrations.mlflow import load_model + + loaded = load_model(f'models:/{model_name}@candidate', + tracking=tracker) + 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() diff --git a/packages/kumo-api/kumoapi/plugin.py b/packages/kumo-api/kumoapi/plugin.py index b649e302b..018755408 100644 --- a/packages/kumo-api/kumoapi/plugin.py +++ b/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`` 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 ' + 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. diff --git a/packages/kumo-api/test/test_plugin_registry.py b/packages/kumo-api/test/test_plugin_registry.py new file mode 100644 index 000000000..52ca7a69f --- /dev/null +++ b/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' + } diff --git a/packages/kumo-core/src/kumo_core/artifacts.py b/packages/kumo-core/src/kumo_core/artifacts.py index 10b33ee11..d9f57fa41 100644 --- a/packages/kumo-core/src/kumo_core/artifacts.py +++ b/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 + #: Every engine artifact's manifest file; underscore-prefixed so Parquet #: readers skip it. PROVENANCE_FILE = '_provenance.json' @@ -382,6 +386,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' @@ -479,6 +503,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. diff --git a/packages/kumo-core/src/kumo_core/automl/auto_trainer.py b/packages/kumo-core/src/kumo_core/automl/auto_trainer.py index 12ebb575b..3dac36e15 100644 --- a/packages/kumo-core/src/kumo_core/automl/auto_trainer.py +++ b/packages/kumo-core/src/kumo_core/automl/auto_trainer.py @@ -12,6 +12,7 @@ import math import os import os.path as osp +import time import traceback import warnings from dataclasses import asdict, dataclass, fields, is_dataclass @@ -39,7 +40,7 @@ AutoTrainerProgress, AutoTrainingProgressLogger, ) -from kumo_core.instrumentation import TrainingInstrumentation +from kumo_core.instrumentation import TrainingInstrumentation, format_error from kumo_core.io import filesystem from kumo_core.io.json import load_json, save_json from kumo_core.job_config import JobConfig, TrainStatus @@ -460,6 +461,7 @@ def _train( metrics_client_job_name: str, pruning: Callback | None = None, status: TrainStatus | None = None, + trial: int | None = None, ) -> list[float]: r"""Config-driven training function. @@ -715,7 +717,7 @@ def _train( trainer.pl_trainer.callbacks.append(pruning) # Add progress logging callback: - self.training_progress_logger.set_current_trainer(trainer) + self.training_progress_logger.set_current_trainer(trainer, trial) trainer.pl_trainer.callbacks.append(self.training_progress_logger) # Save AutoTrainer YAML (Trainer YAML will be saved within @@ -849,7 +851,7 @@ def _run( logger.info(f"Initialized '{self.searcher}'") self.training_progress_logger = AutoTrainingProgressLogger( - self.log_dir, self.num_experiments) + self.log_dir, self.num_experiments, on_event=self._emit) tune_mode = MetricStore.tune_mode(self.tune_metric) metrics_job_name = self.runtime.metrics_job_name() @@ -874,12 +876,17 @@ def _run( logger.info(f"Starting AutoML experiment " f"{i + 1}/{self.num_experiments}") + self._emit('trial_started', trial=i + 1, params={ + **self.searcher.sampled_params, + 'attempt': j + 1, + }) batch_size = cfg.optimization.batch_size[0] with self.instrumentation.timing( "auto_trainer.run.status_add_trial", trial=i + 1): status.add_trial(batch_size) status.save(self.home_dir) + trial_start = time.monotonic() with self.instrumentation.timing( "auto_trainer.run.train_trial", trial=i + 1): val_performance = self._train( @@ -887,7 +894,9 @@ def _run( metrics_client_job_name=metrics_job_name, pruning=pruning, status=status, + trial=i + 1, ) + trial_seconds = time.monotonic() - trial_start best_performance = _eval_fn[tune_mode](val_performance) logger.info(f"AutoML experiment " f"{i + 1}/{self.num_experiments} finished " @@ -904,9 +913,21 @@ def _run( "auto_trainer.run.status_save_success", trial=i + 1): status.save(self.home_dir) + self._emit('trial_ended', trial=i + 1, metrics={ + self.val_tune_metric: best_performance, + 'seconds': trial_seconds, + }, params={ + 'status': 'succeeded', + 'attempt': j + 1, + }) break except RuntimeError as e: + self._emit('trial_ended', trial=i + 1, params={ + 'status': 'failed', + 'attempt': j + 1, + 'error': format_error(e), + }) stacktrace = ''.join( traceback.format_exception(e, e, e.__traceback__)) logger.warning(stacktrace) @@ -959,6 +980,16 @@ def _run( should_init=init_best_trainer): return self.get_best_trainer(should_init=init_best_trainer) + def _emit(self, name: str, **fields: Any) -> None: + event = getattr(self.instrumentation, 'event', None) + if event is None: + return + try: + event(name, **fields) + except Exception: + logger.warning(f"Could not report progress event '{name}'", + exc_info=True) + def _handle_oom_retry( self, cfg: ModelPlan, diff --git a/packages/kumo-core/src/kumo_core/automl/optuna_searcher.py b/packages/kumo-core/src/kumo_core/automl/optuna_searcher.py index 109ead225..f5a23de69 100644 --- a/packages/kumo-core/src/kumo_core/automl/optuna_searcher.py +++ b/packages/kumo-core/src/kumo_core/automl/optuna_searcher.py @@ -66,6 +66,7 @@ def suggest(self) -> PlanT: for _ in range(self.fallback_sampler.max_attempts): self.current_trial = self.study.ask() cfg = self.fallback_sampler.sample_instance() + self.sampled_params = self.fallback_sampler.sampled_params if cfg not in [prev_cfg for prev_cfg, _ in self.state]: return cfg return cfg diff --git a/packages/kumo-core/src/kumo_core/automl/runtime.py b/packages/kumo-core/src/kumo_core/automl/runtime.py index 209680edf..880246456 100644 --- a/packages/kumo-core/src/kumo_core/automl/runtime.py +++ b/packages/kumo-core/src/kumo_core/automl/runtime.py @@ -196,9 +196,10 @@ def resolve_warm_start( encoders; they do not retain hosted AutoML directory layouts. """ if artifacts is not None: - source_schema = load_json(artifacts["prediction_schema"], dict) + from kumo_core.artifacts import bundle_task_type + validate_warm_start_task_type(task.task_type, - TaskType(source_schema["task_type"])) + bundle_task_type(artifacts)) checkpoint = artifacts["selected_checkpoint"] plan_path = artifacts["resolved_model_plan"] encoder_dir = posixpath.dirname(artifacts["encoders"]) diff --git a/packages/kumo-core/src/kumo_core/automl/searcher.py b/packages/kumo-core/src/kumo_core/automl/searcher.py index f8810e7b6..fa22a91d3 100644 --- a/packages/kumo-core/src/kumo_core/automl/searcher.py +++ b/packages/kumo-core/src/kumo_core/automl/searcher.py @@ -39,6 +39,7 @@ def __init__(self, space: PlanT, max_attempts: int = 50) -> None: self.space = space self.max_attempts = max_attempts self.state: list[tuple[PlanT, list[float]]] = [] + self.sampled_params: dict[str, Any] = {} @property def tune_metric(self) -> str: @@ -68,6 +69,7 @@ def sample_value(self, key: str, value: list[Any]) -> Any: def sample_instance(self) -> PlanT: instance = copy.deepcopy(self.space) + self.sampled_params = {} for field in fields(instance): section = getattr(instance, field.name) @@ -76,6 +78,7 @@ def sample_instance(self) -> PlanT: continue assert isinstance(value, list) + searched = len(value) > 1 if isinstance(value[0], PlanMixin): index_of_plan = self.sample_value(key, list(range(len(value)))) @@ -85,15 +88,22 @@ def sample_instance(self) -> PlanT: subvalue, MissingType): continue assert isinstance(subvalue, list) + searched = searched or len(subvalue) > 1 unique_key = f"{key}{index_of_plan}_{subkey}" setattr( plan, subkey, [self.sample_value(unique_key, subvalue)], ) - setattr(section, key, [plan]) + chosen = plan else: - setattr(section, key, [self.sample_value(key, value)]) + chosen = self.sample_value(key, value) + setattr(section, key, [chosen]) + if searched: + self.sampled_params[f"{field.name}.{key}"] = ( + chosen if chosen is None + or isinstance(chosen, (bool, int, float, str)) else + str(chosen)) return instance diff --git a/packages/kumo-core/src/kumo_core/callbacks/train_progress_logger.py b/packages/kumo-core/src/kumo_core/callbacks/train_progress_logger.py index 45f124db5..abdbd7803 100644 --- a/packages/kumo-core/src/kumo_core/callbacks/train_progress_logger.py +++ b/packages/kumo-core/src/kumo_core/callbacks/train_progress_logger.py @@ -5,7 +5,7 @@ import logging import os.path as osp from datetime import datetime, timezone -from typing import Any +from typing import Any, Callable import kumoapi.jobs as jobs_api from kumo_core.io.json import load_json, save_json @@ -140,17 +140,22 @@ class AutoTrainingProgressLogger(MonitoringCallback): # pragma: no cover # Constant to account for overhead not captured in trial estimation: OVERHEAD_CONSTANT = 1.30 - def __init__(self, log_dir: str, num_experiments: int): + def __init__(self, log_dir: str, num_experiments: int, + on_event: Callable[..., None] | None = None): super().__init__() self.path = osp.join(log_dir, self.FILE_NAME) + self.on_event = on_event + self.trial: int | None = None self.progress = load_json(self.path, AutoTrainerProgress, raise_on_missing=False) if self.progress is None: self.progress = AutoTrainerProgress(total_trials=num_experiments) - def set_current_trainer(self, trainer: Any | None): + def set_current_trainer(self, trainer: Any | None, + trial: int | None = None): self.trainer = trainer + self.trial = trial def on_train_start(self, pl_trainer, pl_module): if self.progress.start_time is None: @@ -184,6 +189,20 @@ def on_train_epoch_end(self, pl_trainer, pl_module): save_json(self.progress, self.path) + if self.on_event is not None and self.trial is not None: + try: + progress = self.trainer.training_progress_logger.progress + epoch = progress.num_epochs_completed + metrics = progress.metrics[epoch - 1] + self.on_event('epoch_ended', trial=self.trial, epoch=epoch, + metrics={ + **metrics.train_metrics, + **metrics.validation_metrics + }) + except Exception: + logger.warning('Could not report epoch_ended', + exc_info=True) + def on_validation_end(self, pl_trainer, pl_module): if pl_trainer.state.fn == TrainerFn.VALIDATING: save_json(self.progress, self.path) diff --git a/packages/kumo-core/src/kumo_core/diskgraph/lightning_data.py b/packages/kumo-core/src/kumo_core/diskgraph/lightning_data.py index 3c8e30cec..aa28e8e73 100644 --- a/packages/kumo-core/src/kumo_core/diskgraph/lightning_data.py +++ b/packages/kumo-core/src/kumo_core/diskgraph/lightning_data.py @@ -5,6 +5,7 @@ from __future__ import annotations import ctypes +import errno import itertools import json import logging @@ -235,6 +236,11 @@ def open(cls, process: Any) -> _WorkerHandle: # The duplicated sentinel still gives an authoritative exit # observation for a worker that won the startup/exit race. pass + except OSError as error: + # Sandboxes such as SageMaker hosting refuse pidfd_open. + if error.errno not in (errno.ENOSYS, errno.EPERM): + os.close(sentinel_fd) + raise except BaseException: try: os.close(sentinel_fd) @@ -261,9 +267,8 @@ def send_signal(self, sig: int) -> None: try: if self.pidfd is not None: _pidfd_send_signal(self.pidfd, sig) - elif not sys.platform.startswith('linux'): - # Non-Linux development fallback. Production Linux/Python 3.14 - # uses pidfds, avoiding numeric-PID reuse entirely. + else: + # Without a pidfd the live sentinel above guards the pid. os.kill(self.pid, sig) except ProcessLookupError: pass diff --git a/packages/kumo-core/src/kumo_core/instrumentation.py b/packages/kumo-core/src/kumo_core/instrumentation.py index ab887f8aa..141bcd748 100644 --- a/packages/kumo-core/src/kumo_core/instrumentation.py +++ b/packages/kumo-core/src/kumo_core/instrumentation.py @@ -24,6 +24,9 @@ def span(self, name: str, **fields: Any) -> ContextManager[None]: def timing(self, name: str, **fields: Any) -> ContextManager[None]: ... + def event(self, name: str, **fields: Any) -> None: + return None + class NoOpTrainingInstrumentation: def record_metrics_job_name(self, job_name: str) -> None: @@ -37,6 +40,9 @@ def span(self, name: str, **fields: Any) -> Iterator[None]: def timing(self, name: str, **fields: Any) -> Iterator[None]: yield + def event(self, name: str, **fields: Any) -> None: + return None + _instrumentation: TrainingInstrumentation = NoOpTrainingInstrumentation() @@ -54,6 +60,10 @@ def timing(name: str, **fields: Any) -> ContextManager[None]: return _instrumentation.timing(name, **fields) +def format_error(error: BaseException) -> str: + return f'{type(error).__name__}: {error}'[:200] + + def trace() -> Callable[[Callable[..., T]], Callable[..., T]]: def decorate(function: Callable[..., T]) -> Callable[..., T]: @wraps(function) diff --git a/packages/kumo-core/src/kumo_core/train/formatting.py b/packages/kumo-core/src/kumo_core/train/formatting.py index 53ce12c3a..0d5585d73 100644 --- a/packages/kumo-core/src/kumo_core/train/formatting.py +++ b/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}." diff --git a/packages/kumo-core/src/kumo_core/training/runtime.py b/packages/kumo-core/src/kumo_core/training/runtime.py index ba53e5fe1..96d7a6768 100644 --- a/packages/kumo-core/src/kumo_core/training/runtime.py +++ b/packages/kumo-core/src/kumo_core/training/runtime.py @@ -20,6 +20,7 @@ if TYPE_CHECKING: from kumo_core.embedding_models import EmbeddingModelStore + from kumo_core.instrumentation import TrainingInstrumentation class TrainingDataRuntime(Protocol): @@ -206,6 +207,7 @@ def local_training_bindings( run_name: Optional[str] = None, embedding_models: Optional[EmbeddingModelStore] = None, accelerator: str = 'auto', + instrumentation: Optional[TrainingInstrumentation] = None, ) -> Any: r"""Binds one resolved training input to this process. @@ -224,7 +226,7 @@ def local_training_bindings( run_name=run_name, embedding_models=embedding_models, accelerator=accelerator), - instrumentation=NoOpTrainingInstrumentation(), + instrumentation=instrumentation or NoOpTrainingInstrumentation(), resolve_vital_artifacts=resolve_vital_artifact_sources, # The vital bundle lands beside the run. A platform backend would then # upload it to a registry; locally the path is the identity. diff --git a/packages/kumo-core/tests/test_artifacts.py b/packages/kumo-core/tests/test_artifacts.py index 17e5255b0..fb8d3c23f 100644 --- a/packages/kumo-core/tests/test_artifacts.py +++ b/packages/kumo-core/tests/test_artifacts.py @@ -17,6 +17,7 @@ VitalArtifactSources, add_bundle_file, bundle_id_from_manifest, + bundle_task_type, resolve_vital_artifact_sources, verify_model_bundle, ) @@ -263,3 +264,10 @@ def test_bundle_id_ignores_provenance_only(tmp_path: Path) -> None: (bundle / "metrics.json").write_text('{"auc": 0}') add_bundle_file(bundle, "metrics", "metrics.json") assert verify_model_bundle(bundle) != bundle_id + + +def test_bundle_task_type_reads_the_prediction_schema(tmp_path: Path) -> None: + schema = tmp_path / "prediction-schema.json" + schema.write_text('{"task_type": "regression"}') + assert bundle_task_type({"prediction_schema": str(schema)}).value == ( + "regression") diff --git a/packages/kumo-core/tests/test_prediction_output_columns.py b/packages/kumo-core/tests/test_prediction_output_columns.py new file mode 100644 index 000000000..eee36d837 --- /dev/null +++ b/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 diff --git a/packages/kumo-core/tests/test_searched_params.py b/packages/kumo-core/tests/test_searched_params.py new file mode 100644 index 000000000..c59c8390a --- /dev/null +++ b/packages/kumo-core/tests/test_searched_params.py @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json + +from kumo_core.automl.optuna_searcher import TPESearcher +from kumo_core.model_plan import ModelPlan + + +def test_searcher_records_only_varied_fields_as_json(): + space = ModelPlan() + space.training_job.tune_metric = 'mae' + space.optimization.base_lr = [0.1, 0.01] + space.optimization.batch_size = [512] + space.optimization.optimizer = ['Adam', 'AdamW'] + space.neighbor_sampling.num_neighbors = [[16, 8], [8, 4]] + space.model_architecture.channels = [64, 128] + searcher = TPESearcher(space, seed=0) + + def chosen(cfg): + return { + 'optimization.base_lr': cfg.optimization.base_lr[0], + 'optimization.optimizer': cfg.optimization.optimizer[0], + 'model_architecture.channels': cfg.model_architecture.channels[0], + 'neighbor_sampling.num_neighbors': + str(cfg.neighbor_sampling.num_neighbors[0]), + } + + cfg = searcher.suggest() + assert searcher.sampled_params == chosen(cfg) + json.dumps(searcher.sampled_params) + + searcher.max_attempts = 0 + cfg = searcher.suggest() + assert searcher.sampled_params == chosen(cfg) + assert (searcher.sampled_params + is searcher.fallback_sampler.sampled_params) diff --git a/packages/kumo-core/tests/test_worker_handles.py b/packages/kumo-core/tests/test_worker_handles.py new file mode 100644 index 000000000..03e3f454a --- /dev/null +++ b/packages/kumo-core/tests/test_worker_handles.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import errno +import multiprocessing as mp +import signal +import time + +import pytest +from kumo_core.diskgraph import lightning_data + + +@pytest.mark.parametrize('code', [errno.ENOSYS, errno.EPERM]) +def test_workers_are_stopped_where_pidfds_are_refused(monkeypatch, code): + def refused(pid): + raise OSError(code, 'pidfd_open refused') + + monkeypatch.setattr(lightning_data, '_pidfd_open', refused) + process = mp.get_context('spawn').Process(target=time.sleep, args=(60, )) + process.start() + handle = lightning_data._WorkerHandle.open(process) + try: + assert not handle.exited() + handle.send_signal(signal.SIGKILL) + process.join(10) + assert handle.exited() and process.exitcode == -signal.SIGKILL + finally: + handle.close() + process.kill() + process.join() + + +def test_other_pidfd_errors_still_fail(monkeypatch): + def broken(pid): + raise OSError(errno.EBADF, 'bad') + + monkeypatch.setattr(lightning_data, '_pidfd_open', broken) + process = mp.get_context('spawn').Process(target=time.sleep, args=(60, )) + process.start() + try: + with pytest.raises(OSError, match='bad'): + lightning_data._WorkerHandle.open(process) + finally: + process.kill() + process.join() diff --git a/packages/kumo-databricks/src/kumo_databricks/registry.py b/packages/kumo-databricks/src/kumo_databricks/registry.py index d2d88de09..d51e45474 100644 --- a/packages/kumo-databricks/src/kumo_databricks/registry.py +++ b/packages/kumo-databricks/src/kumo_databricks/registry.py @@ -8,9 +8,9 @@ import json import logging import os +import re 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,13 +20,17 @@ 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' +_UC_FORBIDDEN = re.compile(r'[ /\x00-\x1f\x7f]') _LOG = logging.getLogger(__name__) @@ -41,6 +45,11 @@ def _address(connector: SourceConnector, path: Optional[str], schema, name = normalize_table_coordinate(location.model_name, catalog=config.catalog, schema=config.schema_name) + if len(name) > 255 or _UC_FORBIDDEN.search(name): + raise ValueError( + f'Unity Catalog cannot register a model named {name!r}: use at ' + f'most 255 characters with no space, "/", "." or control ' + f"character, e.g. 'users_churn'.") return f'{config.catalog}.{schema}.{name}' @@ -56,6 +65,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], @@ -99,6 +119,20 @@ def _versions(client: WorkspaceClient, name: str) -> Iterator[dict[str, Any]]: continue +def _registered_version(client: WorkspaceClient, name: str, + bundle_id: str) -> Optional[dict[str, Any]]: + """The ready Kumo version of ``name`` that holds ``bundle_id``, if any.""" + try: + for version in _versions(client, name): + tags = _tags(version) + if (version.get('status') == 'READY' and ENGINE_MARKER in tags + and tags.get(BUNDLE_ID_KEY) == bundle_id): + return version + except NotFound: + pass + return None + + def _validate_catalog(catalog: object) -> None: if not isinstance(catalog, dict) or set(catalog) != {'metadata', 'metrics'}: raise ValueError('Invalid Kumo registry metadata or metrics.') @@ -133,7 +167,15 @@ 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]) + existing = _registered_version(client, name, tags[BUNDLE_ID_KEY]) + if existing is not None: + _LOG.info('%s version %s already holds bundle %s; reused it', name, + existing['version'], tags[BUNDLE_ID_KEY]) + return _published(name, existing, metadata=metadata, + metrics=metrics) try: client.registered_models.create(catalog_name=catalog_name, schema_name=schema, name=model) @@ -200,7 +242,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,13 +261,18 @@ 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.') @@ -233,7 +280,7 @@ def fetch_bundle(client: WorkspaceClient, name: str, local_dir: Path, *, 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 +298,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,13 +327,21 @@ 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))) class DatabricksModelRegistryMixin(_DatabricksSessions): """Publish, discover and reload bundles; prediction remains in Kumo.""" + def check_model_location(self, connector: SourceConnector, + location: ModelRegistryLocation) -> None: + """Raise ``ValueError`` unless ``location`` names a model this + connector's catalog and schema can hold.""" + _address(connector, None, location) + def publish_model( self, connector: SourceConnector, @@ -321,15 +379,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]: diff --git a/packages/kumo-databricks/tests/test_registry.py b/packages/kumo-databricks/tests/test_registry.py index 050fcf18e..e0172b877 100644 --- a/packages/kumo-databricks/tests/test_registry.py +++ b/packages/kumo-databricks/tests/test_registry.py @@ -15,6 +15,8 @@ from kumo_databricks.sink import DatabricksSinkBackend from kumoapi.data_source import DatabricksConnectorResourceConfig from kumoapi.plugin import ( + ENGINE_MARKER, + ModelAliasingSink, ModelFetchingSink, ModelListingSink, ModelPublishingSink, @@ -61,6 +63,7 @@ def test_sink_implements_registry_protocols(): assert isinstance(sink, ModelPublishingSink) assert isinstance(sink, ModelFetchingSink) assert isinstance(sink, ModelListingSink) + assert isinstance(sink, ModelAliasingSink) def test_publish_uses_connector_client_and_reports_assigned_version( @@ -97,10 +100,19 @@ def publish(resolved, name, bundle_dir, *, comment, metadata, metrics): @pytest.mark.parametrize( 'name', - ['other.analytics.churn', 'main.other.churn', 'other.churn', 'a.b.c.d', '']) + ['other.analytics.churn', 'main.other.churn', 'other.churn', 'a.b.c.d', '', + 'users churn', 'users/churn', 'x' * 256]) def test_cross_schema_or_invalid_names_are_refused(connector, name): with pytest.raises(ValueError): - registry._address(connector, None, ModelRegistryLocation(name)) + DatabricksSinkBackend().check_model_location( + connector, ModelRegistryLocation(name)) + + +@pytest.mark.parametrize('name', ['users_churn', 'users-churn', + 'main.analytics.Users-Churn']) +def test_unity_catalog_model_names_are_accepted(connector, name): + DatabricksSinkBackend().check_model_location(connector, + ModelRegistryLocation(name)) @pytest.mark.parametrize('kwargs', [ @@ -141,7 +153,7 @@ def test_invalid_fetch_version_fails_before_authentication( def version(number, kumo=True, status='READY'): - tags = {registry._ENGINE_MARKER: '1'} if kumo else {} + tags = {ENGINE_MARKER: '1'} if kumo else {} return dict(version=number, status=status, tags=[dict(key=k, value=v) for k, v in tags.items()]) @@ -155,6 +167,7 @@ def __init__(self): upload=self.upload, list_directory_contents=self.list_directory, download=self.download) self.objects = {} + self.aliases = {} self.calls = [] self.deleted = [] self.versions = [ @@ -173,6 +186,12 @@ def delete(self, name, number): def do(self, method, path, query=None, body=None): self.calls.append((method, path, query, body)) + if '/aliases/' in path: + alias = path.rsplit('/', 1)[-1] + if method == 'PUT': + self.aliases[alias] = body['version_num'] + return {} + return dict(version=self.aliases[alias]) if method == 'POST': assert all(len(tag['value']) <= 256 for tag in body['tags']) self.versions.append( @@ -207,8 +226,11 @@ def list_directory(self, path): if key.startswith(path + '/'): relative = key[len(path) + 1:] child = relative.split('/')[0] - children[child] = SimpleNamespace(path=path + '/' + child, - is_directory='/' in relative) + directory = '/' in relative + # The Files API lists directories with a trailing slash. + children[child] = SimpleNamespace( + path=f"{path}/{child}{'/' if directory else ''}", + is_directory=directory) return children.values() def download(self, path): @@ -485,3 +507,67 @@ def test_invalid_bundle_metadata_keeps_destination(client, bundle, tmp_path, registry.fetch_bundle(client, name='main.analytics.churn', local_dir=destination, version='7') assert (destination / 'weights.bin').read_bytes() == b'previous' + + +def test_bundle_id_is_tagged_listed_and_fetched_by_version_and_alias( + connector, client, bundle, tmp_path, monkeypatch): + lineage = dict(kumo_bundle_id='ab' * 32) + sink = DatabricksSinkBackend() + monkeypatch.setattr(sink, '_client', lambda _: client) + published = sink.publish_model( + connector, str(bundle), location=ModelRegistryLocation('churn'), + metadata=dict(spec_version='1', kumo_engine='0', kumo_run_id='run-1', + **lineage)) + assert registry._tags(client.versions[-1]) == { + 'kumo_engine': '1', + **lineage + } + listed = registry.list_models(client, catalog='main', schema='analytics') + assert next(m for m in listed if m.version == '7').metadata == lineage + + sink.set_model_alias( + connector, ModelRegistryLocation('churn', version=published.version), + 'prod') + assert client.calls[-1] == ( + 'PUT', + '/api/2.1/unity-catalog/models/main.analytics.churn/aliases/prod', None, + dict(version_num=7)) + for location in (ModelRegistryLocation('churn', version='7'), + ModelRegistryLocation('churn', alias='prod')): + fetched = sink.fetch_model(connector, str(tmp_path / 'restored'), + location=location) + assert fetched.version == '7' + assert fetched.metadata == { + 'spec_version': '1', + 'predictive_query': 'PREDICT COUNT(*)', + 'snapshot_id': 'snapshot', + **lineage + } + + +def test_set_alias_needs_a_positive_version(connector, monkeypatch): + sink = DatabricksSinkBackend() + monkeypatch.setattr(sink, '_client', + lambda _: pytest.fail('resolved identity')) + with pytest.raises(ValueError, match='positive integer'): + sink.set_model_alias(connector, + ModelRegistryLocation('churn', version='0'), + 'prod') + + +def test_publishing_a_registered_bundle_again_reuses_its_version( + client, bundle, caplog): + metadata = dict(kumo_bundle_id='ab' * 32) + first = registry.publish_bundle(client, 'main.analytics.churn', bundle, + metadata=metadata) + posts = sum(call[0] == 'POST' for call in client.calls) + with caplog.at_level('INFO', logger=registry.__name__): + again = registry.publish_bundle(client, 'main.analytics.churn', + bundle, metadata=metadata) + assert again == first and again.version == '7' + assert sum(call[0] == 'POST' for call in client.calls) == posts + assert 'already holds bundle' in caplog.text + other = registry.publish_bundle(client, 'main.analytics.churn', bundle, + metadata=dict(kumo_bundle_id='cd' * 32)) + assert sum(call[0] == 'POST' for call in client.calls) == posts + 1 + assert other.metadata == dict(kumo_bundle_id='cd' * 32) diff --git a/packages/kumo-databricks/tests/test_uc_registry.py b/packages/kumo-databricks/tests/test_uc_registry.py index 8866cf80f..8acca5647 100644 --- a/packages/kumo-databricks/tests/test_uc_registry.py +++ b/packages/kumo-databricks/tests/test_uc_registry.py @@ -76,7 +76,7 @@ def handle_request(self): return self.send(dict(registered_models=[model])) if path.endswith('/models/versions') and self.command == 'POST': assert payload['source'].startswith('file:') - assert payload['tags'] == [{'key': 'kumo_engine', 'value': '1'}] + assert {'key': 'kumo_engine', 'value': '1'} in payload['tags'] versions[identity] = dict(version=1, status='PENDING_REGISTRATION', tags=payload['tags']) diff --git a/packages/kumo-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index 338c2eb73..0a79573eb 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -1,69 +1,32 @@ # MLOps with Kumo Engine -## Quickstart +## Turning MLOps on -With a `graph` and a `pquery` (see the README), turn MLOps on with an artifact root and every snapshot, training table and model is stored under it: +`kumo.MLOps` says where artifacts go and what records them. Use it as a block, which notebooks and scheduled scripts alike can do, or pass it as `mlops=` on a single call: ```python import kumo.engine as kumo -cfg = kumo.MLOps("artifacts") # or s3://..., gs://..., /Volumes/..., (connector, "@DB.SCHEMA.STAGE/kumo") -with kumo.mlops(cfg): - snapshot = graph.snapshot() - result = kumo.Trainer(pquery.suggest_model_plan(snapshot)).fit( - pquery.generate_training_table(snapshot), snapshot=snapshot) -# Later, on any machine that reaches the root: -snapshot = kumo.GraphSnapshot.load_from_connector(cfg, id=result.snapshot_id) -trainer = kumo.Trainer.load_from_connector(cfg, path=f"models/{result.id}") -``` - -## How it fits together - -| artifact | id | stored at (MLOps on) | -|---|---|---| -| snapshot | `snapshot.id` | `/snapshots///` | -| training table | `table.id` | `/training_tables//` | -| model | `result.id` (the bundle id) | `/models//` | - -Each carries a `_provenance.json` manifest that embeds its parents' manifests (see [Lineage](#lineage)). - -## Turning MLOps on - -`kumo.MLOps` says where artifacts go. Use it as a block, or pass it as `mlops=` on a single call: - -```python -cfg = kumo.MLOps("s3://bucket/kumo") +cfg = kumo.MLOps(artifact_root="s3://bucket/kumo") # or (connector, path), e.g. a volume or stage with kumo.mlops(cfg): snapshot = graph.snapshot() -trainer.fit(table, snapshot=snapshot, mlops=cfg) ``` -- **`artifact_root`** is required: a URI (`s3://`, `gs://`, `file://`, or a local path, including `/Volumes/...` on a Databricks cluster) or a `(connector, path)` pair for a store a URI cannot reach, such as `(DatabricksConnector(...), "/Volumes////kumo")` or `(SnowflakeConnector(...), "@DB.SCHEMA.STAGE/kumo")`. A connector whose sink cannot hold files there raises when `kumo.MLOps` is built. -- **`store`** names the artifact kinds stored when they are produced: `snapshot`, `training_table` and `model` (default: all). +- **`artifact_root`** is required. It is a URI (`s3://`, `gs://`, `file://`, or a local path, including `/Volumes/...` on a Databricks cluster) or a `(connector, path)` pair for a store a URI cannot reach. +- **Volumes and stages:** `(DatabricksConnector(...), "/Volumes////kumo")` and `(SnowflakeConnector(...), "@DB.SCHEMA.STAGE/kumo")` work as roots once the installed kumo-databricks or kumo-snowflake sink holds files at those paths; with an older plugin, building `kumo.MLOps` raises. +- **`export`** names the artifact kinds stored under the root when they are produced: `snapshot`, `training_table` and `model` (default: all). +- **`tracking`** and **`registry`** record fits and register models. +- **`on_error`** says what a tracking or registration failure in `fit` does. `"warn"` (default): the call warns, naming the error, and returns as usual. `"raise"`, for scheduled jobs: the call finishes and stores its result first, then raises a `RuntimeError` whose `__cause__` is the failure; after `fit`, `trainer.result` holds the fitted model. - **Precedence:** a call's `mlops=` replaces the enclosing block, and an inner block replaces an outer one. A block covers the calling thread or asyncio task. -- **Off by default.** Without a block or `mlops=`, every call behaves as it always has, and plain exports (`export(config)`, `export_config=`) never write to the artifact root. +- **Off by default.** Without a block or `mlops=`, every call behaves as it always has. Plain exports (`export(config)`, `export_config=`, `output_config=`) never write to the artifact root. A path export that finds a complete copy of the same artifact at its destination reuses it instead of uploading again (logged at info): a snapshot with the same bytes, a training table or prediction directory with `_SUCCESS` and the same `_KUMO_EXPORT` identity, or a model bundle with the same files and `_provenance.json`. -## Exporting and loading artifacts - -Every artifact has one way out, `.export(OutputConfig(...))`, and one way back, `.load_from_connector(...)`. Pass the `kumo.MLOps` in place of the connector to read from the artifact root, with paths relative to it: - -| Artifact | Export | Load | -|---|---|---| -| Snapshot | `snapshot.export(SnapshotOutputConfig(connector, path))`, returns the location | `GraphSnapshot.load_from_connector(connector, path=location)`, or `(cfg, id=snapshot_id)` | -| Training table | `table.export(TrainingTableOutputConfig(connector, name))` | `TrainingTable.load_from_connector(connector, path=name, pquery=pquery, snapshot=snapshot)`, or `(cfg, path=f"training_tables/{id}", ...)` | -| Model | `result.export(ModelOutputConfig(connector, path=... or location=...))` | `Trainer.load_from_connector(connector, path=... or location=...)`, or `(cfg, path=f"models/{id}")` | - -`fit(export_config=...)` is shorthand for fitting, then exporting. Loading needs a file destination (S3, GCS, local files, a volume or stage). Exporting again to a destination that already holds a complete copy of the same artifact reuses it (logged at info) instead of uploading again. - -## Reference - -### Lineage +## Lineage Every artifact carries a `_provenance.json` manifest, metadata only. The leading underscore keeps Spark and Hive Parquet readers from treating it as data. A graph snapshot's manifest, next to its index, records each source table's connector, location and version, plus the engine build that wrote it. `graph.snapshot(record_sources=False)` keeps only connector and table names, plus a SHA-256 of each full location so a source moved elsewhere is still noticed, and keeps Custom SQL only as a hash. -#### Training tables +### Training tables A generated or caller-supplied training table carries `_provenance.json` in its staging directory: - the query, the effective generation plan and the task type; @@ -71,11 +34,11 @@ A generated or caller-supplied training table carries `_provenance.json` in its - the snapshot it was generated from (that snapshot's manifest, embedded); - for `TrainingTable.from_source_table`, the source table's connector, location and version. -`TrainingTable.id` is deterministic: the same snapshot, query, plan, split seed and engine build give the same id. A caller-supplied source without a version gets a random id. An export to files (`TrainingTable.export`) writes the manifest next to the Parquet parts. `TrainingTable.load_from_connector(connector, path=..., pquery=..., snapshot=...)` materializes an export again. +`TrainingTable.id` is deterministic: the same snapshot, query, plan, split seed and engine build give the same id. A caller-supplied source without a version gets a random id. An export to files (`TrainingTable.export`) writes the manifest next to the Parquet parts. `TrainingTable.load_from_connector(connector, path=..., pq=..., snapshot=...)` materializes an export again. With MLOps on, each table is stored once under `/training_tables//`. -#### Models +### Models Every fitted bundle carries `_provenance.json` (`ModelManifest`): - its `id`, the bundle id (a content hash that leaves the manifest file out); @@ -88,6 +51,75 @@ With MLOps on, each bundle is stored under `/models//`. A compl The bundle also carries the published model's holdout predictions (`holdout.parquet`, with its truncation flag), so `result.holdout_df()` works on a model loaded with `Trainer.load` or `Trainer.load_from_connector` (a path or a registry), not only on the result of `fit`. Bundles fitted before this have no copy, and `holdout_df()` raises `FileNotFoundError` for them. +### Scoring defaults + +A model can carry the prediction arguments `predict` uses when a call omits them: + +```python +from kumo.engine import ModelOutputConfig, ScoringDefaults + +tuned = result.with_scoring_defaults( + ScoringDefaults(binary_classification_threshold=0.37)) +tuned.predict(snapshot) # thresholds at 0.37 +tuned.scoring_defaults # ScoringDefaults(binary_classification_threshold=0.37, ...) +``` + +`with_scoring_defaults` returns a `TrainingResult` for a copy of the bundle with its own id. The copy's manifest embeds the original's in `parents.derived_from`. The threshold applies to binary and multilabel classification, `num_classes_to_return` to multiclass, multilabel and link prediction. Explicit `predict` arguments and a query's `TOP k` win; `ScoringDefaults()` removes the defaults. `Trainer.export(config)` publishes a fitted or loaded bundle like `fit(export_config=...)`. + +### Registering models + +With `registry=ModelOutputConfig(connector, location=...)`, every fit in the block registers its model in the connector's registry (Snowflake Model Registry, or Unity Catalog through a `DatabricksConnector`) after it is stored and tracked: + +```python +from kumo.engine import ModelOutputConfig, ModelRegistryLocation + +cfg = kumo.MLOps("s3://bucket/kumo", registry=ModelOutputConfig( + registry, location=ModelRegistryLocation("CHURN_GNN", alias="challenger"))) +with kumo.mlops(cfg): + result = trainer.fit(training_table, snapshot=snapshot) +result.registered_model # PublishedModel(name='CHURN_GNN', version='V3', uri=...) + +# Scoring jobs load whichever version an alias points at. +trainer = kumo.Trainer.load_from_connector(registry, location=ModelRegistryLocation( + "CHURN_GNN", alias="challenger")) +``` + +- Registration is independent of `export`. If it fails, `fit` warns with the error and returns the result with `registered_model=None` (with `on_error="raise"`, it raises instead), and a tracked fit run ends failed with `kumo.error`; an invalid alias or a connector that cannot publish models raises before training, and an invalid alias or model name raises when `kumo.MLOps` is built. Snowflake model names are Snowflake identifiers (a letter or underscore, then letters, digits, underscores or `$`, e.g. `USERS_CHURN`; or a double-quoted name); Unity Catalog model names are at most 255 characters with no space, `/`, `.` or control character, optionally qualified with the connector's own catalog and schema. +- `result.export(ModelOutputConfig(connector, location=...))` on a `TrainingResult`, or `trainer.export(...)` for its latest result, publishes the model with or without a block and also accepts an `alias`; `Trainer.load_from_connector` loads it back. `fit(export_config=...)` is shorthand for fitting then exporting, and its version is `result.published_model`. +- An alias is a letter or underscore followed by letters, digits or underscores; `DEFAULT`, `FIRST`, `LAST`, `LATEST` and `v` are reserved. The alias is set on the new version after publishing, moving it from any version that held it; if that fails, the publish still stands, with a warning. Snowflake allows one alias per version; if moving an alias fails, it is put back on the version that held it. Snowflake's `DEFAULT`, `FIRST` and `LAST` can be loaded but not set. File, S3 and GCS connectors have no aliases. +- Each registered version records only pointers: the `kumo_engine` marker and `kumo_bundle_id`, as version metrics in the Snowflake Model Registry and as version tags in Unity Catalog through a `DatabricksConnector`. (A version registered through MLflow, below, gets a `kumo_bundle_id` version tag only.) The bundle's own `_provenance.json` holds its lineage. `list_models` and `fetch_model` return `kumo_bundle_id` in `PublishedModel.metadata`; `load_from_connector` warns if the fetched bundle differs from it. Through a `DatabricksConnector`, publishing a bundle that a ready version of the model already holds (same `kumo_bundle_id`) returns that version, logged at info, instead of adding another; an alias then moves to it. + +### Promote a Kumo model + +To register every fit as a candidate and promote only a better one, register under a candidate alias (or none) and move the serving alias yourself with the registry's own API. Compare the new model's test metrics (`result.metrics`, also on its tracked fit run) with the current champion's (`load_from_connector(...).result.metrics`): + +```python +champion = kumo.Trainer.load_from_connector( + registry, location=ModelRegistryLocation("CHURN_GNN", alias="champion")) +better = result.metrics["test_auroc"] > champion.result.metrics["test_auroc"] +version = result.registered_model.version +``` + +Unity Catalog (registered through MLflow or a `DatabricksConnector`) and SageMaker managed MLflow: + +```python +from mlflow import MlflowClient + +if better: + MlflowClient().set_registered_model_alias("main.ml.churn_gnn", "champion", version) +``` + +`set_registered_model_alias` moves the alias from the version that held it. For models registered through MLflow, the champion's metrics are also on its run: `client.get_run(client.get_model_version_by_alias(name, "champion").run_id).data.metrics`. + +Snowflake Model Registry, where a version holds at most one alias, so the alias is unset first: + +```sql +ALTER MODEL DB.SCHEMA.CHURN_GNN VERSION champion UNSET ALIAS; +ALTER MODEL DB.SCHEMA.CHURN_GNN VERSION V3 SET ALIAS = champion; +``` + +Scoring jobs that load `alias="champion"` (or `models:/@champion`) pick up the new version on their next run. + ### Moving snapshots between machines To build an index on one machine and train or score on another, export it to a file destination and load it there: @@ -112,11 +144,11 @@ With MLOps on (`graph.snapshot(mlops=...)` or a `kumo.mlops` block), each snapsh ### Scoring without the graph -A scoring job does not need the training job's `Graph` code. It loads a snapshot and the model from the artifact root, then scores: +A scoring job does not need the training job's `Graph` code. It loads the snapshot the training job exported and the model, then scores: ```python -snapshot = GraphSnapshot.load_from_connector(cfg, id=snapshot_id) -trainer = Trainer.load_from_connector(cfg, path=f"models/{bundle_id}") +snapshot = GraphSnapshot.load_from_connector(store, path=location) +trainer = Trainer.load_from_connector(store, path="models/") result = trainer.predict(snapshot, binary_classification_threshold=0.5) ``` @@ -125,3 +157,221 @@ result = trainer.predict(snapshot, binary_classification_threshold=0.5) - It warns when the snapshot was taken of a different graph than the one the model was trained on, naming the schema changes, for example `Schema changes: events edges.` A graph that differs only in table connectors or locations (a dev graph and a prod graph of the same schema) warns that the schema is the same. The model records its training graph from this release on; for an older model the warning cannot say what differs. - `predict(prediction_table, snapshot=...)` scores a prediction table generated beforehand, as before. - `fit` results score too: `result.predict(snapshot, ...)`. + +## Progress events + +`fit(callbacks=[...])` calls each callback with every progress +`kumo.engine.tracking.Event`: `fit_started`, then per search trial +`trial_started`, `epoch_ended` and `trial_ended`, then `fit_ended`. Ended events +carry `params["status"]` and, on failure, a short `params["error"]`. A callback +that raises is logged and not called again; it never fails the fit. For JSON +lines on stdout: + +```python +import json + +trainer.fit(training_table, snapshot=snapshot, + callbacks=[lambda e: print(json.dumps(e.to_dict()))]) +``` + +## Experiment tracking + +`kumo.MLOps(tracking=...)` records each fit in an experiment tracker: one run +per fit, with each hyperparameter-search trial attempt as a child run (MLflow) +or inside the fit run (Snowflake). Tracking is on only +inside a `kumo.mlops(cfg)` block or for a `fit` given `mlops=cfg`. + +```python +from kumo.engine.integrations.mlflow import MlflowTracking # kumo-engine[mlflow] + +cfg = kumo.MLOps(artifact_root="s3://bucket/kumo", + tracking=MlflowTracking(experiment_name="users-spend")) +with kumo.mlops(cfg): + result = trainer.fit(training_table, snapshot=snapshot) +``` + +- A tracking or registration failure does not fail the fit: it is logged, and reported once as a `UserWarning` naming the full error, so both notebooks and job logs show it. With `on_error="raise"` the call raises once instead, after storing its result, with the error as the cause. Only global rank 0 tracks. +- Tracker calls run on a background thread, so a slow or unreachable tracking server does not hold up training. At the end of the fit, `fit` waits at most 30 s for tracking to finish; the rest continues in the background. +- During the fit, MLflow HTTP calls use 2 retries and a 30 s timeout unless `MLFLOW_HTTP_REQUEST_MAX_RETRIES` / `MLFLOW_HTTP_REQUEST_TIMEOUT` are set. +- `MlflowTracking()` records where MLflow's own configuration points when it is built (`MLFLOW_TRACKING_URI`, or `mlflow.set_tracking_uri`): on Databricks this is automatic, and with nothing configured MLflow writes a local store in the working directory. `MlflowTracking(target=...)` names the server and the credentials instead; see [Running from a laptop](#running-from-a-laptop). +- `kumo.MLOps(...)` checks the tracker when it is built: one experiment lookup, plus one registry listing and (Unity Catalog) the model's catalog and schema when it registers. An unreachable server, refused credentials or a missing catalog or schema raise there, before any work. +- The `View run ... at` links MLflow prints do not open on SageMaker; after a tracked fit kumo-engine warns how to open the UI: from SageMaker Studio, or at the URL `aws sagemaker create-presigned-mlflow-tracking-server-url --tracking-server-name ` prints. +- MLflow rejects `@` in metric names, so link-prediction metrics such as `test_map@10` are logged as `test_map_at_10`. +- Per-epoch metrics use step = epoch. After the fit, `tracking.run_id` is the fit run's MLflow run id. +- The fit run is named `fit-` (for example `fit-20261027T061838Z`) unless `run_name` is given; trial runs are `trial--attempt-`. +- The fit's source tables are logged as MLflow run inputs, read from the training table's embedded snapshot manifest (plus a custom training table's source). Databricks tables are logged as Delta tables at the snapshot's Delta version, so Unity Catalog links each table to a model registered from the run; file and object-store sources by URI. Other warehouse tables and Custom SQL sources are not logged as inputs; they stay in the model's `_provenance.json`. + +Snowflake records the fit in Snowflake Experiment Tracking (Snowsight: AI & ML +» Experiments), using a Snowpark session: + +```python +from kumo.engine.integrations.snowflake import SnowflakeTracking # kumo-engine[snowflake] + +cfg = kumo.MLOps(artifact_root="s3://bucket/kumo", + tracking=SnowflakeTracking(session, experiment_name="users_spend")) +with kumo.mlops(cfg): + result = trainer.fit(training_table, snapshot=snapshot) +``` + +- Snowflake has no child runs or tags, so trial attempts are logged in the fit run as `trial-1-attempt-1/` params and metrics, and tags are params of the same name. +- `native_model_lineage=True` links a model the fit publishes through Snowflake lineage. +- Names beyond Snowflake's 200 metrics and 1000 params per run are counted in `kumo.dropped_metrics` / `kumo.dropped_params`. +- If a Snowflake run is already active in the process, the fit is not tracked there; an existing run is never resumed (a taken `run_name` gets a suffix). +- Snowflake Experiment Tracking has no dataset inputs, so source tables are not recorded there. +- The Snowpark session's parameters are left as they are. After the fit, `tracking.run_id` is the fit run's name. + +Runs carry pointers, not lineage: + +| tag | on | value | +|---|---|---| +| `kumo.kind` | every run | `fit` or `trial` | +| `kumo.run_id` | every run | the kumo 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 (including a failed model logging or registration) | 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`. + +Which id is which: + +| id | what it names | where to read it | +|---|---|---| +| kumo run id | one `fit` call | the `kumo.run_id` tag; `tracking.run_id` in the model's `_provenance.json` | +| bundle id | the fitted model's content | `result.id`; the `kumo.bundle_id` tag; `kumo_bundle_id` in the registry; `models//` under the artifact root | +| tracker run id | the run in the tracker | the tracker's `run_id` attribute (`tracking.run_id` on the `MlflowTracking` object: the MLflow run id; Snowflake: the run name) | +| registry version | the registered model version | `result.registered_model.version` | + +To find a bundle's MLflow fit run, search for its tag: ``mlflow.search_runs(filter_string="tags.`kumo.bundle_id` = '' AND tags.`kumo.kind` = 'fit'", search_all_experiments=True)``. + +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 +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`, `check(registry)` (called when `kumo.MLOps` is built) and `log_model(result, snapshot, registry)` (see MLflow models below). 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 bundle the registry already holds under `name` (same `kumo_bundle_id` version tag) is not registered again: its version is returned, the alias moves to it, and an INFO log says it was reused. A bare `ModelRegistryLocation` without such a tracker, a fixed `version`, an invalid alias, or a name the registry (`mlflow.get_registry_uri()`) rejects raises when `kumo.MLOps` is built: + +- SageMaker managed MLflow (a tracking server or MLflow app ARN): letters, digits and hyphens, starting and ending with a letter or digit, at most 57 letters and digits (`^[a-zA-Z0-9](-*[a-zA-Z0-9]){0,56}$`), so no `_` or `.`: `users-churn`. +- Unity Catalog (`databricks-uc`): `catalog.schema.model`, each part at most 255 characters with no space, `/` or control character: `main.ml.users_spend`. On a cluster a shorter name takes the current catalog and schema. +- Other MLflow registries: no extra rule. + +The rules conflict, so no one name fits every registry: SageMaker forbids the `.` Unity Catalog needs between its parts and the `_` Snowflake needs instead of `-`. The examples use `users-churn` for SageMaker and other MLflow registries, `..users_churn` for Unity Catalog and `USERS_CHURN` for Snowflake. + +```python +import mlflow +import pandas as pd +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(pd.DataFrame({"ENTITY": [1, 2, 3]}), params={ + "snapshot": "s3://bucket/kumo/snapshots//", + "prediction_time": "2024-06-01", # temporal queries + "binary_classification_threshold": 0.5, # binary classification only +}) +``` + +- `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 once (`Could not log the model: ...` or `Could not register the model: ...`, with the full error) and still returns its result with `registered_model=None` (with `on_error="raise"`, it raises instead), the run ends `FAILED` with the error in `kumo.error`, 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, the registered version's tags and `result.registered_model.metadata` hold `kumo_bundle_id`. The model's `provenance` artifact is the bundle's `_provenance.json`. Its pip requirement is `kumo-engine[s3,gcs]` of the same version, plus `[text]` when the model encodes a text column. A model trained with a custom embedding model also needs that model's package in the serving environment. +- 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=` (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. +- Input rows have an `ENTITY` column and, optionally, a `TIMESTAMP` column. Scores come back in the input's row order, one per input row (repeated rows included); rows whose entity is not in the snapshot are left out. A temporal query scores entities at a time: pass `params={"prediction_time": ...}` or a `TIMESTAMP` column; without either, `predict` fails with `temporal query: pass params={'prediction_time': ...} or a TIMESTAMP column`. Binary classification needs `params={"binary_classification_threshold": ...}` unless the model's `scoring_defaults` set one. +- Serving needs an image with kumo-engine already installed; MLflow's stock serving containers cannot install it. Serve with `mlflow models serve -m --env-manager local` (on SageMaker endpoints, `-m /opt/ml/model -p 8080`). SageMaker batch transform takes JSON input (`dataframe_split` plus `params`), not CSV, which carries no params. +- GPU real-time endpoints on SageMaker need an inference image whose NVIDIA driver supports CUDA 12.6. The default GPU hosting image ships driver 470, so set a newer `InferenceAmiVersion` on the production variant, for example `al2023-ami-sagemaker-inference-gpu-4-1`. Not yet tested live. +- The `mlflow` extra installs `mlflow-skinny` (3.0.1 or later), the client without the MLflow server's dependencies. 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, use Databricks Runtime 17.3 LTS ML or later: its torch already satisfies kumo-core, so installing kumo-engine keeps the runtime's torch, torchvision and MLflow. kumo-engine accepts pyarrow 16 through 24, so a runtime's pyarrow in that range (19 on 17.3 LTS ML) is kept too; a custom CA bundle for S3 needs pyarrow 21 or later. On 16.4 LTS ML pip replaces torch 2.6 and breaks the preinstalled torchvision. +- 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. On Databricks both work; use the MLflow model unless only kumo-engine jobs load it, since it also serves through `mlflow.pyfunc` (Model Serving, `spark_udf`). 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 | the registry's own API (see [Promote a Kumo model](#promote-a-kumo-model)) | MLflow's `set_registered_model_alias` | +| Databricks | `DatabricksConnector` (Unity Catalog) | Unity Catalog through MLflow | +| SageMaker | not available | SageMaker managed MLflow (`MlflowTracking(target=SageMakerMLflow(arn), ...)`) | +| Snowflake | `SnowflakeConnector` (Snowflake Model Registry) | not available | +| Plain VM | not available (store under the artifact root instead) | any MLflow tracking server | + +#### SageMaker managed MLflow + +1. Create a tracking server once and reuse it; creation takes about 15-25 minutes: + `aws sagemaker create-mlflow-tracking-server --tracking-server-name --artifact-store-uri s3:///mlflow --role-arn ` + (the server role needs read and write on that S3 prefix). +2. Give the identity that runs the job (notebook, VM role or pipeline execution role) `sagemaker-mlflow:*` on the tracking server ARN, and S3 read and write on the artifact store prefix and on the kumo artifact root. +3. Install `pip install "kumo-engine[mlflow,sagemaker]"` (the MLflow client, S3 and the `sagemaker-mlflow` plugin) and pass `MlflowTracking(target=SageMakerMLflow("arn:aws:sagemaker:::mlflow-tracking-server/"), experiment_name=...)`; the registry is the same server. Without an identity, the job's AWS identity (a role or `aws configure`) signs; `MLFLOW_TRACKING_URI` need not be set. The plugin signs and routes with the process-wide tracking URI, so each MLflow call kumo-engine makes sets it to the ARN and restores it afterwards: other threads that call MLflow during a fit see the SageMaker URI. +4. Model names use letters, digits and hyphens only (`users-churn`); `kumo.MLOps` checks this when it is built. +5. The run links MLflow prints do not open on SageMaker. Open the UI from SageMaker Studio or at the URL `aws sagemaker create-presigned-mlflow-tracking-server-url --tracking-server-name ` prints; a tracked fit warns with this command. + +On S3 the artifact root, source tables and prediction outputs are plain `s3://` locations read and written with the job's AWS identity (instance profile, `AWS_PROFILE` or the job role): + +```python +from kumo.engine import PredictionOutputConfig, S3Connector + +data = S3Connector(root_dir="s3:///data") # data["orders.parquet"] reads s3:///data/orders.parquet; data["orders"] a directory of files +cfg = kumo.MLOps("s3:///kumo", + tracking=MlflowTracking(target=SageMakerMLflow(""), experiment_name="users-churn"), + registry=ModelRegistryLocation("users-churn", alias="candidate")) +output = PredictionOutputConfig(S3Connector(), output_paths={"predictions": "s3:///scores/"}) +store = S3Connector() # monitoring: load with absolute s3:// paths +``` + +### Running from a laptop + +`MlflowTracking(target=...)` and `SnowflakeTracking(connector=...)` take the +same identities as the connectors, so a fit on a laptop or any VM can record +to Databricks, SageMaker or Snowflake. An identity with no credentials uses +the platform's ambient chain (a notebook, job or instance role, `aws +configure`, `databricks auth login`); one with credentials uses exactly +those. + +```python +import kumo.engine as kumo +from kumo.engine import ModelRegistryLocation +from kumo.engine.connector import DatabricksConnector, DatabricksIdentity, S3Identity, SnowflakeConnector, SnowflakeIdentity +from kumo.engine.integrations.mlflow import MlflowTracking, SageMakerMLflow +from kumo.engine.integrations.snowflake import SnowflakeTracking + +# Databricks MLflow and Unity Catalog, as a service principal (kumo-engine[mlflow,databricks]). +workspace = DatabricksConnector(catalog="main", schema_name="ml", warehouse_id="", + identity=DatabricksIdentity(host="https://", client_id="", client_secret="")) +cfg = kumo.MLOps("s3:///kumo", + tracking=MlflowTracking(target=workspace, experiment_name="/Users//users_churn"), + registry=ModelRegistryLocation("main.ml.users_churn", alias="candidate")) + +# SageMaker managed MLflow, with an access key or a role to assume (kumo-engine[mlflow,sagemaker]). +cfg = kumo.MLOps("s3:///kumo", + tracking=MlflowTracking(target=SageMakerMLflow("", identity=S3Identity(role_arn="")), + experiment_name="users-churn"), + registry=ModelRegistryLocation("users-churn", alias="candidate")) + +# Snowflake Experiment Tracking, as a key-pair user (kumo-engine[snowflake]). +snowflake = SnowflakeConnector(warehouse="", database="ML", schema_name="RUNS", + identity=SnowflakeIdentity(account="", username="", private_key="")) +cfg = kumo.MLOps("s3:///kumo", tracking=SnowflakeTracking(connector=snowflake, experiment_name="USERS_CHURN")) +``` + +- Each target is checked when `kumo.MLOps` is built, so wrong credentials or a missing catalog, schema or experiment folder parent fail before the fit starts. +- MLflow has no per-client credentials: its Databricks store reads `DATABRICKS_*` variables or a profile, and the SageMaker plugin signs with the process-wide tracking URI. So with explicit credentials, each MLflow call kumo-engine makes sets `DATABRICKS_HOST`, a short-lived `DATABRICKS_TOKEN` (never the client secret) and the process-wide tracking URI `databricks` (Unity Catalog registration downloads the logged model through it), or the process-wide tracking URI and the AWS session, and restores them when the call returns. With explicit credentials, don't call MLflow or Databricks yourself on other threads during a fit. Ambient identities change no credentials; a SageMaker target still sets the process-wide tracking URI for each call, so with a SageMaker target, other threads calling MLflow during a fit see the SageMaker URI. +- SageMaker model artifacts are uploaded and downloaded with the identity's credentials. `load_model(uri, tracking=MlflowTracking(target=..., experiment_name=...))` and `log_model(..., tracking=...)` take the same target. +- `examples/mlops/smoke.py databricks|sagemaker|snowflake` is a manual live check per platform: a small fit recorded and registered with an explicit identity, then loaded back, in a few minutes on a CPU. It is not run in CI. diff --git a/packages/kumo-engine/kumo/engine/__init__.py b/packages/kumo-engine/kumo/engine/__init__.py index db8c0743e..b9e467763 100644 --- a/packages/kumo-engine/kumo/engine/__init__.py +++ b/packages/kumo-engine/kumo/engine/__init__.py @@ -57,6 +57,7 @@ ModelOutputConfig, PredictionOutputConfig, PredictionResult, + ScoringDefaults, SnapshotOutputConfig, Trainer, TrainingResult, @@ -101,6 +102,7 @@ "Trainer", "TrainingResult", "PredictionResult", + "ScoringDefaults", "ModelOutputConfig", "PredictionOutputConfig", "TrainingTableOutputConfig", diff --git a/packages/kumo-engine/kumo/engine/_events.py b/packages/kumo-engine/kumo/engine/_events.py new file mode 100644 index 000000000..7309aab19 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/_events.py @@ -0,0 +1,86 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""Progress events reported by `Trainer.fit`.""" + +import hashlib +import logging +import math +import time +from dataclasses import asdict, dataclass, field +from typing import ( + Any, + Callable, + Dict, + Literal, + Mapping, + Optional, + Sequence, + cast, +) + +from kumo_core.instrumentation import NoOpTrainingInstrumentation + +logger = logging.getLogger(__name__) + +EventName = Literal['fit_started', 'trial_started', 'epoch_ended', + 'trial_ended', 'fit_ended'] + + +@dataclass(frozen=True) +class Event: + r"""One progress event. ``trial`` and ``epoch`` count from 1.""" + + name: EventName + run_id: str + time: float + trial: Optional[int] = None + epoch: Optional[int] = None + metrics: Mapping[str, float] = field(default_factory=dict) + params: Mapping[str, Any] = field(default_factory=dict) + + def to_dict(self) -> Dict[str, Any]: + return asdict(self) + + +def query_sha256(query: str) -> str: + return hashlib.sha256(query.encode()).hexdigest() + + +def _finite_floats(metrics: Mapping[str, Any]) -> Dict[str, float]: + result = {} + for key, value in metrics.items(): + try: + number = float(value) + except (TypeError, ValueError): + continue + if math.isfinite(number): + result[key] = number + return result + + +class EventEmitter(NoOpTrainingInstrumentation): + def __init__(self, run_id: str, + callbacks: Sequence[Callable[[Event], None]]) -> None: + self.run_id = run_id + self._callbacks = list(callbacks) + + def emit(self, name: EventName, *, trial: Optional[int] = None, + epoch: Optional[int] = None, + metrics: Optional[Mapping[str, float]] = None, + params: Optional[Mapping[str, Any]] = None) -> None: + event = Event(name=name, run_id=self.run_id, time=time.time(), + trial=trial, epoch=epoch, + metrics=_finite_floats(metrics or {}), + params=dict(params or {})) + for callback in list(self._callbacks): + try: + callback(event) + except Exception: + logger.warning( + 'Progress callback %r failed; it is disabled for the ' + 'rest of run %s', callback, self.run_id, exc_info=True) + self._callbacks.remove(callback) + + def event(self, name: str, **fields: Any) -> None: + self.emit(cast(EventName, name), **fields) diff --git a/packages/kumo-engine/kumo/engine/_manifest.py b/packages/kumo-engine/kumo/engine/_manifest.py index 64d73bddc..b0ee0b015 100644 --- a/packages/kumo-engine/kumo/engine/_manifest.py +++ b/packages/kumo-engine/kumo/engine/_manifest.py @@ -303,6 +303,7 @@ class TrainingTableManifest(Manifest): class ModelParents(Record): training_table: Optional[TrainingTableManifest] = None warm_start: Optional['ModelManifest'] = None + derived_from: Optional['ModelManifest'] = omit_none() @dataclass(frozen=True, kw_only=True) @@ -312,6 +313,11 @@ class FitTiming(Record): seconds: float +@dataclass(frozen=True, kw_only=True) +class TrackingRef(Record): + run_id: str + + @dataclass(frozen=True, kw_only=True) class ModelManifest(Manifest): """How a model bundle was fitted: from which training table (embedded, @@ -323,6 +329,9 @@ class ModelManifest(Manifest): parents: Optional[ModelParents] = None fit: Optional[FitTiming] = None fit_mode: Optional[str] = None + tracking: Optional[TrackingRef] = omit_none() + + EMBED_OMITS = ('tracking', ) def manifest_bytes(manifest: Manifest) -> bytes: diff --git a/packages/kumo-engine/kumo/engine/_mlops.py b/packages/kumo-engine/kumo/engine/_mlops.py index e38d35329..ac4480fab 100644 --- a/packages/kumo-engine/kumo/engine/_mlops.py +++ b/packages/kumo-engine/kumo/engine/_mlops.py @@ -7,8 +7,20 @@ import contextvars import os import posixpath +import warnings from dataclasses import dataclass, field -from typing import Any, FrozenSet, Iterable, Iterator, Optional, Tuple, Union +from typing import ( + Any, + FrozenSet, + Iterable, + Iterator, + Literal, + Optional, + Tuple, + Union, +) + +from kumoapi.plugin import ModelRegistryLocation, check_model_alias from kumo.engine._manifest import KINDS @@ -31,24 +43,38 @@ def uri(self, relative: str) -> str: @dataclass(frozen=True) class MLOps: - r"""Where artifacts go, for :func:`mlops` blocks and ``mlops=`` - arguments. + r"""Where artifacts go and what records them, for :func:`mlops` blocks + and ``mlops=`` arguments. Args: artifact_root: an ``s3://``, ``gs://``, ``file://`` or local path (``/Volumes/...`` on a Databricks cluster), or a ``(connector, path)`` pair for a store a URI cannot reach, such as a volume through a ``DatabricksConnector`` or a Snowflake stage. + tracking: the experiment tracker, a + ``kumo.engine.tracking.Tracking``. Its ``check(registry)``, if + any, runs here, so a server or registry that is unreachable or + refuses these credentials fails before any work. + registry: where fits register models: a ``ModelOutputConfig`` with a + registry ``location``, or a ``ModelRegistryLocation`` registered + through the MLflow tracker. store: the artifact kinds stored under the root when produced (``snapshot``, ``training_table``, ``model``); ``True`` for all, ``False`` for none. + on_error: what a failure to track a run or register a model does: + ``'warn'`` warns and the call returns; ``'raise'`` raises once + the call has finished and stored its result, with the failure as + the cause. In ``load_from_connector`` calls, an ``MLOps`` stands in for the connector, with paths relative to its artifact root. """ artifact_root: Union[str, Tuple[Any, str]] + tracking: Any = field(default=None, kw_only=True) + registry: Any = field(default=None, kw_only=True) store: Union[bool, Iterable[str]] = field(default=True, kw_only=True) + on_error: Literal['warn', 'raise'] = field(default='warn', kw_only=True) root: _Root = field(init=False, repr=False, compare=False) kinds: FrozenSet[str] = field(init=False, repr=False, compare=False) @@ -60,6 +86,9 @@ def __post_init__(self) -> None: if not self.artifact_root: raise ValueError('MLOps needs an artifact_root.') + if self.on_error not in ('warn', 'raise'): + raise ValueError(f"on_error must be 'warn' or 'raise', not " + f"{self.on_error!r}.") if isinstance(self.artifact_root, str): if self.artifact_root.startswith('@'): raise ValueError( @@ -83,6 +112,28 @@ def __post_init__(self) -> None: f'{list(KINDS)}.') if kinds: _require_writable(root) + from kumo.engine.trainer.output_config import ( + ModelOutputConfig, + check_model_location, + ) + + if isinstance(self.registry, ModelOutputConfig): + check_model_location(self.registry) + 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) + check = getattr(self.tracking, 'check', None) + if callable(check): + check(self.registry if isinstance( + self.registry, ModelRegistryLocation) else None) object.__setattr__(self, 'root', root) object.__setattr__(self, 'kinds', kinds) @@ -156,7 +207,8 @@ def mlops(settings: MLOps) -> Iterator[MLOps]: r"""Turns MLOps on for the calls in the block (in this thread or asyncio task): each snapshot, training table and model is stored under ``settings.artifact_root``. An inner block replaces an outer one; - an ``mlops=`` argument replaces both for its call.""" + an ``mlops=`` argument replaces both for its call. Runs are tracked and + models registered as ``settings`` says.""" if not isinstance(settings, MLOps): raise TypeError('kumo.mlops() takes a kumo.MLOps') token = _CURRENT.set(settings) @@ -172,6 +224,39 @@ def resolve(settings: Optional[MLOps] = None) -> Optional[MLOps]: return settings if settings is not None else _CURRENT.get() +def raise_on_error(settings: Optional[MLOps], error: Optional[BaseException], + call: str, stacklevel: int = 2) -> None: + r"""Reports ``error``, a tracking or registration failure of a finished + ``call``: raises when ``settings.on_error`` is ``'raise'``, else warns + unless `warn_failure` already warned about it.""" + if error is None or settings is None: + return + message = (f'{call} finished, but MLOps tracking or registration ' + f'failed: {full_error(error)}') + if settings.on_error == 'raise': + raise RuntimeError(message) from error + warn_failure(message, error, stacklevel + 1) + + +def full_error(error: BaseException) -> str: + r"""``Type: message``, not shortened.""" + return f'{type(error).__name__}: {error}' + + +def warn_failure(message: str, error: BaseException, stacklevel: int = 2, + also: Optional[BaseException] = None) -> None: + r"""Warns ``message`` (a :class:`UserWarning`) once per ``error``; it + also reports ``also``, an earlier error that caused ``error``.""" + if getattr(error, '_kumo_warned', False): + return + warnings.warn(message, UserWarning, stacklevel=stacklevel + 1) + for reported in (error, also): + try: + setattr(reported, '_kumo_warned', True) + except AttributeError: + pass + + def on_root(source: Any, path: str) -> Tuple[Any, str]: r"""The connector and path a ``load_from_connector(source, path=...)`` reads: an `MLOps` ``source`` reads ``path`` under its artifact root.""" diff --git a/packages/kumo-engine/kumo/engine/_scratch.py b/packages/kumo-engine/kumo/engine/_scratch.py index 7381e5cb5..8a73ce5fa 100644 --- a/packages/kumo-engine/kumo/engine/_scratch.py +++ b/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: diff --git a/packages/kumo-engine/kumo/engine/_tracking.py b/packages/kumo-engine/kumo/engine/_tracking.py new file mode 100644 index 000000000..4160e97aa --- /dev/null +++ b/packages/kumo-engine/kumo/engine/_tracking.py @@ -0,0 +1,482 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""Platform-neutral run records and the `Tracking` protocol.""" + +import dataclasses +import json +import logging +import os +import queue +import threading +import time +import warnings +from dataclasses import dataclass, field +from pathlib import Path +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + Iterator, + List, + Mapping, + Optional, + Protocol, + Sequence, + Tuple, +) + +from kumo_core.artifacts import PROVENANCE_FILE +from kumo_core.instrumentation import format_error + +from kumo.engine._events import Event, _finite_floats +from kumo.engine._mlops import full_error, warn_failure + +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__) + +DROPPED_PARAMS = 'kumo.dropped_params' +DROPPED_METRICS = 'kumo.dropped_metrics' +# Tracker calls run on one background thread so a slow tracking server +# 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) +class FitContext: + r"""What a fit trains, known once its model plan is resolved.""" + + run_id: str + query: str + model_plan: Mapping[str, Any] + inputs: Mapping[str, Mapping[str, Optional[str]]] = field( + default_factory=dict) + + +@dataclass(frozen=True) +class RunRecord: + r"""One tracked run; ``steps`` are per-epoch ``(epoch, value)`` + histories, ``artifacts`` map file names to text and ``inputs`` map + source table names to their ``connector``, ``location``, ``version`` and + ``version_kind``.""" + + name: Optional[str] = None + params: Mapping[str, str] = field(default_factory=dict) + metrics: Mapping[str, float] = field(default_factory=dict) + steps: Mapping[str, Tuple[Tuple[int, float], + ...]] = field(default_factory=dict) + tags: Mapping[str, str] = field(default_factory=dict) + artifacts: Mapping[str, str] = field(default_factory=dict) + inputs: Mapping[str, Mapping[str, Optional[str]]] = field( + default_factory=dict) + + +INPUT_FIELDS = ('connector', 'location', 'version', 'version_kind') + + +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.""" + try: + table = manifest + if getattr(manifest, 'kind', None) == 'model': + table = getattr(_parents(manifest), 'training_table', None) + parents = _parents(table) + snapshot = _parents(getattr(parents, 'snapshot', None)) + tables = dict(getattr(snapshot, 'source_tables', None) or {}) + source = getattr(parents, 'source_table', None) + if source is not None: + tables['training_table'] = source + return { + str(name): {key: _text(getattr(entry, key)) + for key in INPUT_FIELDS} + for name, entry in tables.items() + } + except Exception as error: + logger.warning('Could not read the fit\'s source tables: %s', + format_error(error), exc_info=_debug()) + return {} + + +def _parents(manifest: Any) -> Any: + if manifest is None or getattr(manifest, 'reconstructed', False): + return None + return getattr(manifest, 'parents', None) + + +def _text(value: Any) -> Optional[str]: + return None if value is None else str(value) + + +def run_record(context: FitContext, result: Optional['TrainingResult'] = None, + fit_ended: Optional[Event] = None, *, + tags: Optional[Mapping[str, str]] = None) -> RunRecord: + r"""The fit run's record; ``tags`` are the user's extra tags.""" + metrics: Dict[str, float] = {} + pointers: Dict[str, str] = {} + artifacts = { + 'query.pql': context.query, + 'model_plan.json': json.dumps(dict(context.model_plan), indent=2, + default=str), + } + if result is not None: + metrics.update(result.metrics) + if fit_ended is not None and 'seconds' in fit_ended.metrics: + metrics['seconds'] = fit_ended.metrics['seconds'] + if result.id: + pointers['kumo.bundle_id'] = str(result.id) + provenance = Path(result.bundle_dir) / PROVENANCE_FILE + if provenance.is_file(): + artifacts[PROVENANCE_FILE] = provenance.read_text( + encoding='utf-8') + params = { + key: _param_value(value) + for key, value in _flatten(context.model_plan) + } + return RunRecord( + params=params, + metrics=_finite_floats(metrics), + tags={ + **_user_tags(tags, params), + 'kumo.kind': 'fit', + 'kumo.run_id': context.run_id, + **pointers, + **_error_tag(fit_ended), + }, + artifacts=artifacts, + inputs=dict(context.inputs), + ) + + +def _user_tags(tags: Optional[Mapping[str, Any]], + params: Mapping[str, str]) -> Dict[str, str]: + dropped = sorted(key for key in tags or {} + if key.startswith('kumo.') or key in params) + if dropped: + warnings.warn( + f'Tracking tags {dropped} are reserved and not recorded.', + UserWarning, stacklevel=3) + return { + key: str(value) + for key, value in (tags or {}).items() if key not in dropped + } + + +def _error_tag(ended: Optional[Event]) -> Dict[str, str]: + if ended is None or 'error' not in ended.params: + return {} + return {'kumo.error': str(ended.params['error'])} + + +def trial_record(context: FitContext, events: Sequence[Event]) -> RunRecord: + r"""One trial attempt's record from its events so far.""" + started = next(event for event in events if event.name == 'trial_started') + trial, attempt = started.trial, int(started.params['attempt']) + params: Dict[str, str] = { + key: _param_value(value) + for key, value in started.params.items() + } + metrics: Dict[str, float] = {} + steps: Dict[str, List[Tuple[int, float]]] = {} + ended: Dict[str, str] = {} + for event in events: + if event.name == 'epoch_ended' and event.epoch is not None: + for name, value in event.metrics.items(): + steps.setdefault(name, []).append((event.epoch, value)) + elif event.name == 'trial_ended': + ended = { + 'kumo.status': str(event.params.get('status')), + **_error_tag(event), + } + metrics.update({ + name if name == 'seconds' else f'best_{name}': value + for name, value in event.metrics.items() + }) + return RunRecord( + name=f'trial-{trial}-attempt-{attempt}', + params=params, + metrics=metrics, + steps={ + name: tuple(history) + for name, history in steps.items() + }, + tags={ + 'kumo.kind': 'trial', + 'kumo.run_id': f'{context.run_id}-trial-{trial}-attempt-{attempt}', + 'kumo.parent_run': context.run_id, + **ended, + }, + ) + + +class Tracking(Protocol): + r"""Records one `Trainer.fit` in an experiment tracker. + + Calls come from one background thread. `start` gets the fit run's + record without metrics, `on_trial` a trial attempt's record so far with + 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); ``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``. + """ + def start(self, record: RunRecord) -> None: + ... + + def on_trial(self, record: RunRecord, event: Event) -> None: + ... + + def finish(self, record: Optional[RunRecord], + error: Optional[BaseException]) -> None: + ... + + +class RegistrationError(Exception): + r"""Raised by a tracker's ``log_model`` from the error that registering + the logged model raised.""" + + +class TrackingSession: + def __init__(self, tracking: Optional[Tracking], + warn: bool = True) -> None: + self._tracker = tracking if _rank() == 0 else None + self._warn = warn + self._context: Optional[FitContext] = None + self._trials: Dict[int, List[Event]] = {} + self._fit_ended: Optional[Event] = None + self._jobs: 'queue.Queue[Optional[Callable[[], Any]]]' = queue.Queue( + QUEUE_SIZE) + self._worker: Optional[threading.Thread] = None + self._behind = False + self.error: Optional[BaseException] = None + + @property + def enabled(self) -> bool: + return self._tracker is not None + + def start(self, context: Callable[[], FitContext]) -> None: + if self._tracker is None: + return + try: + value = context() + except Exception as error: + logger.warning('Experiment tracking failed: %s; the fit continues', + full_error(error), exc_info=_debug()) + self._fail(error) + self._tracker = None + return + self._context = value + self._call(lambda tracker: getattr(tracker, 'prepare', _none)()) + self._worker = threading.Thread(target=self._work, + name='kumo-tracking', daemon=True) + self._worker.start() + self._submit(lambda: self._call(lambda tracker: tracker.start( + run_record(value, tags=getattr(tracker, 'tags', None))))) + + def on_event(self, event: Event) -> None: + if self._context is None: + return + if event.name == 'fit_ended': + self._fit_ended = event + if event.trial is None: + return + if event.name == 'trial_started': + self._trials[event.trial] = [] + events = self._trials.setdefault(event.trial, []) + if 'attempt' not in event.params and events: + event = dataclasses.replace(event, params={ + **event.params, 'attempt': events[0].params['attempt'] + }) + events.append(event) + record = trial_record(self._context, events) + 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 (self._tracker is not None and 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 None: + raise RuntimeError( + 'tracking stopped after an earlier failure' + ( + '' if self.error is None else + f' ({full_error(self.error)})')) + 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: + error = failed[0] + verb = 'log' + if isinstance(error, RegistrationError) and error.__cause__: + verb, error = 'register', error.__cause__ + if self._warn: + warn_failure( + f'Could not {verb} the model: {full_error(error)}', + error, 3, also=self.error) + self._fail(error) + return None, error + return (registered or [None])[0], None + + def finish(self, result: Optional['TrainingResult'], + error: Optional[BaseException]) -> None: + worker = self._worker + if self._context is None or worker is None: + return + try: + record: Optional[RunRecord] = run_record(self._context, result, + self._fit_ended) + if error is not None and 'kumo.error' not in record.tags: + record = dataclasses.replace(record, tags={ + **record.tags, 'kumo.error': format_error(error)}) + except Exception as failure: + logger.warning('Could not build the tracking record: %s', + full_error(failure), exc_info=_debug()) + self._fail(failure) + record = None + + def finish_tracker() -> None: + tracker, self._tracker = self._tracker, None + if tracker is not None: + try: + tracker.finish(record, error) + except Exception as failure: + self._fail(failure) + raise + + deadline = time.monotonic() + FINISH_SECONDS + try: + for job in (finish_tracker, None): + self._jobs.put(job, timeout=max(0.0, + deadline - time.monotonic())) + except queue.Full: + pass + worker.join(max(0.0, deadline - time.monotonic())) + if worker.is_alive(): + logger.warning( + 'Experiment tracking did not finish within %ss; it goes on ' + 'in the background and the fit returns', FINISH_SECONDS) + + def _submit(self, job: Callable[[], Any]) -> bool: + if self._behind: + return False + try: + self._jobs.put_nowait(job) + return True + except queue.Full: + self._behind = True + logger.warning('Experiment tracking is falling behind; the rest ' + 'of this fit is not tracked') + return False + + def _work(self) -> None: + while True: + job = self._jobs.get() + if job is None: + return + try: + job() + except Exception as error: + logger.warning('Experiment tracking failed: %s', + full_error(error), exc_info=_debug()) + + def _call(self, call: Callable[[Tracking], Any]) -> None: + tracker = self._tracker + if tracker is None: + return + try: + call(tracker) + except Exception as error: + logger.warning( + 'Experiment tracking failed: %s; it is off for the rest of ' + 'this run', full_error(error), exc_info=_debug()) + self._fail(error) + self._tracker = None + _abandon(tracker) + + def _fail(self, error: BaseException) -> None: + if self.error is None: + self.error = error + + +def _debug() -> bool: + r"""Tracebacks of tracking failures only at DEBUG: the failure itself is + reported once, by the call that tracked it.""" + return logger.isEnabledFor(logging.DEBUG) + + +def _none(*args: Any) -> None: + return None + + +def _abandon(tracker: Tracking) -> None: + try: + tracker.finish(None, None) + except Exception: + logger.debug('Could not end the run of %r', tracker, exc_info=True) + + +def _rank() -> int: + for name in ('RANK', 'LOCAL_RANK'): + if os.environ.get(name, '').strip() not in ('', '0'): + return 1 + return 0 + + +def _flatten(value: Any, prefix: str = '') -> Iterator[Tuple[str, Any]]: + if isinstance(value, Mapping): + for key, item in value.items(): + yield from _flatten(item, f'{prefix}{key}.') + else: + yield prefix[:-1], value + + +def _param_value(value: Any) -> str: + if isinstance(value, str): + return value + return json.dumps(value, default=str) diff --git a/packages/kumo-engine/kumo/engine/connector/snowflake.py b/packages/kumo-engine/kumo/engine/connector/snowflake.py index 0e75a9947..bdfc97387 100644 --- a/packages/kumo-engine/kumo/engine/connector/snowflake.py +++ b/packages/kumo-engine/kumo/engine/connector/snowflake.py @@ -11,7 +11,7 @@ import os from functools import cached_property from pathlib import Path -from typing import Dict, Optional +from typing import Any, Dict, Optional from kumoapi.data_source import ( DataSourceType, @@ -100,6 +100,12 @@ def sink_backend(self) -> ArtifactSink: "snowflake").SnowflakeSinkBackend return SnowflakeSinkBackend(identity=self.identity) + def session(self) -> Any: + r"""The Snowpark session this connector reads through: its identity, + pointed at its warehouse, database and schema, kept alive.""" + with self.source_backend._session(self) as session: + return session + def source_table_for(self, table_name: str) -> SnowflakeSourceTable: return SnowflakeSourceTable( snowflake_connector_id=self.id, diff --git a/packages/kumo-engine/kumo/engine/integrations/__init__.py b/packages/kumo-engine/kumo/engine/integrations/__init__.py new file mode 100644 index 000000000..f2aa5b939 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/integrations/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py b/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py new file mode 100644 index 000000000..da763012f --- /dev/null +++ b/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py @@ -0,0 +1,217 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +import logging +import tempfile +from pathlib import Path +from typing import Any, Dict, Mapping, Optional, Tuple + +import pandas as pd + +from kumo.engine.connector._extras import import_provider +from kumo.engine.connector.base import Connector +from kumo.engine.graph import GraphSnapshot + +_pyfunc = import_provider('mlflow.pyfunc', 'mlflow') +_schema = import_provider('mlflow.types.schema', 'mlflow') +_signature = import_provider('mlflow.models.signature', 'mlflow') + +logger = logging.getLogger(__name__) +_ANCHOR = 'ANCHOR_TIMESTAMP' +_DTYPES = { + 'PREDICTED': 'boolean', + 'SCORE': 'double', + 'TIMESTAMP': 'datetime', + 'False_PROB': 'double', + 'True_PROB': 'double', +} + + +class KumoModel(_pyfunc.PythonModel): + def load_context(self, context: Any) -> None: + from kumo.engine.trainer.trainer import Trainer + + self._trainer = Trainer.load(context.artifacts['bundle']) + self._columns = list(context.model_config['columns']) + self._snapshots: Dict[str, GraphSnapshot] = {} + self._default = context.model_config.get('snapshot') + if self._default: + self._snapshot(self._default) + + def predict(self, context: Any, model_input: pd.DataFrame, + params: Optional[Dict[str, Any]] = None) -> pd.DataFrame: + params = params or {} + if len(model_input) == 0 and not params.get('population'): + return pd.DataFrame(columns=self._columns) + uri = params.get('snapshot') or self._default + if not uri: + raise ValueError("Pass params={'snapshot': }.") + return predict(self._trainer, self._snapshot(uri), model_input, + params) + + def _snapshot(self, uri: str) -> GraphSnapshot: + if uri not in self._snapshots: + connector, path = snapshot_location(uri) + self._snapshots[uri] = GraphSnapshot.load_from_connector( + connector, path=path) + return self._snapshots[uri] + + +def predict(trainer: Any, snapshot: GraphSnapshot, rows: pd.DataFrame, + params: Mapping[str, Any]) -> pd.DataFrame: + from kumo.engine._scratch import scoped_scratch_dir + from kumo.engine.connector.file import FileConnector + + time = params.get('prediction_time') + time = None if time is None else pd.Timestamp(time) + options = dict( + binary_classification_threshold=params.get( + 'binary_classification_threshold'), + num_classes_to_return=params.get('num_classes_to_return'), + ) + _require_inputs(trainer, snapshot, rows, params, time) + trainer.scratch_dir.mkdir(parents=True, exist_ok=True) + with tempfile.TemporaryDirectory(dir=trainer.scratch_dir) as tmp, \ + scoped_scratch_dir(Path(tmp) / 'scratch'): + if len(rows) == 0: + return trainer.predict( + snapshot, prediction_time=None + if time is None else time.to_pydatetime(), **options).frame + rows = rows.rename(columns={'ENTITY': 'ENTITY1'}) + if time is not None and 'TIMESTAMP' not in rows: + rows = rows.assign(TIMESTAMP=time) + (Path(tmp) / 'rows').mkdir() + rows.to_parquet(Path(tmp) / 'rows' / 'rows.parquet', index=False) + entities = FileConnector(root_dir=str(Path(tmp) / 'rows'))[ + 'rows.parquet'] + timed = 'TIMESTAMP' in rows and _temporal(trainer, snapshot) + scored = trainer._predict( + trainer._prediction_table(snapshot, entities, None), + snapshot=snapshot, extra_fields=[_ANCHOR] if timed else [], + **options).frame + return _in_input_order(rows, scored) + + +def _in_input_order(rows: pd.DataFrame, scored: pd.DataFrame) -> pd.DataFrame: + r"""``scored`` in the order of ``rows``, which scoring does not keep; + rows scoring skipped stay out, and repeated rows pair up in order.""" + keys = ['ENTITY'] + left = pd.DataFrame({'ENTITY': rows['ENTITY1'].to_numpy()}) + right = scored[['ENTITY']].copy() + if _ANCHOR in scored: + keys.append(_ANCHOR) + left[_ANCHOR] = _seconds(rows['TIMESTAMP']) + right[_ANCHOR] = _seconds(scored[_ANCHOR]) + try: + left = left.astype({'ENTITY': right['ENTITY'].dtype}) + left['_n'] = left.groupby(keys, sort=False, dropna=False).cumcount() + right['_n'] = right.groupby(keys, sort=False, dropna=False).cumcount() + right['_row'] = range(len(right)) + order = left.merge(right, on=[*keys, '_n'])['_row'] + except (TypeError, ValueError): + order = None + scored = scored.drop(columns=_ANCHOR, errors='ignore') + if order is None or len(order) != len(scored): + logger.warning('Could not match the scores to the input rows; they ' + 'are returned in scoring order: join them on ENTITY') + return scored + return scored.iloc[order.to_numpy()].reset_index(drop=True) + + +def _temporal(trainer: Any, snapshot: GraphSnapshot) -> bool: + from kumo_core.pquery import query_is_temporal + + from kumo.engine.pquery import PredictiveQuery + + return query_is_temporal( + PredictiveQuery._from_snapshot(snapshot, + trainer._training_query)._parsed, + snapshot._native_graph) + + +def _seconds(values: pd.Series) -> pd.Series: + stamps = pd.to_datetime(values) + if stamps.dt.tz is not None: + stamps = stamps.dt.tz_convert(None) + return stamps.astype('datetime64[s]') + + +def _require_inputs(trainer: Any, snapshot: GraphSnapshot, rows: pd.DataFrame, + params: Mapping[str, Any], + time: Optional[pd.Timestamp]) -> None: + from kumo_core.pquery import query_is_temporal + from kumo_core.types import TaskType + + from kumo.engine.pquery import PredictiveQuery + + untimed = len(rows) > 0 and time is None and 'TIMESTAMP' not in rows + unthresholded = ( + params.get('binary_classification_threshold') is None and + trainer.result.scoring_defaults.binary_classification_threshold is None) + if not (untimed or unthresholded): + return + parsed = PredictiveQuery._from_snapshot(snapshot, + trainer._training_query)._parsed + if untimed and query_is_temporal(parsed, snapshot._native_graph): + raise ValueError("temporal query: pass params={'prediction_time': " + "...} or a TIMESTAMP column") + if unthresholded and (TaskType(parsed.task_type) + is TaskType.BINARY_CLASSIFICATION): + raise ValueError("binary classification: pass " + "params={'binary_classification_threshold': ...}") + + +def snapshot_location(uri: str) -> Tuple[Connector, str]: + from kumo.engine.connector import connector_for_root + + parent, slash, name = uri.rstrip('/').rpartition('/') + return connector_for_root(parent or slash or '.'), name + + +def signature(snapshot: GraphSnapshot, query: str) -> Tuple[Any, pd.DataFrame]: + from kumo_core.train.formatting import prediction_output_columns + from kumo_core.types import TaskType + + from kumo.engine.pquery import PredictiveQuery + + bound = PredictiveQuery._from_snapshot(snapshot, query) + parsed = bound._parsed + entity = snapshot._native_graph.table(parsed.entity_table) + entity_type = _entity_type(str(entity.column(entity.primary_key).dtype)) + task_type = TaskType(parsed.task_type) + dtypes = { + **_DTYPES, 'TARGET_PRED': + 'boolean' if task_type == TaskType.BINARY_CLASSIFICATION else 'double' + } + inputs = _schema.Schema([ + _schema.ColSpec(entity_type, 'ENTITY'), + _schema.ColSpec('datetime', 'TIMESTAMP', required=False), + ]) + outputs = _schema.Schema([ + _schema.ColSpec(entity_type, 'ENTITY'), + *(_schema.ColSpec(dtypes.get(name) or _schema.AnyType(), name) + for name in prediction_output_columns(task_type)), + ]) + params = _schema.ParamSchema([ + _schema.ParamSpec('snapshot', 'string', None), + _schema.ParamSpec('prediction_time', 'string', None), + _schema.ParamSpec('binary_classification_threshold', 'double', None), + _schema.ParamSpec('num_classes_to_return', 'long', None), + _schema.ParamSpec('population', 'boolean', False), + ]) + example = pd.DataFrame({ + 'ENTITY': { + 'long': [0, 1], + 'double': [0.0, 1.0], + }.get(entity_type, ['entity-0', 'entity-1']) + }) + return _signature.ModelSignature(inputs, outputs, params), example + + +def _entity_type(dtype: str) -> str: + if dtype.startswith('int') or dtype == 'byte': + return 'long' + if dtype.startswith('float'): + return 'double' + return 'string' diff --git a/packages/kumo-engine/kumo/engine/integrations/_mlflow_model_code.py b/packages/kumo-engine/kumo/engine/integrations/_mlflow_model_code.py new file mode 100644 index 000000000..03cedae08 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/integrations/_mlflow_model_code.py @@ -0,0 +1,7 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +from kumo.engine.connector._extras import import_provider +from kumo.engine.integrations._mlflow_model import KumoModel + +import_provider('mlflow.models', 'mlflow').set_model(KumoModel()) diff --git a/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py b/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py new file mode 100644 index 000000000..95348b466 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py @@ -0,0 +1,230 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""Where `MlflowTracking` records, and the credentials that reach it. + +MLflow clients take a tracking and a registry URI but no credentials: the +Databricks store reads them from ``DATABRICKS_*`` variables, a profile or the +notebook (``mlflow.utils.databricks_utils.get_databricks_host_creds``), and +the SageMaker plugin signs with the process-wide tracking URI and a boto3 +session (``sagemaker_mlflow.auth_provider.AuthProvider.get_auth``). With +explicit credentials, `Target.scope` sets exactly that state around each +call, under one lock, and restores it. Ambient credentials set nothing, +except that a SageMaker target still sets the process-wide tracking URI. +""" + +import os +import re +import threading +from contextlib import ExitStack, contextmanager +from dataclasses import dataclass +from typing import Any, Callable, ContextManager, Iterator, Mapping, Optional + +from kumo.engine.connector._extras import import_provider + +SAGEMAKER_ARN = re.compile(r'arn:aws[a-z-]*:sagemaker:([^:]*):[^:]*:' + r'(mlflow-tracking-server|mlflow-app)/(.+)') +_LOCK = threading.RLock() + + +@dataclass(frozen=True) +class SageMakerMLflow: + r"""A SageMaker managed MLflow tracking server or app, for + ``MlflowTracking(target=...)``. + + Args: + tracking_server_arn: the server's (or app's) ARN. + identity: a ``kumo.engine.connector.S3Identity`` with an access key + or a role to assume; `None` or ``S3Identity()`` uses the ambient + AWS credential chain (an instance or execution role, + ``aws configure``). + """ + + tracking_server_arn: str + identity: Any = None + + def __post_init__(self) -> None: + if not SAGEMAKER_ARN.fullmatch(self.tracking_server_arn): + raise ValueError( + f'{self.tracking_server_arn!r} is not a SageMaker MLflow ' + f'tracking server or app ARN.') + + +class Target: + r"""A tracking URI, its registry URI and, for explicit credentials, the + state MLflow reads them from.""" + + def __init__(self, tracking_uri: str, registry_uri: Optional[str], + swap: Optional[Callable[[], ContextManager[Any]]] = None, + aws_session: Any = None) -> None: + self.tracking_uri = tracking_uri + self.registry_uri = registry_uri + self._swap = swap + self._aws_session = aws_session + + @property + def explicit(self) -> bool: + r"""Whether calls run with process-wide state set around them.""" + return self._swap is not None + + @contextmanager + def scope(self) -> Iterator[None]: + r"""Runs the block with this target's credentials.""" + if self._swap is None: + yield + return + with _LOCK, self._swap(): + yield + + def client(self) -> Any: + mlflow = import_provider('mlflow', 'mlflow') + return mlflow.MlflowClient(self.tracking_uri, self.registry_uri) + + def artifact_repository(self, uri: str) -> Any: + r"""The repository for artifacts at ``uri``, or `None` to go through + the client: MLflow's S3 store signs with the default AWS chain.""" + if self._aws_session is None or not uri.startswith('s3://'): + return None + s3 = import_provider('mlflow.store.artifact.s3_artifact_repo', + 'mlflow') + frozen = self._aws_session.get_credentials().get_frozen_credentials() + return s3.S3ArtifactRepository(uri, access_key_id=frozen.access_key, + secret_access_key=frozen.secret_key, + session_token=frozen.token) + + +def resolve(target: Any, tracking_uri: Optional[str]) -> Target: + mlflow = import_provider('mlflow', 'mlflow') + if target is not None and tracking_uri is not None: + raise ValueError('Pass target or tracking_uri, not both.') + if target is None: + if tracking_uri is None: + return Target(mlflow.get_tracking_uri(), mlflow.get_registry_uri()) + return Target(tracking_uri, registry_for(tracking_uri)) + if isinstance(target, SageMakerMLflow): + return _sagemaker(target) + from kumo.engine.connector.databricks import DatabricksConnector + + if isinstance(target, DatabricksConnector): + return _databricks(target) + raise TypeError(f'MlflowTracking(target=...) takes a DatabricksConnector ' + f'or a SageMakerMLflow, not {type(target).__name__}.') + + +def registry_for(tracking_uri: str) -> str: + r"""The registry a tracking URI implies: Unity Catalog on Databricks, + the tracking server itself elsewhere.""" + scheme, colon, rest = tracking_uri.partition(':') + return f'databricks-uc{colon}{rest}' if scheme == 'databricks' else ( + tracking_uri) + + +def _principal(identity: Any) -> Any: + from kumo_core.sources.principals import store + + return None if identity is None else store.get(identity.slug) + + +def _databricks(connector: Any) -> Target: + from kumoapi.data_source import DatabricksAuthenticationKind + + principal = _principal(connector.identity) + if (principal is None or principal.config.authentication + is DatabricksAuthenticationKind.AMBIENT): + return Target('databricks', 'databricks-uc') + client = import_provider('kumo_databricks.session', + 'databricks')._client_for(principal) + return Target('databricks', 'databricks-uc', + swap=lambda: _databricks_token(client)) + + +@contextmanager +def _databricks_token(client: Any) -> Iterator[None]: + r"""``DATABRICKS_HOST`` and a short-lived ``DATABRICKS_TOKEN`` minted + from the identity's own client, with every other Databricks auth + variable unset, and the process-wide tracking URI set to + ``databricks``; the client secret never enters the environment.""" + config_type = import_provider('databricks.sdk.config', + 'databricks').Config + authorization = client.config.authenticate()['Authorization'] + names = { + attribute.env + for attribute in config_type.attributes() + if (attribute.env or '').startswith('DATABRICKS_') + } + with environ({ + **{name: None + for name in names}, 'DATABRICKS_HOST': client.config.host, + 'DATABRICKS_TOKEN': authorization.partition(' ')[2] + }), global_tracking_uri('databricks'): + # Unity Catalog registration downloads the logged model through the + # process-wide tracking URI, so it must point at the workspace too. + yield + + +def _sagemaker(target: SageMakerMLflow) -> Target: + from kumoapi.data_source import S3AmbientAuthentication + + arn = target.tracking_server_arn + principal = _principal(target.identity) + if principal is None or isinstance(principal.config.authentication, + S3AmbientAuthentication): + return Target(arn, arn, swap=lambda: global_tracking_uri(arn)) + session = import_provider('kumo_s3.backend', + 's3')._session_for(principal) + return Target(arn, arn, swap=lambda: _sagemaker_session(arn, session), + aws_session=session) + + +@contextmanager +def _sagemaker_session(arn: str, session: Any) -> Iterator[None]: + r"""The plugin's signer and router read the process-wide tracking URI + and take a boto3 session from a per-thread context variable; + ``SAGEMAKER_MLFLOW_ASSUME_ROLE_ARN`` would make it assume a role.""" + plugin = import_provider('sagemaker_mlflow.session', 'sagemaker') + with ExitStack() as stack: + stack.enter_context(environ({'SAGEMAKER_MLFLOW_ASSUME_ROLE_ARN': None})) + stack.enter_context(global_tracking_uri(arn)) + stack.enter_context(plugin.use_session(session)) + yield + + +@contextmanager +def global_tracking_uri(uri: str) -> Iterator[None]: + r"""``mlflow.set_tracking_uri(uri)`` for the block. MLflow's own + ``_use_tracking_uri`` restores by unsetting ``MLFLOW_TRACKING_URI``, + so this restores the module value and the variable itself.""" + mlflow = import_provider('mlflow', 'mlflow') + utils = import_provider('mlflow.tracking._tracking_service.utils', + 'mlflow') + if mlflow.get_tracking_uri() == uri: + yield + return + saved, saved_env = utils._tracking_uri, os.environ.get( + 'MLFLOW_TRACKING_URI') + mlflow.set_tracking_uri(uri) + try: + yield + finally: + mlflow.set_tracking_uri(saved) + if saved_env is not None: + os.environ['MLFLOW_TRACKING_URI'] = saved_env + + +@contextmanager +def environ(values: Mapping[str, Optional[str]]) -> Iterator[None]: + r"""Sets (or, for `None`, unsets) environment variables for the block.""" + saved = {name: os.environ.get(name) for name in values} + try: + _apply(values) + yield + finally: + _apply(saved) + + +def _apply(values: Mapping[str, Optional[str]]) -> None: + for name, value in values.items(): + if value is None: + os.environ.pop(name, None) + else: + os.environ[name] = value diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py new file mode 100644 index 000000000..0eee9e20b --- /dev/null +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -0,0 +1,859 @@ +# 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]'``).""" + +import hashlib +import json +import logging +import os +import posixpath +import re +import tempfile +import threading +import time +import warnings +from contextlib import ExitStack, contextmanager +from pathlib import Path +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + FrozenSet, + Iterator, + List, + Mapping, + Optional, + Set, + Tuple, +) +from urllib.parse import quote + +from kumo_core.instrumentation import format_error +from kumoapi.plugin import BUNDLE_ID_KEY, ModelRegistryLocation, PublishedModel + +from kumo.engine._events import Event +from kumo.engine._mlops import full_error +from kumo.engine._tracking import DROPPED_PARAMS, RegistrationError, RunRecord +from kumo.engine._version import __version__ +from kumo.engine.connector._extras import import_provider +from kumo.engine.integrations._mlflow_target import ( + SAGEMAKER_ARN, + SageMakerMLflow, +) +from kumo.engine.integrations._mlflow_target import resolve as resolve_target + +__all__ = ['MlflowTracking', 'SageMakerMLflow', 'load_model', 'log_model'] + +if TYPE_CHECKING: + from kumo_core.embedding_models import EmbeddingModelStore + + from kumo.engine.graph import GraphSnapshot + from kumo.engine.trainer.trainer import Trainer, TrainingResult + +logger = logging.getLogger(__name__) + +_MAX_PARAM_KEY = 250 +_MAX_PARAM_VALUE = 6000 +# MLflow's portable name charset (its Windows rule, which omits ':'). +_PARAM_KEY = re.compile(r'[/\w.\- ]+') +# MLflow's defaults (7 retries, 120 s timeout) can hold a call for minutes. +_HTTP_LIMITS = { + 'MLFLOW_HTTP_REQUEST_MAX_RETRIES': '2', + 'MLFLOW_HTTP_REQUEST_TIMEOUT': '30', +} +# Model upload only, and only where the limits above apply. +_UPLOAD_HTTP = {'MLFLOW_HTTP_REQUEST_TIMEOUT': '600'} +_PATH_CONNECTORS = ('LOCAL_FILE', 'S3', 'GCS') +_SAGEMAKER_MODEL = re.compile(r'[a-zA-Z0-9](-*[a-zA-Z0-9]){0,56}') +_UC_FORBIDDEN = re.compile(r'[ /\x00-\x1f\x7f]') + + +class MlflowTracking: + r"""Records `Trainer.fit` in MLflow: one run per fit (under the active + run, if any, on the same server) with a child run per trial attempt. + Source tables are logged as run inputs: Databricks tables as Delta + tables (Unity 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. + + Args: + experiment_name: experiment for the fit run, created if missing. + `None` uses MLflow's current experiment; required with + ``target`` or ``tracking_uri``. + run_name: name of the fit run. `None` names it after its UTC start + time, as ``fit-20261027T061838Z``. + log_system_metrics: sample CPU, memory, disk, network and (with + ``nvidia-ml-py``) GPU usage on the fit run. + tags: extra tags for the fit run. + log_model: log the fitted model on the fit run (see `log_model`) + before the fit returns. ``kumo.MLOps(registry= + ModelRegistryLocation(...))`` also registers it, which implies + this. + serving_snapshot: exported snapshot URI the served model loads at + startup (see `log_model`). + target: where to record, with the credentials that reach it: a + ``DatabricksConnector`` (Databricks MLflow and Unity Catalog, + as its identity) or a `SageMakerMLflow`. `None` uses MLflow's + own tracking and registry URIs, read once here. + tracking_uri: another MLflow server, instead of ``target``. + + With explicit credentials (a Databricks service principal, an AWS + access key or role), each MLflow call sets the state MLflow reads them + from for its duration (Databricks: ``DATABRICKS_HOST`` and a short-lived + ``DATABRICKS_TOKEN``; SageMaker: the process-wide tracking URI), so do + not call MLflow or Databricks yourself on other threads during a fit. + Ambient credentials (a notebook, job or instance role) change nothing, + except that a SageMaker target still sets the process-wide tracking URI + for each call, since the plugin signs and routes with it. + + ``run_id`` is the MLflow id of the latest fit run, once it has started. + """ + + def __init__( + self, + experiment_name: Optional[str] = None, + run_name: Optional[str] = None, + log_system_metrics: bool = True, + tags: Optional[Mapping[str, str]] = None, + log_model: bool = False, + serving_snapshot: Optional[str] = None, + *, + target: Any = None, + tracking_uri: Optional[str] = None, + ) -> None: + self._mlflow = import_provider('mlflow', 'mlflow') + self.target = resolve_target(target, tracking_uri) + if experiment_name is None and (target is not None + or tracking_uri is not None): + raise ValueError('MlflowTracking(target=...) and ' + 'MlflowTracking(tracking_uri=...) need an ' + 'experiment_name.') + self._default_experiment: Optional[str] = ( + None if experiment_name is not None else + import_provider('mlflow.tracking.fluent', + 'mlflow')._get_experiment_id()) + self.experiment_name = experiment_name + self.run_name = run_name + self.log_system_metrics = log_system_metrics + self.tags = dict(tags or {}) + self.serving_snapshot = serving_snapshot + self.logs_model = log_model + self._outer: Optional[str] = None + self.run_id: Optional[str] = None + self._reset() + + def _reset(self) -> None: + self._client: Any = None + self._parent: Optional[str] = None + self._experiment_id: Optional[str] = None + self._children: Dict[Tuple[int, int], str] = {} + self._artifacts: Set[str] = set() + self._model_id: Optional[str] = None + self._http_owned: FrozenSet[str] = frozenset() + self._guard = ExitStack() + + def check_model_name(self, name: str) -> None: + r"""Raises :class:`ValueError` unless this tracker's registry + accepts ``name``.""" + check_model_name(name, self.target.registry_uri or '') + + def check(self, registry: Optional[ModelRegistryLocation] = None) -> None: + r"""Raises before any work unless the server is reachable with these + credentials and, given ``registry``, can register its model: one + experiment lookup, one registry listing and, on Unity Catalog, the + model's schema.""" + if registry is not None: + self.check_model_name(registry.model_name) + uri = self.target.tracking_uri + owned = frozenset(name for name in _HTTP_LIMITS + if name not in os.environ) + try: + with _limited_http(owned), self.target.scope(): + client = self.target.client() + if self.experiment_name is not None: + client.get_experiment_by_name(self.experiment_name) + else: + client.search_experiments(max_results=1) + if registry is not None: + client.search_registered_models(max_results=1) + check_registry_schema(registry.model_name, + self.target.registry_uri or '') + except ValueError: + raise + except Exception as error: + raise RuntimeError(f'MLflow at {uri} is not usable with these ' + f'credentials: {full_error(error)}') from error + + def prepare(self) -> None: + active = self._mlflow.active_run() + same = self._mlflow.get_tracking_uri() == self.target.tracking_uri + self._outer = (active.info.run_id + if active is not None and same else None) + + def start(self, record: RunRecord) -> None: + self._reset() + self.run_id = None + self._http_owned = frozenset(name for name in _HTTP_LIMITS + if name not in os.environ) + self._guard.enter_context(_limited_http(self._http_owned)) + try: + autologging = import_provider('mlflow.utils.autologging_utils', + 'mlflow') + # Not reentrant: leaving it clears the flag. + if not getattr(autologging, '_AUTOLOGGING_GLOBALLY_DISABLED', + False): + self._guard.enter_context(autologging.disable_autologging()) + except Exception as error: + logger.warning('Could not suspend MLflow autologging: %s', + format_error(error)) + self._client = client = self.target.client() + context = import_provider('mlflow.tracking.context.registry', + 'mlflow') + with self.target.scope(): + experiment_id = self._default_experiment + if self.experiment_name is not None: + experiment = client.get_experiment_by_name( + self.experiment_name) + experiment_id = (experiment.experiment_id if experiment else + client.create_experiment( + self.experiment_name)) + parent = {} if self._outer is None else { + 'mlflow.parentRunId': self._outer + } + run = client.create_run( + experiment_id or '0', + run_name=record.name or self.run_name + or time.strftime('fit-%Y%m%dT%H%M%SZ', time.gmtime()), + tags=context.resolve_tags({**record.tags, **parent})) + self._parent = self.run_id = run.info.run_id + self._experiment_id = run.info.experiment_id + if self.log_system_metrics and record.tags.get('kumo.kind') == 'fit': + self._guard.enter_context( + _SystemMetrics(self._log_system_metrics)) + self._log_params(self._parent, record.params) + self._log_artifacts(record.artifacts) + with self.target.scope(): + _log_inputs(self._mlflow, client, self._parent, record.inputs) + + def on_trial(self, record: RunRecord, event: Event) -> None: + key = (event.trial or 0, int(event.params['attempt'])) + if event.name == 'trial_started': + with self.target.scope(): + run = self._client.create_run( + self._experiment_id, run_name=record.name, + tags={**record.tags, 'mlflow.parentRunId': self._parent}) + self._children[key] = run.info.run_id + self._log_params(run.info.run_id, record.params) + elif event.name == 'epoch_ended': + self._log_metrics(self._children.get(key), event.metrics, + event.epoch or 0) + elif event.name == 'trial_ended': + run_id = self._children.pop(key) + self._log_metrics(run_id, record.metrics) + self._log_tags(run_id, record.tags) + succeeded = event.params.get('status') == 'succeeded' + with self.target.scope(): + self._client.set_terminated( + run_id, 'FINISHED' if succeeded else 'FAILED') + + def log_model( + self, + result: 'TrainingResult', + snapshot: 'GraphSnapshot', + registry: Optional[ModelRegistryLocation] = None, + ) -> Optional[PublishedModel]: + from kumo.engine.trainer.trainer import registry_pointers + + run_id = self._parent + if not (self.logs_model or registry) or run_id is None: + return None + upload = { + name: value + for name, value in _UPLOAD_HTTP.items() if name in self._http_owned + } + os.environ.update(upload) + try: + info = _log_model(self.target, self._client, result, snapshot, + name='model', run_id=run_id, + experiment_id=self._experiment_id, + serving_snapshot=self.serving_snapshot) + finally: + if self._parent == run_id: + os.environ.update({name: _HTTP_LIMITS[name] for name in upload}) + if self._parent != run_id: + logger.warning('The fit run %s ended before its model was ' + 'logged; the model is not registered', run_id) + return None + self._model_id = info.model_id + if registry is None: + return None + try: + return _register(self.target, self._client, info, registry, + registry_pointers(result)) + except Exception as error: + raise RegistrationError(full_error(error)) from error + + def finish(self, record: Optional[RunRecord], + error: Optional[BaseException]) -> None: + status = 'FAILED' + fit_run = (record is not None and self._parent is not None + and record.tags.get('kumo.kind') == 'fit') + try: + if record is not None and self._parent is not None: + self._log_metrics(self._parent, record.metrics, + model_id=self._model_id) + self._log_tags(self._parent, record.tags) + self._log_artifacts(record.artifacts) + if error is None: + status = 'FINISHED' + finally: + self._guard.close() + with self.target.scope(): + for run_id in self._children.values(): + _quietly(self._client.set_terminated, run_id, 'FAILED') + if self._parent is not None: + _quietly(self._client.set_terminated, self._parent, + status) + self._reset() + hint = _sagemaker_ui_hint(self.target.tracking_uri) + if fit_run and hint: + warnings.warn(hint, UserWarning, stacklevel=2) + + def _log_system_metrics(self, metrics: Mapping[str, float], + step: int) -> None: + self._log_metrics(self._parent, metrics, step) + + def _log_artifacts(self, artifacts: Mapping[str, str]) -> None: + for name, text in artifacts.items(): + if name in self._artifacts: + continue + with self.target.scope(): + repository = self.target.artifact_repository( + self._client.get_run(self._parent).info.artifact_uri + ) if self.target.explicit else None + if repository is None: + self._client.log_text(self._parent, text, name) + else: + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / name + path.write_text(text) + repository.log_artifact(str(path)) + self._artifacts.add(name) + + def _log_tags(self, run_id: str, tags: Mapping[str, str]) -> None: + entities = self._mlflow.entities + with self.target.scope(): + self._client.log_batch( + run_id, tags=[ + entities.RunTag(key, value) + for key, value in tags.items() + ]) + + def _log_params(self, run_id: str, params: Mapping[str, str]) -> None: + entities = self._mlflow.entities + kept = _mlflow_params(params) + dropped = len(params) - len(kept) + tags = [entities.RunTag(DROPPED_PARAMS, str(dropped))] if dropped else [] + with self.target.scope(): + self._client.log_batch( + run_id, params=[ + entities.Param(key, value) + for key, value in kept.items() + ], tags=tags) + + def _log_metrics(self, run_id: Optional[str], + metrics: Mapping[str, float], step: int = 0, + model_id: Optional[str] = None) -> None: + if run_id is None or not metrics: + return + entities = self._mlflow.entities + timestamp = int(time.time() * 1000) + with self.target.scope(): + self._client.log_batch( + run_id, metrics=[ + entities.Metric(metric_key(key), float(value), timestamp, + step, model_id=model_id) + for key, value in metrics.items() + ]) + + +class _SystemMetrics: + r"""Samples MLflow's CPU, disk, network and GPU monitors every + ``interval`` seconds on its own thread and hands each aggregate to + ``log(metrics, step)``, which logs it through the tracker's + credentials. (MLflow's own ``SystemMetricsMonitor`` logs with + process-wide ones.)""" + + def __init__(self, log: Callable[[Mapping[str, float], int], None], + interval: float = 10.0) -> None: + self._log = log + self._interval = interval + self._monitors = _metric_monitors() + self._stop = threading.Event() + self._thread = threading.Thread(target=self._run, daemon=True, + name='kumo-system-metrics') + + def __enter__(self) -> '_SystemMetrics': + if self._monitors: + self._thread.start() + return self + + def __exit__(self, *exc: Any) -> None: + self._stop.set() + if self._thread.is_alive(): + self._thread.join(self._interval + 5) + + def _run(self) -> None: + step = 0 + while not self._stop.wait(self._interval): + metrics: Dict[str, float] = {} + for monitor in self._monitors: + monitor.collect_metrics() + metrics.update(monitor.aggregate_metrics()) + monitor.clear_metrics() + try: + self._log({f'system/{k}': v + for k, v in metrics.items()}, step) + except Exception as error: + logger.debug('Could not log system metrics: %s', + format_error(error)) + return + step += 1 + + +def _metric_monitors() -> List[Any]: + monitors = [] + for module, name in (('cpu_monitor', 'CPUMonitor'), + ('disk_monitor', 'DiskMonitor'), + ('network_monitor', 'NetworkMonitor'), + ('gpu_monitor', 'GPUMonitor')): + try: + monitors.append( + getattr( + import_provider(f'mlflow.system_metrics.metrics.{module}', + 'mlflow'), name)()) + except Exception: + logger.debug('No %s system metrics', name, exc_info=True) + return monitors + + +def check_model_name(name: str, registry_uri: str) -> None: + r"""Raises :class:`ValueError` unless the MLflow registry at + ``registry_uri`` accepts ``name``: SageMaker managed MLflow takes letters, + digits and hyphens; Unity Catalog takes ``catalog.schema.model``, each + part at most 255 characters without spaces, ``/``, ``.`` or control + characters.""" + is_databricks_unity_catalog_uri = import_provider( + 'mlflow.utils.uri', 'mlflow').is_databricks_unity_catalog_uri + + if SAGEMAKER_ARN.fullmatch(registry_uri or ''): + if not _SAGEMAKER_MODEL.fullmatch(name): + raise ValueError( + f'SageMaker MLflow cannot register a model named {name!r}: ' + f'use letters, digits and hyphens, starting and ending with ' + f'a letter or digit, at most 57 letters and digits ' + f'(^[a-zA-Z0-9](-*[a-zA-Z0-9]){{0,56}}$), e.g. ' + f"'users-churn'.") + elif is_databricks_unity_catalog_uri(registry_uri or ''): + parts = name.split('.') + if len(parts) > 3 or not all(parts) or any( + len(part) > 255 or _UC_FORBIDDEN.search(part) + for part in parts): + raise ValueError( + f'Unity Catalog cannot register a model named {name!r}: use ' + f'catalog.schema.model (on a cluster, a shorter name takes ' + f'its current catalog and schema), each part at most 255 ' + f'characters with no space, "/" or control character, e.g. ' + f"'main.ml.users_churn'.") + + +def _sagemaker_ui_hint(tracking_uri: str) -> Optional[str]: + match = SAGEMAKER_ARN.fullmatch(tracking_uri or '') + if match is None: + return None + region, kind, name = match.groups() + command = (f'aws sagemaker create-presigned-mlflow-tracking-server-url ' + f'--tracking-server-name {name} --region {region}' + if kind == 'mlflow-tracking-server' else + f'aws sagemaker create-presigned-mlflow-app-url ' + f'--arn {tracking_uri} --region {region}') + return ("MLflow's run links do not open on SageMaker: open the MLflow UI " + f'from SageMaker Studio or at the URL `{command}` prints.') + + +def metric_key(name: str) -> str: + r"""The MLflow name of a Kumo metric: MLflow rejects ``@``, so + ``map@10`` is logged as ``map_at_10``.""" + return name.replace('@', '_at_') + + +def log_model( + trainer: 'Trainer', + *, + snapshot: 'GraphSnapshot', + name: str = 'model', + run_id: Optional[str] = None, + serving_snapshot: Optional[str] = None, + tracking: Optional[MlflowTracking] = None, +) -> Any: + r"""Log a trainer's bundle as an MLflow pyfunc model in ``run_id`` (or + the active run, else a new one) with its source tables as run inputs; + returns its `ModelInfo`. ``tracking`` says where, with what credentials + (its experiment holds a new run); `None` uses MLflow's own tracking + URI. Scores with ``predict(entities, params={'snapshot': uri})``; with + ``serving_snapshot`` (an exported snapshot URI), each serving replica + downloads that snapshot once at startup and scores against it when no + ``snapshot`` param is given.""" + from kumo.engine._tracking import fit_inputs + + mlflow = import_provider('mlflow', 'mlflow') + result = trainer.result + if result is None: + raise ValueError('log_model requires a fitted or loaded Trainer') + tracking = tracking or MlflowTracking() + target = tracking.target + client = target.client() + active = mlflow.active_run() + if run_id is None and active is not None and ( + mlflow.get_tracking_uri() == target.tracking_uri): + run_id = active.info.run_id + created = run_id is None + with target.scope(): + if created: + experiment = (client.get_experiment_by_name( + tracking.experiment_name) if tracking.experiment_name else + None) + run_id = client.create_run( + experiment.experiment_id if experiment else + tracking._default_experiment or '0').info.run_id + experiment_id = client.get_run(run_id).info.experiment_id + _log_inputs(mlflow, client, run_id, fit_inputs(result.manifest)) + info = _log_model(target, client, result, snapshot, name=name, + run_id=run_id, experiment_id=experiment_id, + serving_snapshot=serving_snapshot) + with target.scope(): + if result.metrics: + client.log_batch(run_id, metrics=[ + mlflow.entities.Metric(metric_key(key), float(value), + int(time.time() * 1000), 0, + model_id=info.model_id) + for key, value in result.metrics.items() + ]) + if created: + client.set_terminated(run_id, 'FINISHED') + return info + + +def load_model( + model_uri: str, + embedding_models: Optional['EmbeddingModelStore'] = None, + *, + tracking: Optional[MlflowTracking] = None, +) -> 'Trainer': + r"""Restore the `Trainer` of a model `log_model` logged, from the + server ``tracking`` reaches (`None`: MLflow's own tracking URI).""" + from kumo.engine._scratch import new_run_id, trainer_scratch_dir + from kumo.engine.trainer.trainer import Trainer, warn_if_other_bundle + + mlflow = import_provider('mlflow', 'mlflow') + target = (tracking or MlflowTracking()).target + destination = trainer_scratch_dir(new_run_id()) / 'mlflow-model' + destination.mkdir(parents=True) + with target.scope(): + local = Path(_download(mlflow, target, model_uri, destination)) + model = mlflow.models.Model.load(str(local)) + bundle = model.flavors['python_function']['artifacts']['bundle']['path'] + trainer = Trainer.load(local / bundle, embedding_models) + warn_if_other_bundle(model.metadata, trainer) + return trainer + + +def _download(mlflow: Any, target: Any, model_uri: str, + destination: Path) -> str: + if target.explicit and model_uri.startswith('models:/'): + client = target.client() + source = model_uri + reference = model_uri[len('models:/'):] + name, at, alias = reference.partition('@') + if at: + source = client.get_model_version_by_alias(name, alias).source + elif '/' in reference: + name, _, version = reference.partition('/') + source = client.get_model_version(name, version).source + if source.startswith('models:/') and '/' not in source[8:]: + location = client.get_logged_model(source[8:]).artifact_location + repository = target.artifact_repository(location) + if repository is not None: + return repository.download_artifacts('', str(destination)) + return mlflow.artifacts.download_artifacts( + artifact_uri=model_uri, dst_path=str(destination), + tracking_uri=target.tracking_uri, registry_uri=target.registry_uri) + + +def _log_model(target: Any, client: Any, result: 'TrainingResult', + snapshot: Optional['GraphSnapshot'], *, name: str, run_id: str, + experiment_id: Optional[str], + serving_snapshot: Optional[str]) -> Any: + r"""What ``mlflow.pyfunc.log_model`` does, through ``client``: it only + logs to the process-wide tracking URI.""" + from kumo_core.artifacts import PROVENANCE_FILE + + from kumo.engine.integrations import _mlflow_model, _mlflow_model_code + from kumo.engine.trainer._spec import read_training_spec + from kumo.engine.trainer.trainer import registry_pointers + + mlflow = import_provider('mlflow', 'mlflow') + if snapshot is None: + raise ValueError('log_model needs the snapshot the model trained on') + spec = read_training_spec(result.bundle_dir) + signature, example = _mlflow_model.signature(snapshot, + spec.predictive_query) + provenance = result.bundle_dir / PROVENANCE_FILE + with target.scope(): + model = client.create_logged_model( + experiment_id or client.get_run(run_id).info.experiment_id, + name=name, source_run_id=run_id, + params=dict(client.get_run(run_id).data.params)) + mlflow_model = mlflow.models.Model( + artifact_path=model.artifact_location, model_uuid=model.model_id, + run_id=run_id, metadata=registry_pointers(result), + model_id=model.model_id) + model_logger = logging.getLogger('mlflow.models.model') + model_logger.addFilter(_no_serving_check) + try: + with tempfile.TemporaryDirectory() as tmp: + local = str(Path(tmp) / 'model') + mlflow.pyfunc.save_model( + local, mlflow_model=mlflow_model, + python_model=_mlflow_model_code.__file__, + model_config={ + 'columns': signature.outputs.input_names(), + 'snapshot': serving_snapshot, + }, + artifacts={ + 'bundle': str(result.bundle_dir), + **({'provenance': str(provenance)} + if provenance.is_file() else {}), + }, + signature=signature, + input_example=example, + metadata=registry_pointers(result), + pip_requirements=[_requirement(result)], + ) + with target.scope(): + repository = target.artifact_repository( + model.artifact_location) + if repository is None: + client.log_model_artifacts(model.model_id, local) + else: + repository.log_artifacts(local) + logged = client.finalize_logged_model(model.model_id, + 'READY') + client.log_outputs(run_id, [ + mlflow.entities.LoggedModelOutput(model.model_id, 0) + ]) + except BaseException: + with target.scope(): + _quietly(client.finalize_logged_model, model.model_id, 'FAILED') + raise + finally: + model_logger.removeFilter(_no_serving_check) + return mlflow_model.get_model_info(logged) + + +def _requirement(result: 'TrainingResult') -> str: + from kumo_core.encoders.config import EncodersConfig + from kumo_core.io.json import load_json + from kumoapi.encoder import TextEmbedding + + extras = ['s3', 'gcs'] + path = result.artifact_paths.get('encoders') + config = (load_json(path, EncodersConfig) + if path and Path(path).is_file() else None) + if config is not None and any( + isinstance(encoder, TextEmbedding) + for table in config.encoders.values() + for encoder in table.values()): + extras.append('text') + return f"kumo-engine[{','.join(extras)}]=={__version__}" + + +def _log_inputs(mlflow: Any, client: Any, run_id: str, + inputs: Mapping[str, Mapping[str, Optional[str]]]) -> None: + try: + meta = import_provider('mlflow.data.meta_dataset', 'mlflow') + entities = mlflow.entities + datasets = [] + for name, entry in inputs.items(): + source = _dataset_source(entry) + if source is None: + continue + digest = hashlib.sha256( + json.dumps(dict(entry), sort_keys=True).encode()).hexdigest() + dataset = meta.MetaDataset(source, name=name, digest=digest[:8]) + datasets.append( + entities.DatasetInput( + dataset._to_mlflow_entity(), + tags=[entities.InputTag('mlflow.data.context', 'train')])) + if datasets: + client.log_inputs(run_id, datasets=datasets) + except Exception as error: + logger.warning('Could not log the fit\'s source tables to MLflow: %s', + format_error(error), exc_info=True) + + +def _register(target: Any, client: Any, info: Any, + location: ModelRegistryLocation, + metadata: Mapping[str, str]) -> PublishedModel: + name = location.model_name + with target.scope(): + version = _registered_version(client, name, metadata) + if version is None: + try: + client.create_registered_model(name) + except Exception as error: + if _code(error) != 'RESOURCE_ALREADY_EXISTS': + raise + version = str( + client.create_model_version(name, info.model_uri, + run_id=info.run_id, tags=metadata, + model_id=info.model_id).version) + else: + logger.info('%s version %s already holds bundle %s; reused it', + name, version, metadata.get(BUNDLE_ID_KEY)) + if location.comment: + client.update_model_version(name, version, + description=location.comment) + if location.alias: + client.set_registered_model_alias(name, location.alias, version) + return PublishedModel(name=name, version=version, + uri=f'models:/{name}/{version}', + metadata=dict(metadata)) + + +def _registered_version(client: Any, name: str, + metadata: Mapping[str, str]) -> Optional[str]: + r"""The version of ``name`` that already holds this bundle, if any.""" + bundle = metadata.get(BUNDLE_ID_KEY) + if not bundle: + return None + token = None + while True: + try: + page = client.search_model_versions(f"name = '{name}'", + page_token=token) + except Exception as error: + if _code(error) == 'RESOURCE_DOES_NOT_EXIST': + return None + raise + for version in page: + tags = version.tags or client.get_model_version( + name, version.version).tags + if (tags or {}).get(BUNDLE_ID_KEY) == bundle: + return str(version.version) + token = page.token + if not token: + return None + + +def _code(error: BaseException) -> str: + return str(getattr(error, 'error_code', '') or '') + + +def check_registry_schema(name: str, registry_uri: str) -> None: + r"""On Unity Catalog, raises :class:`ValueError` unless the catalog and + schema of a ``catalog.schema.model`` name exist and these credentials + can use them (one metadata call, no warehouse).""" + uri = import_provider('mlflow.utils.uri', 'mlflow') + parts = name.split('.') + if not uri.is_databricks_unity_catalog_uri(registry_uri) or len( + parts) != 3: + return + databricks = import_provider('mlflow.utils.databricks_utils', 'mlflow') + rest = import_provider('mlflow.utils.rest_utils', 'mlflow') + schema = f'{parts[0]}.{parts[1]}' + response = rest.http_request( + databricks.get_databricks_host_creds(registry_uri), + f'/api/2.1/unity-catalog/schemas/{quote(schema, safe="")}', 'GET', + max_retries=1, timeout=30) + if response.status_code in (403, 404): + try: + body = response.json() + except ValueError: + body = {} + raise ValueError( + f"Unity Catalog cannot register {name!r}: " + f"{body.get('error_code', response.status_code)}: " + f"{body.get('message', response.reason)}. The catalog and schema " + f"must exist, and the registering principal needs USE CATALOG, " + f"USE SCHEMA and CREATE MODEL on them.") + response.raise_for_status() + + +def _no_serving_check(record: logging.LogRecord) -> bool: + # MLflow's input-example check predicts without a snapshot; no switch. + return not record.getMessage().startswith( + 'Failed to validate serving input example') + + +def _dataset_source(entry: Mapping[str, Optional[str]]) -> Any: + connector = (entry.get('connector') or '').upper() + location = entry.get('location') + if not location or location.startswith('sql:'): + return None + if connector == 'DATABRICKS': + delta = import_provider('mlflow.data.delta_dataset_source', 'mlflow') + return delta.DeltaDatasetSource( + delta_table_name=_unquoted(location), + delta_table_version=_delta_version(entry)) + if connector in _PATH_CONNECTORS: + registry = import_provider('mlflow.data.dataset_source_registry', + 'mlflow') + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + return registry.resolve_dataset_source(location) + return None + + +def _unquoted(name: str) -> str: + parts = re.findall(r'`((?:[^`]|``)*)`', name) + return '.'.join(part.replace('``', '`') for part in parts) or name + + +def _delta_version(entry: Mapping[str, Optional[str]]) -> Optional[int]: + version = entry.get('version') or '' + head = version.split('_', 1)[0] + if entry.get('version_kind') == 'exact' and head.isdigit(): + return int(head) + return None + + +def _mlflow_params(params: Mapping[str, str]) -> Dict[str, str]: + return { + key: value + for key, value in params.items() + if len(key) <= _MAX_PARAM_KEY and _PARAM_KEY.fullmatch(key) is not None + and posixpath.normpath(key) == key and not key.startswith('/') + and len(value) <= _MAX_PARAM_VALUE + } + + +def _quietly(function: Callable[..., Any], *args: Any) -> None: + try: + function(*args) + except Exception: + logger.debug('MLflow cleanup call failed', exc_info=True) + + +@contextmanager +def _limited_http(names: FrozenSet[str]) -> Iterator[None]: + os.environ.update({name: _HTTP_LIMITS[name] for name in names}) + try: + yield + finally: + for name in names: + os.environ.pop(name, None) diff --git a/packages/kumo-engine/kumo/engine/integrations/snowflake.py b/packages/kumo-engine/kumo/engine/integrations/snowflake.py new file mode 100644 index 000000000..dfa463060 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/integrations/snowflake.py @@ -0,0 +1,281 @@ +# 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]'``).""" + +import logging +import re +import tempfile +from contextlib import ExitStack +from pathlib import Path +from typing import Any, Dict, List, Mapping, Optional, Set, Tuple + +from kumo_core.instrumentation import format_error +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.connector._extras import import_provider + +logger = logging.getLogger(__name__) + +_MAX_METRICS = 200 +_MAX_PARAMS = 1000 +# Metric name budgets: per-epoch metrics give way first, then trial summaries. +_EPOCH_METRICS = _MAX_METRICS - 60 +_SUMMARY_METRICS = _MAX_METRICS - 30 +_FINISH_PARAMS = 20 + + +class SnowflakeTracking: + r"""Records `Trainer.fit` in Snowflake Experiment Tracking as one run; + 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. + + Args: + session: Snowpark session; needs CREATE EXPERIMENT on the schema, or + MODIFY on an existing experiment. + connector: a ``SnowflakeConnector`` whose identity opens the session + instead. + experiment_name: experiment for the fit run, created if missing. + database: `None` uses the session's current database. + schema: `None` uses the session's current schema, or ``PUBLIC`` when + `database` is given or the session has none. + run_name: name of the fit run. `None` lets Snowflake pick one; a + taken name gets a suffix from the kumo run id. + tags: extra tags for the fit run. + 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. + """ + + def __init__( + self, + session: Any = None, + *, + connector: Any = None, + experiment_name: str, + database: Optional[str] = None, + schema: Optional[str] = None, + run_name: Optional[str] = None, + tags: Optional[Mapping[str, str]] = None, + native_model_lineage: bool = False, + ) -> None: + self._module = import_provider('snowflake.ml.experiment', 'snowflake') + if (session is None) == (connector is None): + raise ValueError('SnowflakeTracking takes a session or a ' + 'connector.') + self.session = session if session is not None else connector.session() + self.experiment_name = experiment_name + self.database = database + self.schema = schema + self.run_name = run_name + self.tags = dict(tags or {}) + self.native_model_lineage = native_model_lineage + self.run_id: Optional[str] = None + self._reset() + + def _reset(self) -> None: + self._tracking: Any = None + self._run: Any = None + self._run_context = ExitStack() + self._trial_tags: Dict[Tuple[int, int], Set[str]] = {} + self._metrics: Set[str] = set() + self._params = 0 + self._start_tags: Set[str] = set() + self._artifacts: Set[str] = set() + self._dropped_metrics: Set[str] = set() + self._dropped_params = 0 + + def check(self, registry: Any = None) -> None: + r"""Raises before any work unless the session can see the + experiment's schema (one metadata query, no warehouse).""" + database = self.database or self.session.get_current_database() + schema = self.schema or (None if self.database else + self.session.get_current_schema()) + if database is None: + return + name = resolve_snowflake_identifier(schema or 'PUBLIC') + try: + rows = self.session.sql( + f"SHOW SCHEMAS LIKE '{name.replace(chr(39), chr(39) * 2)}' " + f"IN DATABASE {database}").collect() + except Exception as error: + raise RuntimeError( + f'Snowflake Experiment Tracking cannot use {database}: ' + f'{type(error).__name__}: {error}') from error + if not rows: + raise ValueError( + f'Snowflake Experiment Tracking needs schema {database}.' + f'{schema or "PUBLIC"}, which does not exist or this role ' + f'cannot see.') + + def start(self, record: RunRecord) -> None: + self._reset() + self.run_id = None + cls = self._module.ExperimentTracking + tracking = getattr(cls, '_instance', None) + if tracking is None: + tracking = cls(self.session) + elif (not hasattr(tracking, '_run') or tracking._run is not None + or getattr(tracking, '_session', None) is not self.session): + logger.warning('Snowflake ExperimentTracking is in use in this ' + 'process; this fit is not tracked in Snowflake') + return + self._tracking = tracking + database = self.database or self.session.get_current_database() + schema = self.schema or (None if self.database else + self.session.get_current_schema()) + experiment = tracking.set_experiment(self.experiment_name, + database_name=database, + schema_name=schema or 'PUBLIC') + self._run = tracking.start_run( + self._new_run_name(tracking, experiment.name, + record.tags['kumo.run_id'])) + self.run_id = str(self._run.name) + if self.native_model_lineage: + self._run_context.enter_context(self._run) + self._start_tags = set(record.tags) + self._log_params({**record.params, **record.tags}) + self._log_artifacts(record.artifacts) + + def on_trial(self, record: RunRecord, event: Event) -> None: + if self._run is None: + return + prefix = f'{record.name}/' + key = (event.trial or 0, int(event.params['attempt'])) + if event.name == 'trial_started': + self._trial_tags[key] = set(record.tags) + self._log_params({prefix + k: v + for k, v in {**record.params, + **record.tags}.items()}) + elif event.name == 'epoch_ended': + self._log_metrics({prefix + k: v + for k, v in event.metrics.items()}, + _EPOCH_METRICS, step=event.epoch or 0) + elif event.name == 'trial_ended': + self._log_metrics({prefix + k: v + for k, v in record.metrics.items()}, + _SUMMARY_METRICS) + logged = self._trial_tags.pop(key, set()) + self._log_params({prefix + k: v + for k, v in record.tags.items() + if k not in logged}) + + def finish(self, record: Optional[RunRecord], + error: Optional[BaseException]) -> None: + status = 'FAILED' + try: + if self._run is not None and record is not None: + self._log_metrics(record.metrics, _MAX_METRICS) + self._log_params({k: v + for k, v in record.tags.items() + if k not in self._start_tags}, _MAX_PARAMS) + self._log_artifacts(record.artifacts) + self._log_dropped() + if error is None: + status = 'FINISHED' + finally: + try: + if self._run is not None: + self._tracking.end_run(status=status) + except Exception as error: + logger.warning('Could not end the Snowflake run %s: %s', + status, format_error(error), exc_info=True) + finally: + self._run_context.close() + self._reset() + + def _new_run_name(self, tracking: Any, experiment: Any, + run_id: str) -> Optional[str]: + if self.run_name is None: + return None + # start_run(name) resumes, or raises on, any run matching name as a + # SQL LIKE pattern and cannot create an existing run, so pick a name + # that is neither. + existing = [ + str(row['name']) + for row in tracking._sql_client.show_runs_in_experiment( + experiment_name=experiment) + ] + short = run_id.rsplit('-', 1)[-1] + for name in (self.run_name, _suffixed(self.run_name, short[:8]), + _suffixed(self.run_name, short)): + if not _taken(name, existing): + break + else: + name = None + if name != self.run_name: + logger.warning( + 'Snowflake run %s already exists in experiment %s; recording ' + 'this fit as %s', self.run_name, self.experiment_name, + name or 'a generated run name') + return name + + def _log_artifacts(self, artifacts: Mapping[str, str]) -> None: + with tempfile.TemporaryDirectory() as directory: + for name, text in artifacts.items(): + if name in self._artifacts: + continue + path = Path(directory) / name + path.write_text(text, encoding='utf-8') + self._tracking.log_artifact(str(path)) + self._artifacts.add(name) + + def _log_dropped(self) -> None: + dropped = { + DROPPED_METRICS: len(self._dropped_metrics), + DROPPED_PARAMS: self._dropped_params, + } + dropped = {name: count for name, count in dropped.items() if count} + if dropped: + self._tracking.log_params(dropped) + + def _log_params(self, params: Mapping[str, str], + limit: int = _MAX_PARAMS - _FINISH_PARAMS) -> None: + kept = {} + for name, value in params.items(): + if '$$' in name or '$$' in value or self._params >= limit: + self._dropped_params += 1 + else: + kept[name] = value + self._params += 1 + if kept: + self._tracking.log_params(kept) + + def _log_metrics(self, metrics: Mapping[str, float], limit: int, + step: int = 0) -> None: + kept = {} + for name, value in metrics.items(): + if '$$' in name: + self._dropped_metrics.add(name) + continue + if name not in self._metrics: + if len(self._metrics) >= limit: + self._dropped_metrics.add(name) + continue + self._metrics.add(name) + kept[name] = value + if kept: + self._tracking.log_metrics(kept, step=step) + + +def _suffixed(name: str, suffix: str) -> str: + if len(name) > 1 and name[0] == name[-1] == '"': + return f'{name[:-1]}_{suffix}"' + return f'{name}_{suffix}' + + +def _taken(name: str, existing: List[str]) -> bool: + resolved = resolve_snowflake_identifier(name) + regex = ''.join('.*' if c == '%' else '.' if c == '_' else re.escape(c) + for c in name) + return any( + other == resolved + or re.fullmatch(regex, other, re.IGNORECASE | re.DOTALL) is not None + for other in existing) diff --git a/packages/kumo-engine/kumo/engine/tracking.py b/packages/kumo-engine/kumo/engine/tracking.py new file mode 100644 index 000000000..af95f228f --- /dev/null +++ b/packages/kumo-engine/kumo/engine/tracking.py @@ -0,0 +1,13 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""API for writing an experiment-tracking adapter for ``kumo.MLOps(tracking=...)``.""" + +from kumo.engine._events import Event +from kumo.engine._tracking import RunRecord, Tracking + +__all__ = [ + 'Event', + 'Tracking', + 'RunRecord', +] diff --git a/packages/kumo-engine/kumo/engine/trainer/__init__.py b/packages/kumo-engine/kumo/engine/trainer/__init__.py index 0f4c1cfff..abc11c7e3 100644 --- a/packages/kumo-engine/kumo/engine/trainer/__init__.py +++ b/packages/kumo-engine/kumo/engine/trainer/__init__.py @@ -2,6 +2,7 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from kumo.engine.trainer._scoring import ScoringDefaults from kumo.engine.trainer.output_config import ( ModelOutputConfig, PredictionOutputConfig, @@ -18,6 +19,7 @@ "Trainer", "TrainingResult", "PredictionResult", + "ScoringDefaults", "ModelOutputConfig", "PredictionOutputConfig", "SnapshotOutputConfig", diff --git a/packages/kumo-engine/kumo/engine/trainer/_provenance.py b/packages/kumo-engine/kumo/engine/trainer/_provenance.py index 67bb84d56..dfd6f151e 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_provenance.py +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -12,6 +12,7 @@ FitTiming, ModelManifest, ModelParents, + TrackingRef, embedded, read_manifest, stamp, @@ -30,6 +31,7 @@ def write_provenance( warm_start: Optional[ModelManifest], started_at: datetime, fit_mode: Optional[str], + tracked_run_id: Optional[str] = None, ) -> None: from kumo_core.artifacts import PROVENANCE_FILE, add_bundle_file @@ -50,6 +52,8 @@ def write_provenance( seconds=round((finished_at - started_at).total_seconds(), 3), ), fit_mode=fit_mode, + tracking=None if tracked_run_id is None else TrackingRef( + run_id=tracked_run_id), **stamp(), )) add_bundle_file(bundle_dir, PROVENANCE_ROLE, PROVENANCE_FILE) @@ -67,8 +71,9 @@ def cut(model: Optional[ModelManifest]) -> Optional[ModelManifest]: return None if model is None else ModelManifest(id=model.id, reconstructed=True) - return replace(copy, - parents=replace(parents, warm_start=cut(parents.warm_start))) + return replace(copy, parents=replace( + parents, warm_start=cut(parents.warm_start), + derived_from=cut(parents.derived_from))) def parent_model(bundle_dir: Path, bundle_id: Optional[str] @@ -84,3 +89,26 @@ def parent_model(bundle_dir: Path, bundle_id: Optional[str] def read_provenance(bundle_dir: Path) -> Optional[ModelManifest]: path = listed_file(bundle_dir, PROVENANCE_ROLE) return None if path is None else read_manifest(path, ModelManifest) + + +def record_derivation(bundle_dir: Path, parent_dir: Path, + parent_id: Optional[str]) -> None: + r"""Writes the manifest of ``bundle_dir``, a copy of the bundle at + ``parent_dir`` with other scoring defaults.""" + from kumo_core.artifacts import ( + PROVENANCE_FILE, + add_bundle_file, + bundle_id_from_manifest, + ) + + parent = parent_model(parent_dir, parent_id) + write_manifest( + bundle_dir / PROVENANCE_FILE, + ModelManifest( + id=bundle_id_from_manifest(bundle_dir), + parents=ModelParents( + derived_from=None if parent is None else embedded_model( + parent)), + **stamp(), + )) + add_bundle_file(bundle_dir, PROVENANCE_ROLE, PROVENANCE_FILE) diff --git a/packages/kumo-engine/kumo/engine/trainer/_runtime.py b/packages/kumo-engine/kumo/engine/trainer/_runtime.py index f258d6ac1..564bdcc6d 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_runtime.py +++ b/packages/kumo-engine/kumo/engine/trainer/_runtime.py @@ -11,8 +11,10 @@ def local_training_bindings(training_input: Any, run_name: Optional[str] = None, embedding_models: Any = None, - accelerator: str = 'auto') -> Any: + accelerator: str = 'auto', + instrumentation: Any = None) -> Any: from kumo_core.training.runtime import local_training_bindings as bind return bind(training_input, run_name=run_name, - embedding_models=embedding_models, accelerator=accelerator) + embedding_models=embedding_models, accelerator=accelerator, + instrumentation=instrumentation) diff --git a/packages/kumo-engine/kumo/engine/trainer/_scoring.py b/packages/kumo-engine/kumo/engine/trainer/_scoring.py new file mode 100644 index 000000000..c6a1f0655 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/trainer/_scoring.py @@ -0,0 +1,81 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""Scoring defaults stored as ``scoring-defaults.json`` in a model bundle.""" + +import json +import warnings +from dataclasses import asdict, dataclass +from numbers import Integral, Real +from pathlib import Path +from typing import Any, Mapping, Optional + +from kumo.engine.trainer._spec import SCORING_DEFAULTS_ROLE, listed_file + +SCORING_DEFAULTS_FILE = 'scoring-defaults.json' + + +@dataclass(frozen=True) +class ScoringDefaults: + """Prediction arguments a bundle applies when ``predict`` omits them.""" + + binary_classification_threshold: Optional[float] = None + num_classes_to_return: Optional[int] = None + + +def normalized_scoring_defaults(defaults: ScoringDefaults, + task_type: Any) -> ScoringDefaults: + from kumo_core.prediction_options import ( + resolve_classification_threshold, + resolve_num_classes_to_return, + ) + + threshold = defaults.binary_classification_threshold + count = defaults.num_classes_to_return + if threshold is not None: + if isinstance(threshold, bool) or not isinstance(threshold, Real): + raise ValueError('binary_classification_threshold must be a ' + f'number, but got {threshold!r}') + threshold = float(threshold) + if resolve_classification_threshold(task_type, threshold) is None: + raise ValueError( + 'binary_classification_threshold applies only to binary and ' + f'multilabel classification, not {task_type.value}') + if isinstance(count, Integral) and not isinstance(count, bool): + count = int(count) + if count is not None: + resolve_num_classes_to_return(task_type, count) + return ScoringDefaults(threshold, count) + + +def write_scoring_defaults(bundle_dir: Path, defaults: ScoringDefaults) -> None: + from kumo_core.artifacts import add_bundle_file, remove_bundle_file + from kumo_core.io.filesystem import atomic_write_bytes + + if defaults == ScoringDefaults(): + remove_bundle_file(bundle_dir, SCORING_DEFAULTS_ROLE) + return + atomic_write_bytes( + bundle_dir / SCORING_DEFAULTS_FILE, + (json.dumps(asdict(defaults), indent=2, sort_keys=True) + + '\n').encode('utf-8')) + add_bundle_file(bundle_dir, SCORING_DEFAULTS_ROLE, SCORING_DEFAULTS_FILE) + + +def read_scoring_defaults(bundle_dir: Path, + artifact_paths: Mapping[str, str]) -> ScoringDefaults: + from kumo_core.artifacts import bundle_task_type + + path = listed_file(bundle_dir, SCORING_DEFAULTS_ROLE) + if path is None: + return ScoringDefaults() + try: + data = json.loads(path.read_text(encoding='utf-8')) + return normalized_scoring_defaults( + ScoringDefaults(data.get('binary_classification_threshold'), + data.get('num_classes_to_return')), + bundle_task_type(artifact_paths)) + except Exception as error: + warnings.warn(f'Ignoring the scoring defaults in {bundle_dir}: {error}', + UserWarning, stacklevel=3) + return ScoringDefaults() diff --git a/packages/kumo-engine/kumo/engine/trainer/_spec.py b/packages/kumo-engine/kumo/engine/trainer/_spec.py index 6c7d4a660..830937a37 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_spec.py +++ b/packages/kumo-engine/kumo/engine/trainer/_spec.py @@ -26,6 +26,7 @@ MANIFEST_FILE = 'publication-manifest.json' SPEC_ROLE = 'training_spec' PROVENANCE_ROLE = 'provenance' +SCORING_DEFAULTS_ROLE = 'scoring_defaults' HOLDOUT_ROLE = 'holdout' HOLDOUT_META_ROLE = 'holdout_meta' @@ -146,6 +147,7 @@ def read_bundle_artifact_paths( return { entry['role']: str(bundle_dir / entry['relative_path']) for entry in artifacts - if entry['role'] not in (SPEC_ROLE, PROVENANCE_ROLE, HOLDOUT_ROLE, + if entry['role'] not in (SPEC_ROLE, PROVENANCE_ROLE, + SCORING_DEFAULTS_ROLE, HOLDOUT_ROLE, HOLDOUT_META_ROLE) }, bundle_id diff --git a/packages/kumo-engine/kumo/engine/trainer/output_config.py b/packages/kumo-engine/kumo/engine/trainer/output_config.py index c828ab2a2..4aee42f1a 100644 --- a/packages/kumo-engine/kumo/engine/trainer/output_config.py +++ b/packages/kumo-engine/kumo/engine/trainer/output_config.py @@ -8,7 +8,11 @@ from kumo_core.sources.object_store_sessions import ObjectStoreSessions from kumoapi.jobs import MetadataField, PredictionArtifactType, WriteMode -from kumoapi.plugin import ModelRegistryLocation, SourceConnector +from kumoapi.plugin import ( + ModelRegistryLocation, + SourceConnector, + check_model_alias, +) from kumoapi.prediction_outputs import resolve_prediction_output_paths from kumo.engine.connector import is_file_destination, model_sink @@ -110,3 +114,17 @@ class ModelOutputConfig: def __post_init__(self) -> None: if (self.path is None) == (self.location is None): raise ValueError("Exactly one of path or location must be set.") + + +def check_model_location(config: ModelOutputConfig) -> None: + """Raises ``ValueError`` unless ``config.location`` holds an alias and + a model name its registry accepts.""" + location = config.location + if location is None: + return + if location.alias: + check_model_alias(location.alias) + check = getattr(config.connector.sink_backend, 'check_model_location', + None) + if callable(check): + check(config.connector, location) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 4e0708b08..386247183 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -8,6 +8,7 @@ import hashlib import json import logging +import os import shutil import time import warnings @@ -18,9 +19,12 @@ Any, Callable, Dict, + List, Literal, Mapping, Optional, + Sequence, + Tuple, TypeVar, Union, cast, @@ -28,6 +32,7 @@ import pandas as pd from kumo_core.embedding_models import EmbeddingModelStore +from kumo_core.instrumentation import format_error from kumo_core.model_plan import ModelPlan as CoreModelPlan from kumo_core.model_plan import inferred_or_missing_type_to_none from kumo_core.model_plan_conversion import public_model_plan_from_internal @@ -41,6 +46,8 @@ from kumo_core.types import ExplainMethod, PredictionArtifactType from kumoapi.model_plan import ModelPlan, TrainingTableGenerationPlan from kumoapi.plugin import ( + BUNDLE_ID_KEY, + ModelAliasingSink, ModelFetchingSink, ModelRegistryLocation, PublishedModel, @@ -54,14 +61,18 @@ require_embedding_models, require_trained_embedding_models, ) +from kumo.engine._events import Event, EventEmitter, query_sha256 from kumo.engine._manifest import ModelManifest +from kumo.engine._mlops import full_error, warn_failure from kumo.engine._scratch import new_run_id, trainer_scratch_dir +from kumo.engine._tracking import FitContext, TrackingSession, fit_inputs from kumo.engine.connector import model_sink from kumo.engine.connector.base import Connector from kumo.engine.connector.source_table import SourceTable from kumo.engine.graph import GraphSnapshot, report_other_graph from kumo.engine.pquery.prediction_table import PredictionTable from kumo.engine.pquery.training_table import TrainingTable +from kumo.engine.trainer._scoring import ScoringDefaults, read_scoring_defaults from kumo.engine.trainer._spec import ( HOLDOUT_META_ROLE, HOLDOUT_ROLE, @@ -74,6 +85,7 @@ from kumo.engine.trainer.output_config import ( ModelOutputConfig, PredictionOutputConfig, + check_model_location, ) logger = logging.getLogger(__name__) @@ -85,6 +97,7 @@ class TrainingResult: `published_model` records the destination's assigned name, version and URI when `fit()` exports the bundle. It is `None` without an export. + `registered_model` is the version a `kumo.MLOps` ``registry`` registered. `holdout_path` is the published model's holdout Parquet file written by `fit()`, or `None` when it has none (no holdout task, or a model refit on @@ -102,6 +115,8 @@ class TrainingResult: holdout_path: Optional[Path] = None holdout_truncated: Optional[bool] = False bundle_id: Optional[str] = None + registered_model: Optional[PublishedModel] = None + scoring_defaults: ScoringDefaults = field(default_factory=ScoringDefaults) _trainer: Optional['Trainer'] = field(default=None, repr=False, compare=False) @@ -128,6 +143,44 @@ def predict(self, *args: Any, **kwargs: Any) -> 'PredictionResult': """Score with this model, as :meth:`Trainer.predict`.""" return self._model('predict').predict(*args, **kwargs) + def export(self, config: ModelOutputConfig) -> PublishedModel: + """Publish this model's bundle, as :meth:`Trainer.export`.""" + _require_publisher(config) + return _publish(self, config) + + def with_scoring_defaults(self, + defaults: ScoringDefaults) -> 'TrainingResult': + """A copy of this model whose ``predict`` applies ``defaults`` when a + call omits them; a new bundle, whose manifest names this one in + ``parents.derived_from``.""" + from kumo_core.artifacts import bundle_task_type + + from kumo.engine.trainer._provenance import record_derivation + from kumo.engine.trainer._scoring import ( + normalized_scoring_defaults, + write_scoring_defaults, + ) + + if self._trainer is None: + raise RuntimeError('This result has no model; use Trainer.load on ' + 'its bundle_dir.') + defaults = normalized_scoring_defaults( + defaults, bundle_task_type(self.artifact_paths)) + bundle_dir = trainer_scratch_dir(new_run_id()) / 'derived-model' + shutil.copytree(self.bundle_dir, bundle_dir, + copy_function=_link_or_copy) + write_scoring_defaults(bundle_dir, defaults) + try: + record_derivation(bundle_dir, self.bundle_dir, self.bundle_id) + except Exception as error: + warnings.warn( + f"Could not record the parent model in _provenance.json: " + f"{error}", UserWarning, stacklevel=2) + result = Trainer._load_bundle(bundle_dir, + self._trainer._embedding_models).result + assert result is not None + return result + def _model(self, call: str) -> 'Trainer': trainer = self._trainer if trainer is None: @@ -329,6 +382,71 @@ def _to_core_model_plan(model_plan: ModelPlan) -> CoreModelPlan: return internal_model_plan_from_public(model_plan) +def _link_or_copy(source: str, destination: str) -> None: + try: + os.link(source, destination) + except OSError: + shutil.copy2(source, destination) + + +def _require_publisher(export_config: ModelOutputConfig) -> None: + if not callable( + getattr(model_sink(export_config.connector, export_config.path), + 'publish_model', None)): + raise NotImplementedError( + f'{type(export_config.connector).__name__} cannot publish ' + f'models') + check_model_location(export_config) + + +def _mlflow_registry(settings: Any) -> Optional[ModelRegistryLocation]: + registry = None if settings is None else settings.registry + return registry if isinstance(registry, ModelRegistryLocation) else None + + +def _connector_registry(settings: Any) -> Optional[ModelOutputConfig]: + registry = None if settings is None else settings.registry + return registry if isinstance(registry, ModelOutputConfig) else None + + +def _publish(training_result: TrainingResult, + export_config: ModelOutputConfig, + errors: Optional[List[BaseException]] = None, + warn: bool = True) -> PublishedModel: + spec = read_training_spec(training_result.bundle_dir) + metadata = { + name: str(getattr(spec, name)) + for name in ('spec_version', 'predictive_query', 'snapshot_id') + } + metadata.update(registry_pointers(training_result)) + location = export_config.location + alias = location.alias if location else None + if location and alias: + location = replace(location, alias=None) + sink = model_sink(export_config.connector, export_config.path) + published = sink.publish_model( + export_config.connector, + str(training_result.bundle_dir), + path=export_config.path, + location=location, + metadata=metadata, + metrics=training_result.metrics, + ) + if location and alias: + try: + _set_alias(export_config.connector, + replace(location, version=published.version), alias) + except Exception as error: + if warn: + warn_failure( + f"Published {published.name} version " + f"{published.version} without alias {alias}: " + f"{full_error(error)}", error, 4) + if errors is not None: + errors.append(error) + return published + + def _stage_model_bundle(artifacts: Mapping[str, str], dest: Path) -> Path: from kumo_core.artifacts import stage_model_bundle return stage_model_bundle(artifacts, dest) @@ -418,6 +536,19 @@ def result(self) -> Optional[TrainingResult]: """The last fit's result, or the loaded bundle's.""" return self._training_result + def _fitted(self, method: str) -> TrainingResult: + if self._training_result is None: + raise RuntimeError( + f"Trainer.{method}() requires a prior Trainer.fit() call on " + "this same Trainer instance.") + return self._training_result + + def export(self, config: ModelOutputConfig) -> PublishedModel: + """Publish the fitted or loaded model bundle.""" + training_result = self._fitted('export') + _require_publisher(config) + return _publish(training_result, config) + @_without_torch_scatter_hint def fit( self, @@ -427,6 +558,7 @@ def fit( export_config: Optional[ModelOutputConfig] = None, warm_start: Optional[Self] = None, mlops: Optional[Any] = None, + callbacks: Sequence[Callable[[Event], None]] = (), ) -> TrainingResult: r"""Train a model against `train_table`. @@ -456,16 +588,101 @@ def fit( architecture; they support prediction but are unreliable warm-start sources. mlops: a `kumo.MLOps`; with MLOps on (this or a `kumo.mlops` - block), the bundle is also stored under the artifact root. + block), the bundle is also stored under the artifact root and + the fit is recorded by its ``tracking``. + callbacks: called with each progress + `kumo.engine.tracking.Event`; a callback that raises is logged + and not called again. Each call returns a new `TrainingResult`; `Trainer.result` holds the latest. """ + if export_config is not None: + _require_publisher(export_config) + from kumo.engine._mlops import raise_on_error, resolve + + settings = resolve(mlops) + registry = _connector_registry(settings) + if registry is not None: + _require_publisher(registry) + run_id = f'{self._id}-fit-{new_run_id()}' + warn = settings is None or settings.on_error == 'warn' + session = TrackingSession( + None if settings is None else settings.tracking, warn=warn) + events = EventEmitter( + run_id, + [*callbacks, session.on_event] if session.enabled else callbacks) + fit_start = time.monotonic() + try: + started = { + 'snapshot_id': snapshot.id, + 'query_sha256': query_sha256(train_table.query), + } + except Exception: + started = {} + events.emit('fit_started', params=started) + try: + result = self._fit(train_table, snapshot=snapshot, + export_config=export_config, + warm_start=warm_start, run_id=run_id, + events=events, settings=settings, + tracking=session) + except BaseException as error: + events.emit('fit_ended', params={ + 'status': 'failed', + 'error': format_error(error), + }) + session.finish(None, error) + raise + events.emit('fit_ended', metrics={ + **result.metrics, 'seconds': time.monotonic() - fit_start + }, params={'status': 'succeeded'}) + registered, error = session.log_model(result, snapshot, + _mlflow_registry(settings)) + if registered is not None: + result = replace(result, registered_model=registered) + self._training_result = result + if registry is not None: + result, failure = self._register(result, registry, warn) + error = error or failure + session.finish(result, error) + raise_on_error(settings, session.error or error, + 'Trainer.fit (its model is in trainer.result)') + return result + + def _register( + self, result: TrainingResult, registry: ModelOutputConfig, warn: bool + ) -> Tuple[TrainingResult, Optional[BaseException]]: + errors: List[BaseException] = [] + try: + registered = _publish(result, registry, errors, warn) + except Exception as error: + if warn: + warn_failure( + f"Could not register the model: {full_error(error)}", + error, 3) + return result, error + result = replace(result, registered_model=registered) + self._training_result = result + return result, (errors or [None])[0] + + def _fit( + self, + train_table: TrainingTable, + *, + snapshot: GraphSnapshot, + export_config: Optional[ModelOutputConfig], + warm_start: Optional[Self], + run_id: str, + events: EventEmitter, + settings: Optional[Any], + tracking: TrackingSession, + ) -> TrainingResult: from kumo_core.automl.runtime import ( inherit_warm_start_model_plan, validate_warm_start_task_type, ) - from kumo_core.io.json import load_json + from kumo_core.io.json import load_json, to_json_dict from kumo_core.training.orchestration import ( CallableTrainingOptions, ResolvedTrainingInput, @@ -480,21 +697,10 @@ def fit( ) from kumo_core.types import TaskType - from kumo.engine.trainer._runtime import local_training_bindings - - if export_config is not None and not callable( - getattr(model_sink(export_config.connector, - export_config.path), 'publish_model', - None)): - raise NotImplementedError( - f'{type(export_config.connector).__name__} cannot publish ' - f'models') - from kumo.engine._mlops import resolve from kumo.engine.trainer._provenance import parent_model + from kumo.engine.trainer._runtime import local_training_bindings - settings = resolve(mlops) started_at = datetime.now(timezone.utc) - run_id = f'{self._id}-fit-{new_run_id()}' run_dir = self.scratch_dir / 'runs' / run_id model_plan = _to_core_model_plan(self._model_plan) @@ -529,11 +735,11 @@ def fit( search=AutoMLSearchOptions(), ) if warm_start_artifacts is not None: - source_schema = json.loads( - Path(warm_start_artifacts['prediction_schema']).read_text()) + from kumo_core.artifacts import bundle_task_type + validate_warm_start_task_type( TaskType(workflow.query.task_type), - TaskType(source_schema['task_type']), + bundle_task_type(warm_start_artifacts), ) from kumo_core.pquery import ( query_is_temporal, @@ -605,6 +811,12 @@ def fit( and not materialized.task.task_type.is_link_pred): raise ValueError('The test split is empty: shorten the query ' 'window or add data.') + tracking.start(lambda: FitContext( + run_id=run_id, + query=train_table.query, + model_plan=to_json_dict(materialized.model_plan), + inputs=fit_inputs(train_table.manifest), + )) training_input = ResolvedTrainingInput( task=materialized.task, model_plan=materialized.model_plan, @@ -617,7 +829,7 @@ def fit( bindings=local_training_bindings( training_input, run_name=run_id, embedding_models=self._embedding_models, - accelerator=self._accelerator), + accelerator=self._accelerator, instrumentation=events), options=CallableTrainingOptions(), ) @@ -669,7 +881,8 @@ def fit( warm_start=None if warm_start is None else parent_model( warm_start._training_result.bundle_dir, warm_start_bundle_id), - started_at=started_at, fit_mode=fit_mode), + started_at=started_at, fit_mode=fit_mode, + tracked_run_id=run_id if tracking.enabled else None), settings=settings, ) @@ -718,22 +931,9 @@ def _finish_fit( f"Publishing the model bundle without _provenance.json: " f"{error}", UserWarning, stacklevel=3) if export_config is not None: - spec = read_training_spec(training_result.bundle_dir) - sink = model_sink(export_config.connector, export_config.path) - published = sink.publish_model( - export_config.connector, - str(training_result.bundle_dir), - path=export_config.path, - location=export_config.location, - metadata={ - name: str(getattr(spec, name)) - for name in ('spec_version', 'predictive_query', - 'snapshot_id') - }, - metrics=training_result.metrics, - ) - training_result = replace(training_result, - published_model=published) + training_result = replace( + training_result, + published_model=_publish(training_result, export_config)) if settings is not None and settings.wants('model'): _store_model(training_result, settings) @@ -787,6 +987,8 @@ def predict( Binary classification requires an explicit threshold. Multilabel classification defaults to 0.5. Both accept thresholds from 0 to 1; thresholds change predicted labels without changing class scores. + + Omitted arguments fall back to the bundle's :attr:`scoring_defaults`. """ from kumo_core.prediction_options import ( validate_num_classes_to_return, @@ -853,11 +1055,12 @@ def _predict( num_classes_to_return: Optional[int] = None, num_workers: int = 1, binary_classification_threshold: Optional[float] = None, + extra_fields: Sequence[str] = (), ) -> PredictionResult: - if self._training_result is None: - raise RuntimeError( - "Trainer.predict() requires a prior Trainer.fit() call on " - "this same Trainer instance.") + training_result = self._fitted('predict') + if binary_classification_threshold is None: + binary_classification_threshold = ( + training_result.scoring_defaults.binary_classification_threshold) from kumo_core.pquery import normalize_graph, parse_and_validate_query from kumo_core.prediction_options import ( prediction_query_contract, @@ -891,6 +1094,8 @@ def _predict( # population and anchors still come from the immutable PredictionTable. native_graph = normalize_graph(graph_config) parsed = parse_and_validate_query(prediction_table.query, native_graph) + if num_classes_to_return is None and parsed.query.top_k is None: + num_classes_to_return = training_result.scoring_defaults.num_classes_to_return num_classes_to_return = resolve_num_classes_to_return( task_type, num_classes_to_return, @@ -910,7 +1115,7 @@ def _predict( from kumo_core.io.json import load_json # Loaded once: the embedding check and the in-process scorer share it. - encoders_path = self._training_result.artifact_paths.get('encoders') + encoders_path = training_result.artifact_paths.get('encoders') encoders = (load_json(encoders_path, EncodersConfig) if encoders_path and Path(encoders_path).is_file() else None) require_trained_embedding_models(self._embedding_models, encoders) @@ -974,12 +1179,15 @@ def _predict( model_dir = str( _stage_model_bundle( - self._training_result.artifact_paths, + training_result.artifact_paths, run_dir / 'model_bundle', )) - extra_fields = ([f.value for f in output_config.extra_fields] - if output_config is not None else []) + extra_fields = [ + *(f.value for f in (output_config.extra_fields + if output_config is not None else [])), + *extra_fields + ] source_identity = [ snapshot.id, prediction_table.query, timestamps, sorted(kind.value for kind in artifacts), num_classes_to_return, @@ -1071,10 +1279,18 @@ def load_from_connector( if not isinstance(sink, ModelFetchingSink): raise TypeError( 'model_registry.sink_backend must implement ModelFetchingSink') - sink.fetch_model(cast(SourceConnector, model_registry), str(bundle_dir), - path=address.path, location=address.location) - return cls._load_bundle(bundle_dir, embedding_models, - accelerator=accelerator) + alias = None if location is None else location.alias + if alias is not None and not isinstance(sink, ModelAliasingSink): + raise NotImplementedError( + f'{type(model_registry).__name__} cannot fetch models by ' + f'alias') + published = sink.fetch_model(cast(SourceConnector, model_registry), + str(bundle_dir), path=address.path, + location=address.location) + trainer = cls._load_bundle(bundle_dir, embedding_models, + accelerator=accelerator) + warn_if_other_bundle(getattr(published, 'metadata', None), trainer) + return trainer @classmethod def _load_bundle( @@ -1102,5 +1318,30 @@ def _load_bundle( trainer._training_result = TrainingResult( metrics, paths, bundle_dir, holdout_path=holdout_path, holdout_truncated=holdout_truncated, bundle_id=bundle_id, + scoring_defaults=read_scoring_defaults(bundle_dir, paths), _trainer=trainer) return trainer + + +def _set_alias(connector: Connector, location: ModelRegistryLocation, + alias: str) -> None: + sink = connector.sink_backend + if not isinstance(sink, ModelAliasingSink): + raise NotImplementedError( + f'{type(connector).__name__} cannot set model aliases') + sink.set_model_alias(cast(SourceConnector, connector), location, alias) + + +def warn_if_other_bundle(metadata: Optional[Mapping[str, Any]], + trainer: Trainer) -> None: + recorded = (metadata or {}).get(BUNDLE_ID_KEY) + current = None if trainer.result is None else trainer.result.id + if recorded and current and recorded != current: + warnings.warn( + f"The registry records bundle {recorded} for this version, but " + f"the fetched bundle is {current}.", UserWarning, + stacklevel=3) + + +def registry_pointers(result: TrainingResult) -> Dict[str, str]: + return {BUNDLE_ID_KEY: result.bundle_id} if result.bundle_id else {} diff --git a/packages/kumo-engine/pyproject.toml b/packages/kumo-engine/pyproject.toml index fb5741d0b..b051564ae 100644 --- a/packages/kumo-engine/pyproject.toml +++ b/packages/kumo-engine/pyproject.toml @@ -71,6 +71,17 @@ text=[ gcs=[ "kumo-gcs>=0.1.0.dev1", ] +# nvidia-ml-py adds GPU system metrics. +mlflow=[ + "mlflow-skinny>=3.0.1,<4", + "nvidia-ml-py", +] +# SageMaker managed MLflow: 0.5.0 added the per-thread boto3 session +# (sagemaker_mlflow.session.use_session) that explicit AWS identities use. +sagemaker=[ + "kumo-engine[mlflow,s3]", + "sagemaker-mlflow>=0.5.0", +] # Reports this wheel's release version as the Kumo version. The provider is a # top-level module (see `[tool.setuptools] py-modules`) because importing diff --git a/packages/kumo-engine/test/_snowflake_experiment_fake.py b/packages/kumo-engine/test/_snowflake_experiment_fake.py new file mode 100644 index 000000000..27b1cc778 --- /dev/null +++ b/packages/kumo-engine/test/_snowflake_experiment_fake.py @@ -0,0 +1,192 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""A recording stand-in for `snowflake.ml.experiment.ExperimentTracking` that +keeps its observable rules (singleton, one active run, LIKE-matched resume).""" + +import os +import re +import sys +import types +from typing import Any, Dict, List, Optional, Tuple + +from kumo_core.sources.sql import resolve_snowflake_identifier as _identifier + +STATUSES = ('FAILED', 'FINISHED') + + +def _like(pattern: str, name: str) -> bool: + regex = ''.join('.*' if c == '%' else '.' if c == '_' else re.escape(c) + for c in pattern) + return re.fullmatch(regex, name, re.IGNORECASE | re.DOTALL) is not None + + +class FakeRun: + def __init__(self, tracking: 'FakeExperimentTracking', experiment: str, + name: str) -> None: + self._tracking = tracking + self.experiment_name = experiment + self.name = name + + def __enter__(self) -> 'FakeRun': + self._tracking.calls.append(('run.__enter__', self.name)) + self._tracking.lineage_stack.append((self.experiment_name, self.name)) + return self + + def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> None: + self._tracking.calls.append(('run.__exit__', self.name)) + self._tracking.lineage_stack.pop() + if self._tracking._run is self: + self._tracking.end_run( + status='FAILED' if exc_type is not None else None) + + +class FakeExperimentTracking: + _instance: Optional['FakeExperimentTracking'] = None + + def __new__(cls, *args: Any, **kwargs: Any) -> 'FakeExperimentTracking': + if cls._instance is None: + cls._instance = super().__new__(cls) + return cls._instance + + def __init__(self, session: Any, *, database_name: Optional[str] = None, + schema_name: Optional[str] = None, + capture_source_info: bool = True, model_registry: Any = None, + warnings_as_errors: bool = False) -> None: + self.capture_source_info = capture_source_info + self.warnings_as_errors = warnings_as_errors + if hasattr(self, '_initialized'): + self.calls.append(('__init__', database_name, schema_name)) + return + self._session = session + self.database = _identifier(database_name + or session.get_current_database()) + self.schema = _identifier(schema_name or session.get_current_schema() + or 'PUBLIC') + self.calls: List[Tuple[Any, ...]] = [('__init__', database_name, + schema_name)] + self.fail_on: set = set() + self.experiment: Optional[str] = None + self._run: Optional[FakeRun] = None + self.runs: Dict[str, Dict[str, Any]] = {} + self.lineage_stack: List[Tuple[str, str]] = [] + self._sql_client = types.SimpleNamespace( + show_runs_in_experiment=self._show_runs) + self._initialized = True + + def _show_runs(self, *, experiment_name: str, + like: Optional[str] = None) -> List[Dict[str, Any]]: + self._record('show_runs', experiment_name, like) + return self._matching(experiment_name, like) + + def _matching(self, experiment: str, + like: Optional[str]) -> List[Dict[str, Any]]: + return [{'name': name} for name, run in self.runs.items() + if run['experiment'] == experiment + and (like is None or _like(like, name))] + + def _record(self, *call: Any) -> None: + self.calls.append(call) + if call[0] in self.fail_on: + raise RuntimeError(f'fake Snowflake error in {call[0]}') + + def set_experiment(self, experiment_name: str, + database_name: Optional[str] = None, + schema_name: Optional[str] = None) -> Any: + self._record('set_experiment', experiment_name, database_name, + schema_name) + if database_name is not None and schema_name is None: + schema_name = 'PUBLIC' + location = (_identifier(database_name or self.database), + _identifier(schema_name or self.schema)) + name = _identifier(experiment_name) + if (name, *location) != (self.experiment, self.database, self.schema): + self._run = None + self.experiment = name + self.database, self.schema = location + return types.SimpleNamespace(name=name) + + def start_run(self, run_name: Optional[str] = None) -> FakeRun: + self._record('start_run', run_name) + if self._run is not None: + raise RuntimeError('A run is already active. Please end the ' + 'current run before starting a new one.') + if self.experiment is None: + self.experiment = 'DEFAULT' + matches = self._matching(self.experiment, + run_name) if run_name else [] + if matches: + name = matches[0]['name'] + if self.runs[name]['status'] != 'RUNNING': + raise RuntimeError(f'Run {run_name} exists but cannot be ' + f'resumed as it is no longer running.') + else: + name = _identifier(run_name or f'RUN_{len(self.runs) + 1}') + self.runs[name] = {'experiment': self.experiment, + 'status': 'RUNNING', 'params': {}, + 'metrics': [], 'artifacts': {}} + self._run = FakeRun(self, self.experiment, name) + return self._run + + def end_run(self, run_name: Optional[str] = None, + status: Optional[str] = None) -> None: + self._record('end_run', run_name, status) + if self.experiment is None: + raise RuntimeError('No experiment set.') + if self._run is None: + raise RuntimeError('No run is active.') + if status is not None and status not in STATUSES: + raise ValueError(f"Unsupported status value: '{status}'.") + self.runs[self._run.name]['status'] = status or 'FINISHED' + self._run = None + + def _active(self) -> Dict[str, Any]: + if self._run is None: + self.start_run() + return self.runs[self._run.name] + + def log_params(self, params: Dict[str, Any]) -> None: + self._record('log_params', dict(params)) + self._active()['params'].update( + {key: str(value) for key, value in params.items()}) + + def log_metrics(self, metrics: Dict[str, float], step: int = 0) -> None: + self._record('log_metrics', dict(metrics), step) + self._active()['metrics'].extend( + (key, float(value), step) for key, value in metrics.items()) + + def log_artifact(self, local_path: str, + artifact_path: Optional[str] = None) -> None: + self._record('log_artifact', local_path, artifact_path) + with open(local_path, encoding='utf-8') as file: + self._active()['artifacts'][os.path.basename(local_path)] = ( + file.read()) + + +class FakeSession: + def __init__(self, database: Optional[str] = 'DB', + schema: Optional[str] = 'PUBLIC') -> None: + self.database, self.schema = database, schema + self.statements: List[str] = [] + + def sql(self, query: str) -> Any: + self.statements.append(query) + return types.SimpleNamespace(collect=lambda: []) + + def get_current_database(self) -> Optional[str]: + return self.database + + def get_current_schema(self) -> Optional[str]: + return self.schema + + +def install(monkeypatch: Any) -> type: + """Makes `import snowflake.ml.experiment` return a fresh fake.""" + FakeExperimentTracking._instance = None + module = types.ModuleType('snowflake.ml.experiment') + module.ExperimentTracking = FakeExperimentTracking + for name in ('snowflake', 'snowflake.ml'): + if name not in sys.modules: + monkeypatch.setitem(sys.modules, name, types.ModuleType(name)) + monkeypatch.setitem(sys.modules, 'snowflake.ml.experiment', module) + return FakeExperimentTracking diff --git a/packages/kumo-engine/test/test_fit_events.py b/packages/kumo-engine/test/test_fit_events.py new file mode 100644 index 000000000..c9efff26c --- /dev/null +++ b/packages/kumo-engine/test/test_fit_events.py @@ -0,0 +1,154 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json + +import numpy as np +import pandas as pd +import pytest +from kumoapi.model_plan import ModelPlan + +import kumo.engine as kumo +from kumo.engine._scratch import get_scratch_dir +from kumo.engine.tracking import Event + + +@pytest.fixture +def table_and_snapshot(tmp_path): + previous_scratch = get_scratch_dir() + kumo.configure(scratch_dir=tmp_path / 'scratch') + try: + rng = np.random.default_rng(0) + (tmp_path / 'data').mkdir() + pd.DataFrame({ + 'ID': np.arange(400), + 'Y': rng.random(400), + 'F1': rng.normal(size=400), + }).to_parquet(tmp_path / 'data' / 'ents.parquet') + connector = kumo.FileConnector(root_dir=str(tmp_path / 'data')) + ents = kumo.Table.from_source_table(connector['ents.parquet'], + primary_key='ID').infer_metadata() + graph = kumo.Graph(tables={'ents': ents}, edges=[]) + snapshot = graph.snapshot() + query = kumo.PredictiveQuery(graph=graph, + query='PREDICT ents.Y FOR EACH ents.ID') + yield query.generate_training_table(snapshot), snapshot + finally: + kumo.configure(scratch_dir=previous_scratch) + + +def test_fit_reports_progress_events(table_and_snapshot, monkeypatch, + caplog): + from kumo_core.automl.optuna_searcher import TPESearcher + + update = TPESearcher.update + + def fail_second_trial(self, *args, **kwargs): + if self._num_updates: + raise RuntimeError('searcher down') + self._num_updates = 1 + return update(self, *args, **kwargs) + + monkeypatch.setattr(TPESearcher, '_num_updates', 0, raising=False) + monkeypatch.setattr(TPESearcher, 'update', fail_second_trial) + table, snapshot = table_and_snapshot + plan = ModelPlan() + plan.training_job.num_experiments = 2 + plan.training_job.disable_explain = True + plan.optimization.max_epochs = 2 + plan.optimization.max_steps_per_epoch = 2 + events = [] + broken_calls = [] + + def broken(event): + broken_calls.append(event) + raise RuntimeError('callback bug') + + result = kumo.Trainer(plan).fit(table, snapshot=snapshot, + callbacks=[broken, events.append]) + + assert len(broken_calls) == 1 + assert 'callback bug' in caplog.text + assert all(isinstance(event, Event) for event in events) + names = [event.name for event in events] + assert names == [ + 'fit_started', + 'trial_started', 'epoch_ended', 'epoch_ended', 'trial_ended', + 'trial_started', 'epoch_ended', 'epoch_ended', 'trial_ended', + 'fit_ended', + ] + assert len({event.run_id for event in events}) == 1 + assert set(events[0].params) == {'snapshot_id', 'query_sha256'} + assert [(e.trial, e.epoch) for e in events if e.name == 'epoch_ended' + ] == [(1, 1), (1, 2), (2, 1), (2, 2)] + assert any(name.startswith('val_') for name in events[2].metrics) + assert events[1].params['attempt'] == 1 + assert events[4].params == {'status': 'succeeded', 'attempt': 1} + assert {'val_mae', 'seconds'} <= set(events[4].metrics) + assert events[8].params == { + 'status': 'failed', + 'attempt': 1, + 'error': 'RuntimeError: searcher down' + } + assert events[-1].metrics == {**result.metrics, + 'seconds': events[-1].metrics['seconds']} + assert events[-1].params == {'status': 'succeeded'} + assert [json.loads(json.dumps(event.to_dict()))['name'] + for event in events] == names + + +def test_failed_fit_reports_fit_ended(table_and_snapshot, monkeypatch): + import kumo.engine.trainer.trainer as trainer_module + + def fail(*args, **kwargs): + raise ValueError('stop') + + monkeypatch.setattr(trainer_module, 'build_generated_training_task_dataset', + fail) + table, snapshot = table_and_snapshot + events = [] + + with pytest.raises(ValueError, match='stop'): + kumo.Trainer(ModelPlan()).fit(table, snapshot=snapshot, + callbacks=[events.append]) + + assert [event.name for event in events] == ['fit_started', 'fit_ended'] + assert events[-1].params == { + 'status': 'failed', + 'error': 'ValueError: stop' + } + + +def test_event_metrics_keep_finite_numbers(): + from kumo.engine._events import EventEmitter + + received = [] + EventEmitter('run', [received.append]).emit( + 'epoch_ended', metrics={'val_loss': 'n/a', 'val_auc': float('nan'), + 'val_mae': 0.5}) + + assert received[0].metrics == {'val_mae': 0.5} + + +def test_epoch_event_without_recorded_metrics_does_not_raise(tmp_path): + from datetime import datetime, timezone + from types import SimpleNamespace + + from kumo_core.callbacks.train_progress_logger import ( + AutoTrainingProgressLogger, + ) + + received = [] + logger = AutoTrainingProgressLogger( + str(tmp_path), 1, on_event=lambda *a, **k: received.append(a)) + progress = SimpleNamespace(num_epochs_completed=1, metrics=[], + estimated_train_time=0.0) + logger.set_current_trainer( + SimpleNamespace(training_progress_logger=SimpleNamespace( + progress=progress)), trial=1) + logger.progress.start_time = datetime.now(timezone.utc) + + logger.on_train_epoch_end(None, None) + + assert received == [] diff --git a/packages/kumo-engine/test/test_mlflow_model.py b/packages/kumo-engine/test/test_mlflow_model.py new file mode 100644 index 000000000..8937403b3 --- /dev/null +++ b/packages/kumo-engine/test/test_mlflow_model.py @@ -0,0 +1,516 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import dataclasses +import json +import warnings +from types import SimpleNamespace + +import pandas as pd +import pytest +from kumo_core.artifacts import PROVENANCE_FILE +from kumoapi.model_plan import ModelPlan +from kumoapi.plugin import PublishedModel +from test_mlflow_tracking import ( # noqa: F401 + _runs, + table_and_snapshot, + tracking_uri, +) + +import kumo.engine as kumo +from kumo.engine._scratch import get_scratch_dir + +mlflow = pytest.importorskip('mlflow') + + +def _plan(): + plan = ModelPlan() + plan.training_job.num_experiments = 1 + plan.training_job.disable_explain = True + plan.optimization.max_epochs = 2 + plan.optimization.max_steps_per_epoch = 2 + return plan + + +def _sorted(frame): + return frame.sort_values('ENTITY').reset_index(drop=True) + + +def test_fit_logs_registers_and_serves_the_model(table_and_snapshot, + tmp_path, caplog, + monkeypatch): + from kumo.engine.integrations.mlflow import MlflowTracking, load_model + + table, snapshot = table_and_snapshot + plan = _plan() + cfg = kumo.MLOps(str(tmp_path / 'root'), + tracking=MlflowTracking(log_system_metrics=False), + registry=kumo.ModelRegistryLocation('m', + alias='champion')) + trainer = kumo.Trainer(plan) + result = trainer.fit(table, snapshot=snapshot, mlops=cfg) + + assert 'MLflow tracking failed' not in caplog.text + parent, = [ + run for run in _runs(mlflow.MlflowClient()) + if 'mlflow.parentRunId' not in run.data.tags + ] + model, = mlflow.search_logged_models(output_format='list') + assert model.source_run_id == parent.info.run_id + metrics = {metric.key: metric.value for metric in model.metrics} + assert metrics == {**result.metrics, 'seconds': metrics['seconds']} + version = mlflow.MlflowClient().get_model_version_by_alias('m', 'champion') + assert version.source == model.model_uri + assert {k: v for k, v in parent.data.tags.items() + if k.startswith('kumo.')} == { + 'kumo.kind': 'fit', + 'kumo.run_id': parent.data.tags['kumo.run_id'], + 'kumo.bundle_id': result.id, + } + assert result.registered_model == PublishedModel( + name='m', version=str(version.version), + uri=f'models:/m/{version.version}', + metadata={'kumo_bundle_id': result.id}) + assert version.tags == {'kumo_bundle_id': result.id} + assert trainer.result is result + + info = mlflow.models.get_model_info('models:/m@champion') + assert info.metadata == {'kumo_bundle_id': result.id} + provenance = json.loads(mlflow.artifacts.load_text( + f'{model.artifact_location}/artifacts/{PROVENANCE_FILE}')) + assert provenance == result.manifest.to_dict() + assert [(c['name'], c['type']) for c in info.signature.inputs.to_dict() + ] == [('ENTITY', 'long'), ('TIMESTAMP', 'datetime')] + assert [(c['name'], c['type']) for c in info.signature.outputs.to_dict() + ] == [('ENTITY', 'long'), ('TARGET_PRED', 'double')] + + location = snapshot.export( + kumo.SnapshotOutputConfig( + kumo.FileConnector(root_dir=str(tmp_path / 'exports')), + 'snapshots')) + served = mlflow.pyfunc.load_model('models:/m@champion').predict( + pd.DataFrame({'ENTITY': range(0, 400, 7)}), + params={'snapshot': str(tmp_path / 'exports' / location)}) + assert list(served) == ['ENTITY', 'TARGET_PRED'] + assert len(served) == len(range(0, 400, 7)) + assert load_model('models:/m@champion').result.id == result.id + + original = mlflow.models.Model.load + + def other_bundle(path): + model = original(path) + model.metadata['kumo_bundle_id'] = 'other' + return model + + with pytest.MonkeyPatch.context() as patch, pytest.warns( + UserWarning, match='records bundle other'): + patch.setattr(mlflow.models.Model, 'load', other_bundle) + load_model('models:/m@champion') + + +def test_logged_model_scores_the_query_population_and_timed_rows( + fitted, tracking_uri, tmp_path, monkeypatch): + from kumo.engine.integrations.mlflow import log_model + + trainer, table, snapshot = fitted + location = snapshot.export( + kumo.SnapshotOutputConfig( + kumo.FileConnector(root_dir=str(tmp_path / 'exports')), + 'snapshots')) + + info = log_model(trainer, snapshot=snapshot, + serving_snapshot=str(tmp_path / 'exports' / location)) + + assert mlflow.get_run(info.run_id).info.status == 'FINISHED' + requirements = mlflow.pyfunc.get_model_dependencies(info.model_uri) + assert f'kumo-engine[gcs,s3]=={kumo.__version__}' in open( + requirements).read().splitlines() + assert {metric.key + for metric in mlflow.get_logged_model(info.model_id).metrics + } == set(trainer.result.metrics) + assert [c['name'] for c in info.signature.outputs.to_dict() + ] == ['ENTITY', 'TARGET_PRED', 'False_PROB', 'True_PROB'] + pyfunc = mlflow.pyfunc.load_model(info.model_uri) + params = { + 'snapshot': str(tmp_path / 'exports' / location), + 'prediction_time': str(pd.Timestamp(table.timestamps[0])), + 'binary_classification_threshold': 0.5, + } + expected = _sorted( + trainer.predict(table, snapshot=snapshot, + binary_classification_threshold=0.5).frame) + empty = pd.DataFrame({'ENTITY': pd.Series([], dtype='int64')}) + tables = get_scratch_dir() / 'prediction_tables' + before = set(tables.glob('*')) + predict = kumo.Trainer.predict + + def predict_beside_another_call(self, *args, **kwargs): + (tables / 'other-call').mkdir(parents=True, exist_ok=True) + return predict(self, *args, **kwargs) + + monkeypatch.setattr(kumo.Trainer, 'predict', predict_beside_another_call) + + assert pyfunc.predict(empty, params=params).empty + served = pyfunc.unwrap_python_model()._snapshots + assert list(served) == [params['snapshot']] + pd.testing.assert_frame_equal( + _sorted( + pyfunc.predict(empty, params={ + **params, 'snapshot': None, 'population': True + })), expected) + for rows, population in ((empty, True), + (pd.DataFrame({'ENTITY': range(50)}), False)): + pd.testing.assert_frame_equal( + _sorted( + pyfunc.predict(rows, + params={**params, 'population': population})), + expected) + scratch = pyfunc.unwrap_python_model()._trainer.scratch_dir + assert list(scratch.iterdir()) == [] + assert set(tables.glob('*')) == before | {tables / 'other-call'} + + +def test_pyfunc_names_missing_temporal_and_threshold_inputs( + fitted, monkeypatch): + from dataclasses import replace + + from kumo.engine.integrations._mlflow_model import predict + + trainer, table, snapshot = fitted + rows = pd.DataFrame({'ENTITY': range(5)}) + time = str(pd.Timestamp(table.timestamps[0])) + with pytest.raises(ValueError) as error: + predict(trainer, snapshot, rows, + {'binary_classification_threshold': .5}) + assert str(error.value) == ("temporal query: pass params={" + "'prediction_time': ...} or a TIMESTAMP " + "column") + with pytest.raises(ValueError, match=r"^binary classification: pass " + r"params=\{'binary_classification_threshold': " + r"\.\.\.\}$"): + predict(trainer, snapshot, rows, {'prediction_time': time}) + + monkeypatch.setattr( + trainer, '_training_result', + replace(trainer.result, scoring_defaults=kumo.ScoringDefaults( + binary_classification_threshold=.5))) + assert len(predict(trainer, snapshot, rows, {'prediction_time': time + })) == 5 + + +def test_pyfunc_returns_rows_in_input_order(fitted): + from kumo.engine.integrations._mlflow_model import predict + + trainer, table, snapshot = fitted + time = pd.Timestamp(table.timestamps[0]) + entities = [7, 3, 9, 3, 999, 1, 40, 2, 33, 11, 3, 25] * 10 + rows = pd.DataFrame({ + 'ENTITY': entities, + 'TIMESTAMP': [time - pd.Timedelta(days=i % 3) + for i in range(len(entities))], + }) + scored = predict(trainer, snapshot, rows, + {'binary_classification_threshold': .5}) + assert scored['ENTITY'].tolist() == [e for e in entities if e != 999] + assert 'ANCHOR_TIMESTAMP' not in scored + + +def test_scores_pair_with_their_rows_when_scoring_reorders(): + from kumo.engine.integrations._mlflow_model import _in_input_order + + day = pd.Timestamp('2024-05-01 12:00:00.250') + rows = pd.DataFrame({ + 'ENTITY1': ['b', 'a', 'b', 'c', 'a', 'gone'], + 'TIMESTAMP': [day, day, day - pd.Timedelta(days=1), day, day, day], + }) + seconds = day.floor('s') + scored = pd.DataFrame({ + 'ENTITY': ['a', 'c', 'b', 'a', 'b'], + 'ANCHOR_TIMESTAMP': [seconds, seconds, seconds - pd.Timedelta(days=1), + seconds, seconds], + 'SCORE': [1.0, 3.0, 2.5, 1.0, 2.0], + }) + + ordered = _in_input_order(rows, scored) + + assert ordered.to_dict('list') == { + 'ENTITY': ['b', 'a', 'b', 'c', 'a'], + 'SCORE': [2.0, 1.0, 2.5, 3.0, 1.0], + } + + +@pytest.mark.parametrize('tracking,registry', [ + (None, kumo.ModelRegistryLocation('m')), + (SimpleNamespace(start=None), kumo.ModelRegistryLocation('m')), + ('mlflow', kumo.ModelRegistryLocation('m', alias='SageMakerLifecycle-p')), + ('mlflow', kumo.ModelRegistryLocation('m', version='3')), +]) +def test_invalid_mlflow_registry_is_refused(tracking_uri, tmp_path, tracking, + registry): + from kumo.engine.integrations.mlflow import MlflowTracking + + if tracking == 'mlflow': + tracking = MlflowTracking() + with pytest.raises(ValueError): + kumo.MLOps(str(tmp_path), tracking=tracking, registry=registry) + + +def test_fit_waits_for_model_logging_and_fails_the_run_if_it_fails( + monkeypatch): + import time + + from kumo.engine import _tracking + from kumo.engine._tracking import FitContext, TrackingSession + + monkeypatch.setattr(_tracking, 'FINISH_SECONDS', 0.01) + + class Tracker: + def __init__(self, fail): + self.fail, self.calls = fail, [] + + def start(self, record): + self.calls.append('start') + + def on_trial(self, record, event): + pass + + def log_model(self, result, snapshot, registry): + time.sleep(0.2) + self.calls.append('log_model') + if self.fail: + raise RuntimeError('upload failed') + + def finish(self, record, error): + self.calls.append(('finish', error)) + + result = SimpleNamespace(metrics={}, bundle_dir='missing', bundle_id='b') + for fail in (False, True): + tracker = Tracker(fail) + session = TrackingSession(tracker) + session.start(lambda: FitContext('r', 'q', {})) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter('always') + _, error = session.log_model(result, None) + session.finish(result, error) + session._worker.join(5) + assert tracker.calls[:2] == ['start', 'log_model'] + assert tracker.calls[2] == ('finish', error) + assert (error is not None) is fail + assert any('upload failed' in str(w.message) for w in caught) is fail + + +@pytest.mark.parametrize('stall', ['slow upload', 'dead thread']) +def test_model_logging_wait_is_bounded(monkeypatch, stall): + import threading + + from kumo.engine import _tracking + from kumo.engine._tracking import FitContext, TrackingSession + + monkeypatch.setattr(_tracking, 'MODEL_SECONDS', 0.5) + release = threading.Event() + + class Tracker: + calls = [] + + def start(self, record): + pass + + def on_trial(self, record, event): + pass + + def log_model(self, result, snapshot, registry): + release.wait(30) + + def finish(self, record, error): + self.calls.append((record, error)) + + tracker = Tracker() + session = TrackingSession(tracker) + session.start(lambda: FitContext('r', 'q', {})) + if stall == 'dead thread': + session._submit(lambda: (_ for _ in ()).throw(SystemExit())) + session._worker.join(5) + with pytest.warns(UserWarning, match='Could not log the model'): + _, error = session.log_model(SimpleNamespace(), None) + release.set() + session.finish(None, error) + + assert type(error) is (TimeoutError + if stall == 'slow upload' else RuntimeError) + assert tracker.calls == [(None, None)] + + +def test_model_upload_gets_a_longer_http_timeout_and_a_late_one_is_dropped( + tracking_uri, monkeypatch, caplog): + import os + + from test_tracking_contract import CONTEXT + + from kumo.engine._tracking import TrackingSession + from kumo.engine.integrations import mlflow as integration + + for name in integration._HTTP_LIMITS: + monkeypatch.delenv(name, raising=False) + tracker = integration.MlflowTracking(log_system_metrics=False) + location = kumo.ModelRegistryLocation('m') + seen = [] + + def log_model(target, client, result, snapshot, **kwargs): + seen.append(os.environ.get('MLFLOW_HTTP_REQUEST_TIMEOUT')) + if len(seen) == 2: + tracker.finish(None, None) + return SimpleNamespace(model_id=None, model_uri='runs:/x/model') + + monkeypatch.setattr(integration, '_log_model', log_model) + monkeypatch.setattr(integration, '_register', lambda *args: 'registered') + for _ in range(2): + session = TrackingSession(tracker) + session.start(lambda: CONTEXT) + registered, error = session.log_model( + SimpleNamespace(bundle_id=None), None, location) + assert error is None + during = os.environ.get('MLFLOW_HTTP_REQUEST_TIMEOUT') + session.finish(None, None) + session._worker.join(5) + assert registered == ('registered' if len(seen) == 1 else None) + assert during == ('30' if len(seen) == 1 else None) + + assert seen == ['600', '600'] + assert 'MLFLOW_HTTP_REQUEST_TIMEOUT' not in os.environ + assert 'the model is not registered' in caplog.text + + +def test_a_model_with_text_columns_requires_the_text_extra(tmp_path): + from kumo_core.encoders.config import EncodersConfig + from kumo_core.io.json import save_json + from kumoapi.encoder import Numerical, TextEmbedding + + from kumo.engine.integrations.mlflow import _requirement + + def requirement(encoders): + path = tmp_path / 'encoders.json' + save_json(EncodersConfig(encoders={'users': encoders}), str(path)) + return _requirement( + SimpleNamespace(artifact_paths={'encoders': str(path)})) + + assert requirement({'age': Numerical()}) == ( + f'kumo-engine[s3,gcs]=={kumo.__version__}') + assert requirement({'age': Numerical(), 'bio': TextEmbedding('words') + }) == f'kumo-engine[s3,gcs,text]=={kumo.__version__}' + + +def test_a_tracker_dropped_earlier_says_the_model_was_not_logged(): + import time + + from kumo.engine._tracking import FitContext, TrackingSession + + class Tracker: + def start(self, record): + raise RuntimeError('server down') + + def log_model(self, result, snapshot, registry): + raise AssertionError('not called') + + def finish(self, record, error): + pass + + session = TrackingSession(Tracker()) + session.start(lambda: FitContext('r', 'q', {})) + deadline = time.monotonic() + 10 + while session._tracker is not None and time.monotonic() < deadline: + time.sleep(0.01) + with pytest.warns(UserWarning, match=r'Could not log the model: ' + r'RuntimeError: tracking stopped after an earlier ' + r'failure \(RuntimeError: server down\)'): + registered, error = session.log_model(SimpleNamespace(), None) + session.finish(None, error) + assert registered is None and isinstance(error, RuntimeError) + assert str(session.error) == 'server down' + + +@pytest.mark.parametrize('on_error', ['warn', 'raise']) +def test_a_registration_failure_is_tagged_on_the_fit_run_and_can_raise( + table_and_snapshot, tmp_path, monkeypatch, on_error): + from mlflow.exceptions import MlflowException + from mlflow.tracking import MlflowClient + + from kumo.engine.integrations.mlflow import MlflowTracking + + detail = ("CATALOG_DOES_NOT_EXIST: Catalog 'nope' does not exist. " + + 'x' * 300 + ' the actionable end') + + def missing_catalog(*args, **kwargs): + raise MlflowException(detail) + + monkeypatch.setattr(MlflowClient, 'create_model_version', + missing_catalog) + table, snapshot = table_and_snapshot + cfg = kumo.MLOps(str(tmp_path / 'root'), + tracking=MlflowTracking(log_system_metrics=False), + registry=kumo.ModelRegistryLocation('nope.s.m'), + on_error=on_error) + trainer = kumo.Trainer(_plan()) + cause = f'MlflowException: {detail}' + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter('always') + if on_error == 'warn': + result = trainer.fit(table, snapshot=snapshot, mlops=cfg) + else: + with pytest.raises(RuntimeError) as raised: + trainer.fit(table, snapshot=snapshot, mlops=cfg) + assert str(raised.value).endswith(cause) + assert isinstance(raised.value.__cause__, MlflowException) + result = trainer.result + reported = [str(w.message) for w in caught if 'nope' in str(w.message)] + assert reported == ([f'Could not register the model: {cause}'] + if on_error == 'warn' else []) + assert result.id and result.registered_model is None + + run, = [run for run in _runs(mlflow.MlflowClient()) + if 'mlflow.parentRunId' not in run.data.tags] + assert run.info.status == 'FAILED' + assert run.data.tags['kumo.error'].startswith('MlflowException: ' + 'CATALOG_DOES_NOT_EXIST') + assert run.data.tags['kumo.bundle_id'] == result.id + + +def test_a_tracking_start_failure_names_its_cause_and_can_raise( + table_and_snapshot, tmp_path, monkeypatch, caplog): + from mlflow.exceptions import MlflowException + from mlflow.tracking import MlflowClient + + from kumo.engine.integrations.mlflow import MlflowTracking + + def no_folder(self, name, *args, **kwargs): + raise MlflowException(f'RESOURCE_DOES_NOT_EXIST: Parent of {name} ' + f'does not exist.') + + monkeypatch.setattr(MlflowClient, 'create_experiment', no_folder) + table, snapshot = table_and_snapshot + cfg = kumo.MLOps(str(tmp_path / 'root'), + tracking=MlflowTracking('/NoSuchFolder/churn', + log_system_metrics=False), + on_error='raise') + trainer = kumo.Trainer(_plan()) + with pytest.raises(RuntimeError, match='RESOURCE_DOES_NOT_EXIST') as raised: + trainer.fit(table, snapshot=snapshot, mlops=cfg) + + assert isinstance(raised.value.__cause__, MlflowException) + assert trainer.result.id + off, = [record for record in caplog.records + if 'off for the rest' in record.getMessage()] + assert 'MlflowException: RESOURCE_DOES_NOT_EXIST' in off.getMessage() + assert not off.exc_info + + cfg = dataclasses.replace(cfg, on_error='warn') + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter('always') + trainer.fit(table, snapshot=snapshot, mlops=cfg) + reported = [str(w.message) for w in caught + if 'RESOURCE_DOES_NOT_EXIST' in str(w.message)] + assert reported == [ + 'Could not log the model: RuntimeError: tracking stopped after an ' + 'earlier failure (MlflowException: RESOURCE_DOES_NOT_EXIST: Parent ' + 'of /NoSuchFolder/churn does not exist.)' + ] diff --git a/packages/kumo-engine/test/test_mlflow_target.py b/packages/kumo-engine/test/test_mlflow_target.py new file mode 100644 index 000000000..c18a2df68 --- /dev/null +++ b/packages/kumo-engine/test/test_mlflow_target.py @@ -0,0 +1,290 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""How `MlflowTracking` reaches its server: explicit credentials set exactly +the state MLflow reads, only around each call; ambient ones set nothing. The +last tests pin the MLflow and plugin internals this relies on.""" + +import os +import threading +from types import SimpleNamespace + +import pytest + +import kumo.engine as kumo +from kumo.engine.connector import ( + DatabricksConnector, + DatabricksIdentity, + S3Identity, +) +from kumo.engine.integrations import _mlflow_target as target_module +from kumo.engine.integrations._mlflow_target import SageMakerMLflow, resolve + +mlflow = pytest.importorskip('mlflow') + +ARN = 'arn:aws:sagemaker:us-west-2:account:mlflow-tracking-server/demo' + + +@pytest.fixture +def tracking_uri(tmp_path): + previous = mlflow.get_tracking_uri() + uri = f'sqlite:///{tmp_path / "mlflow.db"}' + mlflow.set_tracking_uri(uri) + yield uri + mlflow.set_tracking_uri(previous) + + +def _connector(identity): + return DatabricksConnector(catalog='main', schema_name='ml', + identity=identity, warehouse_id='w') + + +def test_no_target_reads_mlflows_own_uris_once(tracking_uri): + target = resolve(None, None) + assert target.tracking_uri == tracking_uri + assert not target.explicit + assert resolve(None, 'databricks').registry_uri == 'databricks-uc' + assert resolve(None, 'databricks://prod').registry_uri == ( + 'databricks-uc://prod') + assert resolve(None, 'http://mlflow:5000').registry_uri == ( + 'http://mlflow:5000') + with pytest.raises(ValueError, match='not both'): + resolve(SageMakerMLflow(ARN), 'http://x') + + +def test_ambient_identities_set_no_credentials(monkeypatch, tracking_uri): + monkeypatch.setenv('DATABRICKS_TOKEN', 'mine') + monkeypatch.setenv('MLFLOW_TRACKING_URI', tracking_uri) + target = resolve(_connector(DatabricksIdentity()), None) + before = dict(os.environ) + with target.scope(): + assert dict(os.environ) == before + assert mlflow.get_tracking_uri() == tracking_uri + assert not target.explicit + + for target in (resolve(SageMakerMLflow(ARN, S3Identity()), None), + resolve(SageMakerMLflow(ARN), None)): + with target.scope(): + assert mlflow.get_tracking_uri() == ARN + assert target.artifact_repository('s3://b/run') is None + assert mlflow.get_tracking_uri() == tracking_uri + assert dict(os.environ) == before + + +def test_a_databricks_service_principal_lends_a_token_for_each_call( + monkeypatch): + client = SimpleNamespace(config=SimpleNamespace( + host='https://dbc.example', authenticate=lambda: { + 'Authorization': 'Bearer short-lived' + })) + monkeypatch.setattr('kumo_databricks.session._client_for', + lambda principal: client) + monkeypatch.setenv('DATABRICKS_CLIENT_SECRET', 'other') + monkeypatch.setenv('DATABRICKS_CONFIG_PROFILE', 'laptop') + monkeypatch.delenv('DATABRICKS_TOKEN', raising=False) + target = resolve( + _connector( + DatabricksIdentity(host='https://dbc.example', client_id='sp', + client_secret='secret')), None) + + previous = mlflow.get_tracking_uri() + with target.scope(): + inside = {k: v for k, v in os.environ.items() + if k.startswith('DATABRICKS_')} + # Unity Catalog registration downloads through the global URI. + assert mlflow.get_tracking_uri() == 'databricks' + assert mlflow.get_tracking_uri() == previous + assert inside == {'DATABRICKS_HOST': 'https://dbc.example', + 'DATABRICKS_TOKEN': 'short-lived'} + assert 'secret' not in inside.values() + assert os.environ['DATABRICKS_CLIENT_SECRET'] == 'other' + assert os.environ['DATABRICKS_CONFIG_PROFILE'] == 'laptop' + assert 'DATABRICKS_TOKEN' not in os.environ + assert (target.tracking_uri, target.registry_uri) == ('databricks', + 'databricks-uc') + + +def test_an_aws_access_key_signs_sagemaker_calls_and_uploads( + monkeypatch, tracking_uri): + boto3 = pytest.importorskip('boto3') + plugin = pytest.importorskip('sagemaker_mlflow.session') + monkeypatch.setenv('MLFLOW_TRACKING_URI', tracking_uri) + monkeypatch.setenv('SAGEMAKER_MLFLOW_ASSUME_ROLE_ARN', 'arn:aws:iam::x') + target = resolve( + SageMakerMLflow(ARN, S3Identity(access_key_id='AKIA', + secret_access_key='s3cret')), None) + + with target.scope(): + assert mlflow.get_tracking_uri() == ARN + assert 'SAGEMAKER_MLFLOW_ASSUME_ROLE_ARN' not in os.environ + session = plugin._get_current_session() + assert session.get_credentials().access_key == 'AKIA' + repository = target.artifact_repository('s3://bucket/run/artifacts') + assert mlflow.get_tracking_uri() == tracking_uri + assert os.environ['MLFLOW_TRACKING_URI'] == tracking_uri + assert os.environ['SAGEMAKER_MLFLOW_ASSUME_ROLE_ARN'] == 'arn:aws:iam::x' + assert plugin._get_current_session() is None + assert repository._access_key_id == 'AKIA' + assert target.artifact_repository('dbfs:/x') is None + assert isinstance(session, boto3.Session) + + +def test_explicit_calls_take_turns(monkeypatch): + target = target_module.Target('databricks', 'databricks-uc', + swap=lambda: target_module.environ( + {'KUMO_PROBE': 'x'})) + seen = [] + + def call(): + with target.scope(): + seen.append(os.environ.get('KUMO_PROBE')) + + threads = [threading.Thread(target=call) for _ in range(8)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + assert seen == ['x'] * 8 and 'KUMO_PROBE' not in os.environ + + +def test_ambient_sagemaker_calls_see_its_uri_and_restore_the_global_one( + tracking_uri, monkeypatch): + pytest.importorskip('sagemaker_mlflow') # registers the arn store + from kumo.engine.integrations.mlflow import MlflowTracking + + seen = [] + + def lookup(self, name): + seen.append((self.tracking_uri, mlflow.get_tracking_uri())) + + monkeypatch.setattr(mlflow.MlflowClient, 'get_experiment_by_name', + lookup) + tracking = MlflowTracking(target=SageMakerMLflow(ARN), + experiment_name='e') + kumo.MLOps('/tmp/x', tracking=tracking) + + assert seen == [(ARN, ARN)] + assert mlflow.get_tracking_uri() == tracking_uri + + +def test_an_unreachable_server_fails_when_mlops_is_built( + tmp_path, monkeypatch): + from kumo.engine.integrations.mlflow import MlflowTracking + + monkeypatch.setenv('MLFLOW_HTTP_REQUEST_MAX_RETRIES', '0') + tracking = MlflowTracking(tracking_uri='http://127.0.0.1:9', + experiment_name='e') + with pytest.raises(RuntimeError, match='MLflow at http://127.0.0.1:9 ' + 'is not usable with these credentials'): + kumo.MLOps(str(tmp_path), tracking=tracking) + + +def test_a_missing_unity_catalog_schema_fails_when_mlops_is_built( + monkeypatch): + from kumo.engine.integrations import mlflow as integration + + calls = [] + + def http_request(creds, endpoint, method, **kwargs): + calls.append(endpoint) + return SimpleNamespace( + status_code=404, reason='Not Found', json=lambda: { + 'error_code': 'CATALOG_DOES_NOT_EXIST', + 'message': "Catalog 'nope' does not exist." + }) + + monkeypatch.setattr('mlflow.utils.rest_utils.http_request', http_request) + monkeypatch.setattr('mlflow.utils.databricks_utils.' + 'get_databricks_host_creds', lambda uri: None) + with pytest.raises(ValueError, match="cannot register 'nope.ml.churn': " + 'CATALOG_DOES_NOT_EXIST'): + integration.check_registry_schema('nope.ml.churn', 'databricks-uc') + assert calls == ['/api/2.1/unity-catalog/schemas/nope.ml'] + integration.check_registry_schema('churn', 'databricks-uc') + integration.check_registry_schema('a.b.c', 'sqlite:///x.db') + assert len(calls) == 1 + + +# Pins: MLflow and plugin internals the code above relies on. + + +def test_pin_set_tracking_uri_keeps_the_module_value(tracking_uri): + utils = pytest.importorskip('mlflow.tracking._tracking_service.utils') + assert utils._tracking_uri == tracking_uri + + +def test_pin_current_experiment_lookup(): + from mlflow.tracking.fluent import _get_experiment_id + + assert callable(_get_experiment_id) + + +def test_pin_sagemaker_signer_takes_the_thread_session(monkeypatch): + boto3 = pytest.importorskip('boto3') + pytest.importorskip('sagemaker_mlflow') + from sagemaker_mlflow.auth import AuthBoto + from sagemaker_mlflow.session import use_session + + session = boto3.Session(aws_access_key_id='AKIA', + aws_secret_access_key='s3cret') + with use_session(session): + auth = AuthBoto('us-west-2', 'sagemaker-mlflow') + assert auth.creds.access_key == 'AKIA' + + +def test_pin_system_metric_monitors(): + from kumo.engine.integrations.mlflow import _metric_monitors + + monitors = _metric_monitors() + assert len(monitors) >= 3 + for monitor in monitors: + monitor.collect_metrics() + assert isinstance(monitor.aggregate_metrics(), dict) + monitor.clear_metrics() + + +def test_pin_databricks_auth_variables(): + from databricks.sdk.config import Config + + names = {attribute.env for attribute in Config.attributes()} + assert {'DATABRICKS_HOST', 'DATABRICKS_TOKEN', 'DATABRICKS_CLIENT_ID', + 'DATABRICKS_CLIENT_SECRET', 'DATABRICKS_CONFIG_PROFILE'} <= names + + +def test_registering_a_registered_bundle_reuses_its_version( + tracking_uri, caplog): + from kumo.engine.integrations.mlflow import _register + + target = resolve(None, None) + client = target.client() + run = client.create_run('0') + info = SimpleNamespace(model_uri=f'runs:/{run.info.run_id}/model', + run_id=run.info.run_id, model_id=None) + location = kumo.ModelRegistryLocation('churn', alias='champion') + + first = _register(target, client, info, location, + {'kumo_bundle_id': 'a'}) + with caplog.at_level('INFO'): + again = _register(target, client, info, location, + {'kumo_bundle_id': 'a'}) + other = _register(target, client, info, location, + {'kumo_bundle_id': 'b'}) + + assert (first.version, again.version, other.version) == ('1', '1', '2') + assert 'churn version 1 already holds bundle a; reused it' in caplog.text + assert len(client.search_model_versions("name = 'churn'")) == 2 + assert str(client.get_model_version_by_alias('churn', + 'champion').version) == '2' + + +def test_the_smoke_script_compiles(tmp_path): + import py_compile + from pathlib import Path + + for parent in Path(__file__).resolve().parents: + script = parent / 'examples' / 'mlops' / 'smoke.py' + if script.is_file(): + py_compile.compile(str(script), cfile=str(tmp_path / 'out'), + doraise=True) + return + pytest.skip('examples/mlops is not in this checkout') diff --git a/packages/kumo-engine/test/test_mlflow_tracking.py b/packages/kumo-engine/test/test_mlflow_tracking.py new file mode 100644 index 000000000..b28f30b4c --- /dev/null +++ b/packages/kumo-engine/test/test_mlflow_tracking.py @@ -0,0 +1,404 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json +import os +import re +import socket +import threading +import time +import warnings +from pathlib import Path + +import numpy as np +import pandas as pd +import pytest +from kumo_core.artifacts import PROVENANCE_FILE +from kumoapi.encoder import Numerical +from kumoapi.model_plan import ModelPlan +from test_tracking_contract import CONTEXT, join_tracking_threads, replay + +import kumo.engine as kumo +from kumo.engine._scratch import get_scratch_dir +from kumo.engine._tracking import run_record + +mlflow = pytest.importorskip('mlflow') +autologging_utils = pytest.importorskip('mlflow.utils.autologging_utils') + +QUERY = 'PREDICT ents.Y FOR EACH ents.ID' + + +@pytest.fixture +def tracking_uri(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) # the default artifact root is ./mlruns + previous = mlflow.get_tracking_uri() + uri = f'sqlite:///{tmp_path / "mlflow.db"}' + mlflow.set_tracking_uri(uri) + try: + yield uri + finally: + while mlflow.active_run() is not None: + mlflow.end_run() + mlflow.set_tracking_uri(previous) + + +@pytest.fixture +def table_and_snapshot(tmp_path, tracking_uri): + previous_scratch = get_scratch_dir() + kumo.configure(scratch_dir=tmp_path / 'scratch') + try: + rng = np.random.default_rng(0) + (tmp_path / 'data').mkdir() + pd.DataFrame({ + 'ID': np.arange(400), + 'Y': rng.random(400), + 'F1': rng.normal(size=400), + 'amount ($)': rng.normal(size=400), + }).to_parquet(tmp_path / 'data' / 'ents.parquet') + connector = kumo.FileConnector(root_dir=str(tmp_path / 'data')) + ents = kumo.Table.from_source_table(connector['ents.parquet'], + primary_key='ID').infer_metadata() + graph = kumo.Graph(tables={'ents': ents}, edges=[]) + snapshot = graph.snapshot() + query = kumo.PredictiveQuery(graph=graph, query=QUERY) + yield query.generate_training_table(snapshot), snapshot + finally: + kumo.configure(scratch_dir=previous_scratch) + + +def _runs(client): + return client.search_runs( + [e.experiment_id for e in client.search_experiments()]) + + +def test_fit_logs_a_run_per_fit_and_per_trial(table_and_snapshot, tmp_path, + caplog): + from kumo.engine.integrations.mlflow import MlflowTracking + + table, snapshot = table_and_snapshot + client = mlflow.MlflowClient() + plan = ModelPlan() + plan.training_job.num_experiments = 2 + plan.training_job.disable_explain = True + plan.optimization.max_epochs = 2 + plan.optimization.max_steps_per_epoch = 2 + plan.optimization.base_lr = [0.01, 0.001] + plan.column_processing.encoder_overrides = { + 'ents.amount ($)': Numerical('standard') + } + mlflow.pytorch.autolog() + try: + with mlflow.start_run() as outer: + tracking = MlflowTracking(experiment_name='ents') + mlops = kumo.MLOps(artifact_root=str(tmp_path / 'root'), + tracking=tracking) + result = kumo.Trainer(plan).fit(table, snapshot=snapshot, + mlops=mlops) + assert mlflow.active_run().info.run_id == outer.info.run_id + assert autologging_utils._AUTOLOGGING_GLOBALLY_DISABLED is False + finally: + mlflow.pytorch.autolog(disable=True) + + runs = _runs(client) + parent, = [ + run for run in runs + if run.data.tags.get('mlflow.parentRunId') == outer.info.run_id + ] + assert tracking.run_id == parent.info.run_id + assert re.fullmatch(r'fit-\d{8}T\d{6}Z', parent.info.run_name) + children = sorted( + (run for run in runs + if run.data.tags.get('mlflow.parentRunId') == parent.info.run_id), + key=lambda run: run.info.run_name) + assert len(runs) == 4 + assert client.get_experiment(parent.info.experiment_id).name == 'ents' + assert parent.info.status == 'FINISHED' + assert {k for k in parent.data.tags if k.startswith('kumo.')} == { + 'kumo.kind', 'kumo.run_id', 'kumo.bundle_id', 'kumo.dropped_params'} + assert parent.data.tags['kumo.bundle_id'] == result.id + stored = tmp_path / 'root' / 'models' / result.id / PROVENANCE_FILE + assert json.loads(stored.read_text())['tracking'] == { + 'run_id': parent.data.tags['kumo.run_id']} + logged = client.download_artifacts(parent.info.run_id, PROVENANCE_FILE, + str(tmp_path)) + assert json.loads(Path(logged).read_text())['id'] == result.id + assert int(parent.data.tags['kumo.dropped_params']) >= 1 + params = parent.data.params + assert json.loads(params['optimization.base_lr']) == [0.01, 0.001] + assert not any('$' in name for name in params) + metrics = { + name: value + for name, value in parent.data.metrics.items() + if not name.startswith('system/') + } + assert metrics == {**result.metrics, 'seconds': metrics['seconds']} + assert 'off for the rest of this run' not in caplog.text + ents, = parent.inputs.dataset_inputs + assert (ents.dataset.name, ents.dataset.source_type) == ('ents', 'local') + assert json.loads(ents.dataset.source) == { + 'uri': str((tmp_path / 'data' / 'ents.parquet').resolve()) + } + + assert [run.info.run_name + for run in children] == ['trial-1-attempt-1', 'trial-2-attempt-1'] + for child in children: + assert child.info.status == 'FINISHED' + assert child.data.tags['kumo.status'] == 'succeeded' + assert {'best_val_mae', 'seconds', + 'val_mae'} <= set(child.data.metrics) + history = client.get_metric_history(child.info.run_id, 'val_mae') + assert sorted(metric.step for metric in history) == [1, 2] + + +def test_failed_fit_marks_the_run_failed(table_and_snapshot, tmp_path, + monkeypatch): + import kumo_core.training.orchestration as orchestration + + from kumo.engine.integrations.mlflow import MlflowTracking + + def fail(*args, **kwargs): + raise ValueError('stop') + + monkeypatch.setattr(orchestration, 'run_callable_training', fail) + table, snapshot = table_and_snapshot + + with kumo.mlops(kumo.MLOps(artifact_root=str(tmp_path / 'root'), + tracking=MlflowTracking())): + with pytest.raises(ValueError, match='stop'): + kumo.Trainer(ModelPlan()).fit(table, snapshot=snapshot) + + run, = _runs(mlflow.MlflowClient()) + assert run.info.status == 'FAILED' + assert run.data.tags['kumo.error'] == 'ValueError: stop' + assert mlflow.active_run() is None + + +def test_mlflow_error_ends_runs_failed_and_is_not_raised( + tracking_uri, monkeypatch, caplog): + from mlflow.tracking import MlflowClient + + from kumo.engine.integrations.mlflow import MlflowTracking + + original = MlflowClient.log_batch + + def down(self, *args, metrics=(), **kwargs): + if metrics: + raise RuntimeError('tracking server down') + return original(self, *args, **kwargs) + + monkeypatch.setattr(MlflowClient, 'log_batch', down) + replay(MlflowTracking(log_system_metrics=False)) + + assert caplog.text.count('off for the rest of this run') == 1 + assert mlflow.active_run() is None + statuses = {run.info.status for run in _runs(mlflow.MlflowClient())} + assert statuses == {'FAILED'} + + +def test_a_user_run_name_wins_over_the_default(tracking_uri): + from kumo.engine.integrations.mlflow import MlflowTracking + + replay(MlflowTracking(run_name='nightly', log_system_metrics=False)) + join_tracking_threads() + + names = {run.info.run_name for run in _runs(mlflow.MlflowClient()) + if 'mlflow.parentRunId' not in run.data.tags} + assert names == {'nightly'} + + +SAGEMAKER = ('arn:aws:sagemaker:us-west-2:account:' + 'mlflow-tracking-server/demo') + + +@pytest.mark.parametrize('uri,good,bad', [ + (SAGEMAKER, ['users-churn', 'u--1', 'a' * 57], + ['users_churn', 'main.ml.churn', 'bad/name:x', '-churn', 'a' * 58]), + ('arn:aws:sagemaker:us-east-1:account:mlflow-app/app-1', + ['users-churn'], ['users_churn']), + ('databricks-uc', ['main.ml.users_churn', 'ml.users-churn', 'churn'], + ['main.ml.users churn', 'a.b.c.d', 'main..churn', 'main.ml.a/b']), + ('sqlite:///mlflow.db', ['bad/name:x', 'users_churn', 'a.b.c.d'], []), +]) +def test_registry_model_names_are_checked_when_mlops_is_built( + tmp_path, uri, good, bad): + from kumo.engine.integrations.mlflow import ( + MlflowTracking, + check_model_name, + ) + + for name in good: + check_model_name(name, uri) + tracking = MlflowTracking(log_system_metrics=False, + tracking_uri=f'sqlite:///{tmp_path}/m.db', + experiment_name='e') + tracking.target.registry_uri = uri + for name in bad: + with pytest.raises(ValueError, match="e.g. '"): + kumo.MLOps(str(tmp_path), tracking=tracking, + registry=kumo.ModelRegistryLocation(name)) + + +@pytest.mark.parametrize('uri,command', [ + (SAGEMAKER, 'aws sagemaker create-presigned-mlflow-tracking-server-url ' + '--tracking-server-name demo --region us-west-2'), + ('sqlite:///mlflow.db', None), +]) +def test_a_tracked_fit_on_sagemaker_says_how_to_open_the_ui( + tracking_uri, uri, command): + from kumo.engine.integrations.mlflow import MlflowTracking + + tracker = MlflowTracking(log_system_metrics=False) + tracker.start(run_record(CONTEXT)) + tracker.target.tracking_uri = uri + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter('always') + tracker.finish(run_record(CONTEXT), None) + hints = [str(w.message) for w in caught if 'Studio' in str(w.message)] + assert hints == ([] if command is None else [ + "MLflow's run links do not open on SageMaker: open the MLflow UI " + f'from SageMaker Studio or at the URL `{command}` prints.']) + + +def test_user_disabled_autologging_stays_disabled(tracking_uri): + from kumo.engine.integrations.mlflow import MlflowTracking + + tracker = MlflowTracking(log_system_metrics=False) + with autologging_utils.disable_autologging(): + tracker.start(run_record(CONTEXT)) + tracker.finish(None, None) + assert autologging_utils._AUTOLOGGING_GLOBALLY_DISABLED is True + + +@pytest.fixture +def hung_server(): + server = socket.create_server(('127.0.0.1', 0)) + held = [] + + def accept(): + while True: + try: + held.append(server.accept()[0]) + except OSError: + return + + threading.Thread(target=accept, daemon=True).start() + try: + yield f'http://127.0.0.1:{server.getsockname()[1]}' + finally: + server.close() + for connection in held: + connection.close() + join_tracking_threads() + + +def test_a_hung_server_does_not_stall_the_fit(hung_server, monkeypatch): + import kumo.engine._tracking as tracking + from kumo.engine.integrations.mlflow import MlflowTracking + + monkeypatch.setattr(tracking, 'FINISH_SECONDS', 0.5) + monkeypatch.setenv('MLFLOW_HTTP_REQUEST_MAX_RETRIES', '0') + monkeypatch.delenv('MLFLOW_HTTP_REQUEST_TIMEOUT', raising=False) + previous = mlflow.get_tracking_uri() + mlflow.set_tracking_uri(hung_server) + try: + began = time.monotonic() + replay(MlflowTracking(log_system_metrics=False)) + elapsed = time.monotonic() - began + limits = {name: os.environ.get(name) for name in ( + 'MLFLOW_HTTP_REQUEST_MAX_RETRIES', 'MLFLOW_HTTP_REQUEST_TIMEOUT')} + finally: + mlflow.set_tracking_uri(previous) + + assert elapsed < 1.5 + assert limits == {'MLFLOW_HTTP_REQUEST_MAX_RETRIES': '0', + 'MLFLOW_HTTP_REQUEST_TIMEOUT': '30'} + + +def test_source_tables_are_logged_as_native_inputs( + tracking_uri, monkeypatch): + from mlflow.data.delta_dataset_source import DeltaDatasetSource + + from kumo.engine.integrations.mlflow import MlflowTracking + + monkeypatch.setattr(DeltaDatasetSource, '_lookup_table_id', + lambda self, name: 'table-id') + replay(MlflowTracking(log_system_metrics=False)) + + run, = [run for run in _runs(mlflow.MlflowClient()) + if run.data.tags['kumo.kind'] == 'fit'] + inputs = { + item.dataset.name: (item.dataset.source_type, + json.loads(item.dataset.source)) + for item in run.inputs.dataset_inputs + } + assert inputs == { + 'ents': ('s3', {'uri': 's3://bucket/ents.parquet'}), + 'orders': ('delta_table', { + 'delta_table_name': 'main.sales.orders', + 'delta_table_version': 7, + 'is_databricks_uc_table': True, + 'delta_table_id': 'table-id', + }), + } + orders = next(item for item in run.inputs.dataset_inputs + if item.dataset.name == 'orders') + assert mlflow.data.get_source(orders).delta_table_id == 'table-id' + + +def test_fit_inputs_come_from_the_embedded_parents(): + from kumo.engine._manifest import ( + ModelManifest, + ModelParents, + SnapshotManifest, + SnapshotParents, + SourceTableEntry, + TrainingTableManifest, + TrainingTableParents, + stub, + ) + from kumo.engine._tracking import fit_inputs + + entry = {'connector': 'LOCAL_FILE', 'location': '/data/ents.parquet', + 'version': '1790000000', 'version_kind': 'timestamp'} + snapshot = SnapshotManifest(id='s', parents=SnapshotParents( + source_tables={'ents': SourceTableEntry(**entry, sql='SELECT 1')})) + custom = SourceTableEntry(**{**entry, 'version': None}) + table = TrainingTableManifest(id='t', parents=TrainingTableParents( + snapshot=snapshot, source_table=custom)) + model = ModelManifest(id='m', parents=ModelParents( + training_table=table, warm_start=ModelManifest( + id='w', parents=ModelParents(training_table=TrainingTableManifest( + id='u', parents=TrainingTableParents(snapshot=SnapshotManifest( + id='v', parents=SnapshotParents(source_tables={ + 'old': SourceTableEntry(**entry)})))))))) + + expected = {'ents': entry, 'training_table': {**entry, 'version': None}} + assert fit_inputs(table) == expected + assert fit_inputs(model) == expected + assert fit_inputs(TrainingTableManifest(id='t', parents=TrainingTableParents( + snapshot=stub('snapshot', 's'), source_table=custom))) == { + 'training_table': {**entry, 'version': None}} + assert fit_inputs(stub('model', 'm')) == {} + assert fit_inputs(None) == {} + + +def test_link_prediction_metric_names_are_logged(tracking_uri): + from kumo.engine._tracking import RunRecord + from kumo.engine.integrations.mlflow import MlflowTracking + + tracker = MlflowTracking(experiment_name='lp') + tags = {'kumo.kind': 'fit', 'kumo.run_id': 'r'} + tracker.start(RunRecord(name='r', tags=tags)) + tracker.finish( + RunRecord(name='r', tags=tags, metrics={ + 'test_map@10': 0.5, + 'val_precision@10': 0.25 + }), None) + + run = mlflow.get_run(tracker.run_id) + assert run.info.status == 'FINISHED' + assert run.data.metrics == { + 'test_map_at_10': 0.5, + 'val_precision_at_10': 0.25 + } diff --git a/packages/kumo-engine/test/test_mlops.py b/packages/kumo-engine/test/test_mlops.py index 3d3af18d9..ea9974e7d 100644 --- a/packages/kumo-engine/test/test_mlops.py +++ b/packages/kumo-engine/test/test_mlops.py @@ -32,6 +32,30 @@ def test_mlops_needs_an_artifact_root_it_can_write_files_to(tmp_path): kumo.MLOps(artifact_root='@DB.SCHEMA.STAGE/kumo') with pytest.raises(ValueError, match='Unknown store kinds'): kumo.MLOps(artifact_root=str(tmp_path), store={'models'}) + with pytest.raises(ValueError, match='on_error'): + kumo.MLOps(artifact_root=str(tmp_path), on_error='ignore') + + +def test_a_connector_registry_checks_the_name_when_mlops_is_built(tmp_path): + def check_model_location(connector, location): + if '-' in location.model_name: + raise ValueError('bad name') + + connector = SimpleNamespace(sink_backend=SimpleNamespace( + check_model_location=check_model_location)) + good = kumo.ModelRegistryLocation('USERS_CHURN', alias='candidate') + kumo.MLOps(artifact_root=str(tmp_path), + registry=kumo.ModelOutputConfig(connector, location=good)) + with pytest.raises(ValueError, match='bad name'): + kumo.MLOps(artifact_root=str(tmp_path), + registry=kumo.ModelOutputConfig( + connector, + location=kumo.ModelRegistryLocation('users-churn'))) + with pytest.raises(ValueError, match='model alias'): + kumo.MLOps(artifact_root=str(tmp_path), + registry=kumo.ModelOutputConfig( + connector, location=kumo.ModelRegistryLocation( + 'USERS_CHURN', alias='LATEST'))) def test_a_root_is_a_uri_or_a_connector_and_path(tmp_path): diff --git a/packages/kumo-engine/test/test_optional_imports.py b/packages/kumo-engine/test/test_optional_imports.py index c70368fed..c18d0310d 100644 --- a/packages/kumo-engine/test/test_optional_imports.py +++ b/packages/kumo-engine/test/test_optional_imports.py @@ -26,6 +26,7 @@ 's3fs', 'google', 'gcsfs', + 'mlflow', ) PROVIDER_PACKAGES = ( 'kumo_snowflake', @@ -114,6 +115,22 @@ def find_spec(self, name, path=None, target=None): # and only the stand-in `self` is missing. Fine: nothing # broke at import time. pass + + from kumo.engine.integrations.mlflow import MlflowTracking + try: + MlflowTracking() + except ImportError as e: + assert "kumo-engine[mlflow]" in str(e), e + else: + raise AssertionError("MlflowTracking() needs mlflow") + + from kumo.engine.integrations.snowflake import SnowflakeTracking + try: + SnowflakeTracking(None, experiment_name="ents") + except ImportError as e: + assert "kumo-engine[snowflake]" in str(e), e + else: + raise AssertionError("SnowflakeTracking() needs snowflake-ml-python") print("ok") ''') diff --git a/packages/kumo-engine/test/test_prediction_completion.py b/packages/kumo-engine/test/test_prediction_completion.py index 526d87a3e..255ba9a7a 100644 --- a/packages/kumo-engine/test/test_prediction_completion.py +++ b/packages/kumo-engine/test/test_prediction_completion.py @@ -13,6 +13,7 @@ from kumoapi.model_plan import ModelPlan from kumo.engine.pquery.prediction_table import PredictionTable +from kumo.engine.trainer._scoring import ScoringDefaults from kumo.engine.trainer.trainer import Trainer @@ -33,7 +34,8 @@ def prepared_prediction(snapshot, tmp_path, monkeypatch): from kumo_core.types import TaskType trainer = Trainer(ModelPlan()) - trainer._training_result = SimpleNamespace(artifact_paths={}) + trainer._training_result = SimpleNamespace( + artifact_paths={}, scoring_defaults=ScoringDefaults()) trainer._snapshot_id = snapshot.id trainer._training_query = 'PREDICT users.label FOR EACH users.id' table = PredictionTable(trainer._training_query, tmp_path, tmp_path, 8, diff --git a/packages/kumo-engine/test/test_prediction_workers.py b/packages/kumo-engine/test/test_prediction_workers.py index 2b9e7b7c7..d06ee9c92 100644 --- a/packages/kumo-engine/test/test_prediction_workers.py +++ b/packages/kumo-engine/test/test_prediction_workers.py @@ -14,6 +14,7 @@ from kumo_core.types.enums import PredictionArtifactType from kumoapi.model_plan import ModelPlan +from kumo.engine.trainer._scoring import ScoringDefaults from kumo.engine.trainer.trainer import Trainer @@ -80,7 +81,8 @@ def test_bad_class_count_fails_before_prediction_work(count, monkeypatch): def test_valid_class_count_reaches_scorer(task_type, count, tmp_path, monkeypatch): trainer = Trainer(ModelPlan(), accelerator='cpu') - trainer._training_result = SimpleNamespace(artifact_paths={}) + trainer._training_result = SimpleNamespace( + artifact_paths={}, scoring_defaults=ScoringDefaults()) trainer._training_query = 'trained query' monkeypatch.setattr(Trainer, 'scratch_dir', property(lambda self: tmp_path)) monkeypatch.setattr( diff --git a/packages/kumo-engine/test/test_scoring_defaults.py b/packages/kumo-engine/test/test_scoring_defaults.py new file mode 100644 index 000000000..6ccfdf08e --- /dev/null +++ b/packages/kumo-engine/test/test_scoring_defaults.py @@ -0,0 +1,139 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import hashlib +import shutil +from types import SimpleNamespace + +import numpy as np +import pandas as pd +import pytest +from kumo_core.artifacts import add_bundle_file, remove_bundle_file +from kumo_core.types import TaskType +from kumoapi.plugin import PublishedModel + +import kumo.engine as kumo +from kumo.engine.trainer._scoring import normalized_scoring_defaults + + +def _frame(model, table, snapshot, **kwargs): + result = model.predict(table, snapshot=snapshot, **kwargs) + return result.frame.sort_values('ENTITY').reset_index(drop=True) + + +def _digest(directory): + return { + str(path.relative_to(directory)): + hashlib.sha256(path.read_bytes()).hexdigest() + for path in sorted(directory.rglob('*')) if path.is_file() + } + + +def _copy(trainer, destination): + shutil.copytree(trainer.result.bundle_dir, destination) + return destination + + +def test_derived_model_thresholds_unless_explicit(fitted): + trainer, table, snapshot = fitted + parent = trainer.result + original = _digest(parent.bundle_dir) + probs = _frame(parent, table, snapshot, + binary_classification_threshold=.5)['True_PROB'] + threshold = float(probs.median()) + child = parent.with_scoring_defaults( + kumo.ScoringDefaults(binary_classification_threshold=threshold)) + + implicit = _frame(child, table, snapshot) + pd.testing.assert_frame_equal( + implicit, + _frame(parent, table, snapshot, + binary_classification_threshold=threshold)) + assert implicit['TARGET_PRED'].any() + assert not _frame(child, table, snapshot, + binary_classification_threshold=1.0)['TARGET_PRED'].any() + + assert _digest(parent.bundle_dir) == original + assert parent.scoring_defaults == kumo.ScoringDefaults() + assert child.id != parent.id + assert child.manifest.id == child.id + assert child.manifest.parents.training_table is None + derived_from = child.manifest.parents.derived_from + assert derived_from.id == parent.id + assert derived_from.parents.training_table == ( + parent.manifest.parents.training_table) + cleared = child.with_scoring_defaults(kumo.ScoringDefaults()) + assert cleared.id == parent.id + assert cleared.manifest.parents.derived_from.parents.derived_from == ( + kumo.ModelManifest(id=parent.id, reconstructed=True)) + + +def test_equal_defaults_give_equal_bundle_ids(fitted): + trainer, *_ = fitted + children = [ + trainer.result.with_scoring_defaults( + kumo.ScoringDefaults(binary_classification_threshold=value)) + for value in (1, np.float32(1)) + ] + assert children[0].id == children[1].id + assert children[0].scoring_defaults.binary_classification_threshold == 1.0 + + +def test_parent_without_a_manifest_reads_as_a_stub(fitted, tmp_path): + trainer, *_ = fitted + old = _copy(trainer, tmp_path / 'old') + remove_bundle_file(old, 'provenance') + child = kumo.Trainer.load(old).result.with_scoring_defaults( + kumo.ScoringDefaults(binary_classification_threshold=.5)) + assert child.manifest.parents.derived_from == kumo.ModelManifest( + id=trainer.result.id, reconstructed=True) + + +def test_defaults_survive_export_and_connector_load(fitted, tmp_path): + trainer, *_ = fitted + child = trainer.result.with_scoring_defaults( + kumo.ScoringDefaults(binary_classification_threshold=.4)) + registry = kumo.FileConnector(root_dir=str(tmp_path)) + child._trainer.export(kumo.ModelOutputConfig(registry, path='m1')) + loaded = kumo.Trainer.load_from_connector(registry, path='m1') + assert loaded.result.scoring_defaults == child.scoring_defaults + assert loaded.result.id == child.id + + +def test_derived_export_publishes_the_child_bundle_id(fitted): + trainer, *_ = fitted + child = trainer.result.with_scoring_defaults( + kumo.ScoringDefaults(binary_classification_threshold=.4)) + calls = [] + sink = SimpleNamespace(publish_model=lambda *a, **k: calls.append(k) or + PublishedModel(name='m', version='1', uri='u')) + child._trainer.export( + kumo.ModelOutputConfig(SimpleNamespace(sink_backend=sink), path='m')) + assert calls[0]['metadata']['kumo_bundle_id'] == child.id + + +def test_defaults_are_validated_against_the_task(fitted): + trainer, *_ = fitted + with pytest.raises(ValueError, match='num_classes_to_return'): + trainer.result.with_scoring_defaults( + kumo.ScoringDefaults(num_classes_to_return=2)) + assert normalized_scoring_defaults( + kumo.ScoringDefaults(.3, np.int64(2)), + TaskType.MULTILABEL_CLASSIFICATION) == kumo.ScoringDefaults(.3, 2) + + +def test_unlisted_or_unreadable_defaults_are_ignored(fitted, tmp_path): + trainer, table, snapshot = fitted + bundle = _copy(trainer, tmp_path / 'old') + (bundle / 'scoring-defaults.json').write_text('not json') + old = kumo.Trainer.load(bundle) + assert old.result.scoring_defaults == kumo.ScoringDefaults() + assert old.result.id == trainer.result.id + + add_bundle_file(bundle, 'scoring_defaults', 'scoring-defaults.json') + with pytest.warns(UserWarning, match='Ignoring the scoring defaults'): + ignored = kumo.Trainer.load(bundle) + assert ignored.result.scoring_defaults == kumo.ScoringDefaults() + with pytest.raises(ValueError, match='must be specified'): + ignored.predict(table, snapshot=snapshot) diff --git a/packages/kumo-engine/test/test_semantic_carry_forward.py b/packages/kumo-engine/test/test_semantic_carry_forward.py index fc5788a9b..5968ef805 100644 --- a/packages/kumo-engine/test/test_semantic_carry_forward.py +++ b/packages/kumo-engine/test/test_semantic_carry_forward.py @@ -24,6 +24,7 @@ PredictionTable, _compile_prediction_table, ) +from kumo.engine.trainer._scoring import ScoringDefaults from kumo.engine.trainer.trainer import Trainer, _stage_model_bundle _SPLIT = SimpleNamespace(stage_counts={'test': 1}) @@ -136,7 +137,7 @@ def test_prediction_admission_before_materialization(snapshot, tmp_path, query, monkeypatch): from kumoapi.model_plan import ModelPlan trainer = Trainer(ModelPlan()) - trainer._training_result = SimpleNamespace() + trainer._training_result = SimpleNamespace(scoring_defaults=ScoringDefaults()) trainer._snapshot_id = snapshot.id trainer._training_query = 'PREDICT users.label FOR EACH users.id' materialize = Mock(side_effect=AssertionError('must not materialize')) @@ -164,7 +165,8 @@ def test_binary_predictions_preserve_probabilities_and_metadata_contract( _stub_generated_task_dataset_builder(monkeypatch) monkeypatch.setattr(type(trainer), 'scratch_dir', property(lambda _: tmp_path)) - trainer._training_result = SimpleNamespace(artifact_paths={}) + trainer._training_result = SimpleNamespace( + artifact_paths={}, scoring_defaults=ScoringDefaults()) trainer._snapshot_id = snapshot.id trainer._training_query = 'PREDICT users.label FOR EACH users.id' table = PredictionTable(trainer._training_query, tmp_path, tmp_path, 0, @@ -273,7 +275,8 @@ def stats(index_path, table_name, mode): trainer = Trainer(ModelPlan()) monkeypatch.setattr(type(trainer), 'scratch_dir', property(lambda _: tmp_path)) - trainer._training_result = SimpleNamespace(artifact_paths={}) + trainer._training_result = SimpleNamespace( + artifact_paths={}, scoring_defaults=ScoringDefaults()) trainer._snapshot_id = snapshot.id trainer._training_query = query @@ -708,7 +711,8 @@ def run_predict() -> None: trainer = Trainer(ModelPlan()) monkeypatch.setattr(type(trainer), 'scratch_dir', property(lambda _: tmp_path)) - trainer._training_result = SimpleNamespace(artifact_paths={}) + trainer._training_result = SimpleNamespace( + artifact_paths={}, scoring_defaults=ScoringDefaults()) trainer._snapshot_id = snapshot.id trainer._training_query = query trainer.predict(table, snapshot=snapshot, output_config=output, @@ -720,7 +724,8 @@ def run_predict() -> None: same_trainer = Trainer(ModelPlan()) monkeypatch.setattr(type(same_trainer), 'scratch_dir', property(lambda _: tmp_path)) - same_trainer._training_result = SimpleNamespace(artifact_paths={}) + same_trainer._training_result = SimpleNamespace( + artifact_paths={}, scoring_defaults=ScoringDefaults()) same_trainer._snapshot_id = snapshot.id same_trainer._training_query = query same_trainer.predict(table, snapshot=snapshot, output_config=output, diff --git a/packages/kumo-engine/test/test_snowflake_connector.py b/packages/kumo-engine/test/test_snowflake_connector.py index 1e4fcc687..b2b9254c8 100644 --- a/packages/kumo-engine/test/test_snowflake_connector.py +++ b/packages/kumo-engine/test/test_snowflake_connector.py @@ -108,3 +108,18 @@ def test_a_missing_target_field_with_no_override_is_reported( with pytest.raises(ValueError, match="names no warehouse"): SnowflakeConnector.from_spcs_env() + + +def test_session_is_the_readers_session(monkeypatch): + from kumo_snowflake import session as sessions + + connector = SnowflakeConnector(warehouse='WH', database='DB', + schema_name='S') + calls = [] + monkeypatch.setattr( + sessions._sessions, 'get', + lambda principal, target, fresh=False: calls.append(target) or 'live') + monkeypatch.setattr(type(connector.source_backend), '_principal', + lambda self, target: 'principal') + assert connector.session() == 'live' + assert calls == [connector] diff --git a/packages/kumo-engine/test/test_snowflake_tracking.py b/packages/kumo-engine/test/test_snowflake_tracking.py new file mode 100644 index 000000000..3ca74f771 --- /dev/null +++ b/packages/kumo-engine/test/test_snowflake_tracking.py @@ -0,0 +1,250 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import inspect +import threading +import time +from unittest import mock + +import _snowflake_experiment_fake +import pytest +from _snowflake_experiment_fake import FakeExperimentTracking, FakeSession +from test_tracking_contract import ( + CONTEXT, + EVENTS, + RESULT, + join_tracking_threads, +) + +from kumo.engine._tracking import TrackingSession + +FAKED = ('__init__', 'set_experiment', 'start_run', 'end_run', 'log_params', + 'log_metrics', 'log_artifact') + + +@pytest.mark.parametrize('method', FAKED) +def test_fake_matches_installed_experiment_tracking(method): + experiment = pytest.importorskip('snowflake.ml.experiment') + real = inspect.signature(getattr(experiment.ExperimentTracking, method)) + fake = inspect.signature(getattr(FakeExperimentTracking, method)) + assert [(p.name, p.kind, p.default) for p in fake.parameters.values() + ] == [(p.name, p.kind, p.default) for p in real.parameters.values()] + + +def test_installed_singleton_exposes_what_the_adapter_reads(monkeypatch): + experiment = pytest.importorskip('snowflake.ml.experiment') + cls = experiment.ExperimentTracking + monkeypatch.setattr(cls, '_instance', None) + session = mock.MagicMock() + session.get_current_database.return_value = 'DB' + session.get_current_schema.return_value = 'PUBLIC' + + tracking = cls(session) + + assert cls(mock.MagicMock()) is tracking + assert tracking._run is None + assert tracking._session is session + assert 'experiment_name' in inspect.signature( + tracking._sql_client.show_runs_in_experiment).parameters + + +@pytest.fixture +def fake(monkeypatch): + return _snowflake_experiment_fake.install(monkeypatch) + + +def _tracking(session=FakeSession, **kwargs): + from kumo.engine.integrations.snowflake import SnowflakeTracking + return SnowflakeTracking( + session() if session is FakeSession else session, + experiment_name='ents', **kwargs) + + +def _replay(tracking, result=RESULT, error=None, during=None): + session = TrackingSession(tracking) + session.start(lambda: CONTEXT) + for event in EVENTS: + session.on_event(event) + if during is not None: + caught_up = threading.Event() + session._submit(caught_up.set) + caught_up.wait(30) + during() + session.finish(result, error) + + +def test_native_model_lineage_holds_the_run_open_during_the_fit(fake): + lineage = [] + _replay( + _tracking(database='ml', schema='runs', run_name='fit', + native_model_lineage=True), + during=lambda: lineage.extend(fake._instance.lineage_stack)) + + assert lineage == [('ENTS', 'FIT')] + assert fake._instance.lineage_stack == [] + assert fake._instance.calls[1] == ('set_experiment', 'ents', 'ml', 'runs') + + +def test_snowflake_errors_warn_once_and_end_the_run_failed(fake, caplog): + session = FakeSession() + fake(session).fail_on = {'log_metrics'} + + _replay(_tracking(session)) + + tracking = fake._instance + assert caplog.text.count('off for the rest of this run') == 1 + run, = tracking.runs.values() + assert run['status'] == 'FAILED' + assert tracking._run is None + + +def test_failed_fit_ends_the_run_failed(fake): + _replay(_tracking(), None, ValueError('stop')) + + run, = fake._instance.runs.values() + assert run['status'] == 'FAILED' + + +@pytest.mark.parametrize('busy', ['active run', 'other session']) +def test_an_instance_in_use_is_left_alone(fake, busy, caplog): + session = FakeSession() + user = fake(session, warnings_as_errors=True) + user.set_experiment('mine') + if busy == 'active run': + user.start_run('mine_1') + else: + session = FakeSession() + calls = len(user.calls) + + _replay(_tracking(session)) + + assert user.calls[calls:] == [] + assert user.warnings_as_errors is True + assert caplog.text.count('not tracked in Snowflake') == 1 + + +def test_location_defaults_to_the_session_not_the_last_fit(fake): + session = FakeSession(database='DEV', schema='WORK') + _replay(_tracking(session, database='ml')) + _replay(_tracking(session)) + + assert [ + call[1:] for call in fake._instance.calls if call[0] == 'set_experiment' + ] == [('ents', 'ml', 'PUBLIC'), ('ents', 'DEV', 'WORK')] + assert not any(call[0].startswith('run.') for call in fake._instance.calls) + + +@pytest.mark.parametrize('taken', ['FIT', 'MYXRUN']) +def test_an_existing_run_is_never_resumed(fake, taken, caplog): + session = FakeSession() + user = fake(session) + user.set_experiment('ents') + user.start_run(taken) + if taken == 'FIT': + user.end_run() + else: + user._run = None # still running, in another process + run_name = 'fit' if taken == 'FIT' else 'my_run' + + tracking = _tracking(session, run_name=run_name) + _replay(tracking) + + new = f"{run_name}_{CONTEXT.run_id.rsplit('-', 1)[-1][:8]}".upper() + assert tracking.run_id == new + assert set(user.runs) == {taken, new} + assert user.runs[taken]['params'] == {} + assert user.runs[new]['status'] == 'FINISHED' + assert 'recording this fit as' in caplog.text + + +def test_metric_and_param_limits(fake, monkeypatch): + import kumo.engine.integrations.snowflake as snowflake + + monkeypatch.setattr(snowflake, '_EPOCH_METRICS', 2) + monkeypatch.setattr(snowflake, '_SUMMARY_METRICS', 3) + monkeypatch.setattr(snowflake, '_MAX_METRICS', 5) + result = type(RESULT)(metrics={**RESULT.metrics, 'r2': float('nan')}, + id=RESULT.id, bundle_dir=RESULT.bundle_dir) + _replay(_tracking(tags={'note': 'a$$b', 'level': 2}), result) + + run, = fake._instance.runs.values() + assert sorted({name for name, _, _ in run['metrics']}) == [ + 'test_mae', 'trial-1-attempt-1/best_val_mae', + 'trial-1-attempt-1/train_loss', 'trial-1-attempt-1/val_mae', 'val_mae' + ] + assert run['params']['kumo.dropped_metrics'] == '6' + assert run['params']['kumo.dropped_params'] == '1' + assert 'note' not in run['params'] + assert run['params']['level'] == '2' + + +def test_quoted_run_names_are_taken_only_by_their_exact_name(): + from kumo.engine.integrations.snowflake import _taken + + assert _taken('"MyRun"', ['MyRun']) + assert not _taken('"MyRun"', ['MYRUN']) + assert _taken('my_run', ['MYXRUN']) + + +def test_the_users_session_is_not_altered(fake): + session = FakeSession() + + _replay(_tracking(session)) + + assert session.statements == [] + + +def test_a_hung_snowflake_call_does_not_stall_the_fit(fake, monkeypatch, + caplog): + import kumo.engine._tracking as tracking + + monkeypatch.setattr(tracking, 'FINISH_SECONDS', 0.5) + release = threading.Event() + monkeypatch.setattr(FakeExperimentTracking, 'log_metrics', + lambda *args, **kwargs: release.wait(30)) + began = time.monotonic() + try: + _replay(_tracking()) + elapsed = time.monotonic() - began + finally: + release.set() + join_tracking_threads() + + assert elapsed < 1.5 + assert 'did not finish within' in caplog.text + run, = fake._instance.runs.values() + assert run['status'] == 'FINISHED' + + +class _SchemaSession: + def __init__(self, schemas): + self.schemas = schemas + self.statements = [] + + def get_current_database(self): + return 'DB' + + def get_current_schema(self): + return 'RUNS' + + def sql(self, statement): + self.statements.append(statement) + return mock.Mock(collect=lambda: self.schemas) + + +def test_a_connector_opens_the_session_and_the_schema_is_checked_first( + fake): + import kumo.engine as kumo + + session = _SchemaSession([('RUNS', )]) + tracking = _tracking(None, connector=mock.Mock(session=lambda: session)) + tracking.check() + assert tracking.session is session + assert session.statements == ["SHOW SCHEMAS LIKE 'RUNS' IN DATABASE DB"] + + tracking.session = _SchemaSession([]) + with pytest.raises(ValueError, match='needs schema DB.RUNS'): + kumo.MLOps('/tmp/root', tracking=tracking) + with pytest.raises(ValueError, match='a session or a connector'): + _tracking(None) diff --git a/packages/kumo-engine/test/test_tracking_contract.py b/packages/kumo-engine/test/test_tracking_contract.py new file mode 100644 index 000000000..2efe3cb35 --- /dev/null +++ b/packages/kumo-engine/test/test_tracking_contract.py @@ -0,0 +1,400 @@ +# 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.""" + +import json +import threading +import time +from dataclasses import dataclass, field +from pathlib import Path +from types import SimpleNamespace +from typing import Any, Callable, Dict, List, Mapping, Optional, Tuple + +import pytest +from kumo_core.artifacts import PROVENANCE_FILE + +from kumo.engine._tracking import ( + DROPPED_METRICS, + DROPPED_PARAMS, + FitContext, + TrackingSession, + run_record, +) +from kumo.engine.tracking import Event, RunRecord + +RUN = 'trainer-fit-1' +USER_TAGS = {'team': 'ents', 'kumo.run_id': 'overridden'} +QUERY = 'PREDICT ents.Y FOR EACH ents.ID' +INPUTS = { + 'ents': {'connector': 'S3', 'location': 's3://bucket/ents.parquet', + 'version': '1790000000', 'version_kind': 'timestamp'}, + 'orders': {'connector': 'DATABRICKS', + 'location': '`main`.`sales`.`orders`', + 'version': '7_1790000000', 'version_kind': 'exact'}, + 'users': {'connector': 'SNOWFLAKE', 'location': '"DB"."SCH"."USERS"', + 'version': '1790000000', 'version_kind': 'timestamp'}, +} +CONTEXT = FitContext( + run_id=RUN, query=QUERY, inputs=INPUTS, + model_plan={ + 'optimization': {'base_lr': [0.01, 0.001], 'max_epochs': 2}, + 'training_job': {'tune_metric': 'mae'}, + 'column_processing': { + 'encoder_overrides': {'ents.amount ($)': 'numerical'} + }, + }) + + +def _event(name, trial=None, epoch=None, metrics=None, **params): + return Event(name=name, run_id=RUN, time=0.0, trial=trial, epoch=epoch, + metrics=metrics or {}, params=params) + + +def _epochs(trial, *values): + return [ + _event('epoch_ended', trial, epoch, {'train_loss': loss, + 'val_mae': mae}) + for epoch, (loss, mae) in enumerate(values, 1) + ] + + +# Trial 2's first attempt fails; epoch events carry no attempt, as in core. +EVENTS = [ + _event('trial_started', 1, attempt=1, **{'optimization.base_lr': 0.01}), + *_epochs(1, (0.5, 0.4), (0.3, 0.35)), + _event('trial_ended', 1, metrics={'val_mae': 0.35, 'seconds': 3.0}, + status='succeeded', attempt=1), + _event('trial_started', 2, attempt=1, **{'optimization.base_lr': 0.001}), + _event('trial_ended', 2, status='failed', attempt=1, + error='OutOfMemoryError: CUDA out of memory'), + _event('trial_started', 2, attempt=2, **{'optimization.base_lr': 0.001, + 'optimization.batch_size': 256}), + *_epochs(2, (0.6, 0.5), (0.4, 0.3)), + _event('trial_ended', 2, metrics={'val_mae': 0.3, 'seconds': 4.0}, + status='succeeded', attempt=2), + _event('fit_ended', metrics={'val_mae': 0.3, 'test_mae': 0.32, + 'seconds': 9.0}, status='succeeded'), +] +BUNDLE = 'b' * 64 +RESULT = SimpleNamespace(metrics={'val_mae': 0.3, 'test_mae': 0.32}, + id=BUNDLE, bundle_dir=Path('no-bundle')) + + +class _Recorder: + def __init__(self, tracker=None): + self.calls: List[Tuple[str, Any]] = [] + self.tracker = tracker + self.tags = getattr(tracker, 'tags', None) + + def prepare(self): + getattr(self.tracker, 'prepare', lambda: None)() + + def start(self, record): + self.calls.append(('start', record)) + if self.tracker is not None: + self.tracker.start(record) + + def on_trial(self, record, event): + self.calls.append(('on_trial', record)) + if self.tracker is not None: + self.tracker.on_trial(record, event) + + def finish(self, record, error): + self.calls.append(('finish', record)) + if self.tracker is not None: + self.tracker.finish(record, error) + + +def join_tracking_threads(): + for thread in threading.enumerate(): + if thread.name == 'kumo-tracking': + thread.join(30) + + +def replay(tracker=None, events=EVENTS, result=RESULT, + error=None) -> Dict[str, RunRecord]: + r"""Feeds the recorded fit to `tracker`; returns the records expected + of it.""" + recorder = _Recorder(tracker) + session = TrackingSession(recorder) + session.start(lambda: CONTEXT) + for event in events: + session.on_event(event) + session.finish(result, error) + records = { + record.tags['kumo.run_id']: record + for hook, record in recorder.calls if hook == 'on_trial' + } + records[RUN] = run_record(CONTEXT, result, events[-1], tags=USER_TAGS) + return records + + +@dataclass +class Observed: + name: Optional[str] + params: Mapping[str, str] + metrics: Mapping[str, float] + steps: Mapping[str, Tuple[Tuple[int, float], ...]] + tags: Mapping[str, str] + parent: Optional[str] = None + artifacts: Mapping[str, str] = field(default_factory=dict) + + +@dataclass(frozen=True) +class Adapter: + make: Callable[[], Any] + observe: Callable[[], Dict[str, Observed]] + kept_params: Callable[[Mapping[str, str]], Mapping[str, str]] = dict + nested_runs: bool = False + step_metrics: bool = False + artifacts: bool = False + + +def gaps(expected: Dict[str, RunRecord], observed: Dict[str, Observed], + adapter: Adapter) -> List[str]: + kept_params = adapter.kept_params + found = [] + if set(observed) != set(expected): + return [f'runs {sorted(observed)} != {sorted(expected)}'] + dropped = sum(len(record.params) - len(kept_params(record.params)) + for record in expected.values()) + counted = sum(int(seen.tags.get(DROPPED_PARAMS, 0)) + for seen in observed.values()) + if counted != dropped: + found.append(f'{DROPPED_PARAMS}: {counted} != {dropped}') + for run_id, record in expected.items(): + seen = observed[run_id] + + def differ(what, have, want): + if have != want: + found.append(f'{run_id} {what}: {have!r} != {want!r}') + + if record.name is not None: + differ('name', seen.name, record.name) + differ('params', dict(seen.params), dict(kept_params(record.params))) + differ('metrics', {k: seen.metrics.get(k) for k in record.metrics}, + dict(record.metrics)) + names = set(record.metrics) + if adapter.step_metrics: + names |= set(record.steps) + differ('steps', {k: seen.steps.get(k) for k in record.steps}, + dict(record.steps)) + differ('metric names', sorted(seen.metrics), sorted(names)) + differ('tags', {k: seen.tags.get(k) for k in record.tags}, + dict(record.tags)) + differ( + 'kumo tags', + sorted(k for k in seen.tags if k.startswith('kumo.') + and k not in (DROPPED_PARAMS, DROPPED_METRICS)), + sorted(k for k in record.tags if k.startswith('kumo.'))) + if adapter.nested_runs: + differ('parent', seen.parent, record.tags.get('kumo.parent_run')) + if adapter.artifacts: + differ('artifacts', dict(seen.artifacts), dict(record.artifacts)) + return found + + +@pytest.fixture +def mlflow_adapter(tmp_path, monkeypatch): + mlflow = pytest.importorskip('mlflow') + from kumo.engine.integrations.mlflow import MlflowTracking, _mlflow_params + + monkeypatch.chdir(tmp_path) + previous = mlflow.get_tracking_uri() + mlflow.set_tracking_uri(f'sqlite:///{tmp_path / "mlflow.db"}') + + def observe(): + client = mlflow.MlflowClient() + runs = client.search_runs( + [e.experiment_id for e in client.search_experiments()]) + kumo_ids = {r.info.run_id: r.data.tags['kumo.run_id'] for r in runs} + + def history(run_id, key): + return tuple((m.step, m.value) for m in sorted( + client.get_metric_history(run_id, key), + key=lambda m: (m.step, m.timestamp))) + + return { + kumo_ids[run.info.run_id]: Observed( + name=run.info.run_name, params=run.data.params, + metrics=run.data.metrics, tags=run.data.tags, + steps={k: history(run.info.run_id, k) + for k in run.data.metrics}, + parent=kumo_ids.get(run.data.tags.get('mlflow.parentRunId')), + artifacts={ + a.path: mlflow.artifacts.load_text( + f'{run.info.artifact_uri}/{a.path}') + for a in client.list_artifacts(run.info.run_id) + }) + for run in runs + } + + try: + yield Adapter( + make=lambda: MlflowTracking(log_system_metrics=False, + tags=USER_TAGS), observe=observe, + kept_params=_mlflow_params, nested_runs=True, step_metrics=True, + artifacts=True) + finally: + while mlflow.active_run() is not None: + mlflow.end_run() + mlflow.set_tracking_uri(previous) + + +@pytest.fixture +def snowflake_adapter(monkeypatch): + import _snowflake_experiment_fake + + fake = _snowflake_experiment_fake.install(monkeypatch) + from kumo.engine.integrations.snowflake import SnowflakeTracking + + def observe(): + observed = {} + for run in fake._instance.runs.values(): + names: Dict[Optional[str], Dict[str, Dict[str, Any]]] = {} + + def bucket(key): + prefix, _, name = key.rpartition('/') + return names.setdefault(prefix or None, { + 'params': {}, 'tags': {}, 'metrics': {}, 'steps': {} + }), name + + for key, value in run['params'].items(): + runs, name = bucket(key) + # Snowflake has no tags: they are params of the same name. + if name.startswith('kumo.') or name in USER_TAGS: + runs['tags'][name] = value + else: + runs['params'][name] = value + for key, value, step in run['metrics']: + runs, name = bucket(key) + runs['metrics'][name] = value + runs['steps'][name] = (*runs['steps'].get(name, ()), + (step, value)) + for prefix, runs in names.items(): + observed[runs['tags']['kumo.run_id']] = Observed( + name=prefix, + artifacts=run['artifacts'] if prefix is None else {}, + **runs) + return observed + + return Adapter( + make=lambda: SnowflakeTracking(_snowflake_experiment_fake.FakeSession(), + experiment_name='ents', tags=USER_TAGS), + observe=observe, step_metrics=True, artifacts=True) + + +@pytest.fixture(params=['mlflow', 'snowflake']) +def adapter(request) -> Adapter: + return request.getfixturevalue(f'{request.param}_adapter') + + +def test_records_of_the_recorded_fit(): + records = replay() + + fit = records[RUN] + assert fit.tags == {'team': 'ents', 'kumo.kind': 'fit', 'kumo.run_id': RUN, + 'kumo.bundle_id': BUNDLE} + assert fit.params == { + 'optimization.base_lr': '[0.01, 0.001]', + 'optimization.max_epochs': '2', + 'training_job.tune_metric': 'mae', + 'column_processing.encoder_overrides.ents.amount ($)': 'numerical', + } + with pytest.warns(UserWarning, + match="'kumo.bundle_id', 'training_job.tune_metric'"): + spoofed = run_record(CONTEXT, tags={'kumo.bundle_id': 'x', + 'training_job.tune_metric': 'y', + 'a': 1}) + assert spoofed.tags.keys() - fit.tags.keys() == {'a'} + assert fit.metrics == {'val_mae': 0.3, 'test_mae': 0.32, 'seconds': 9.0} + assert fit.artifacts['query.pql'] == QUERY + assert json.loads(fit.artifacts['model_plan.json']) == CONTEXT.model_plan + assert fit.inputs == INPUTS + + oom = records[f'{RUN}-trial-2-attempt-1'] + assert oom == RunRecord( + name='trial-2-attempt-1', + params={'optimization.base_lr': '0.001', 'attempt': '1'}, tags={ + 'kumo.kind': 'trial', + 'kumo.run_id': f'{RUN}-trial-2-attempt-1', + 'kumo.parent_run': RUN, + 'kumo.status': 'failed', + 'kumo.error': 'OutOfMemoryError: CUDA out of memory', + }) + retry = records[f'{RUN}-trial-2-attempt-2'] + assert retry.params['attempt'] == '2' + assert retry.metrics == {'best_val_mae': 0.3, 'seconds': 4.0} + assert retry.steps == {'train_loss': ((1, 0.6), (2, 0.4)), + 'val_mae': ((1, 0.5), (2, 0.3))} + + +def test_adapter_records_the_same_runs(adapter, tmp_path): + (tmp_path / PROVENANCE_FILE).write_text('{"bundle": "b"}') + result = SimpleNamespace(metrics=RESULT.metrics, id=BUNDLE, + bundle_dir=tmp_path) + expected = replay(adapter.make(), result=result) + + assert expected[RUN].artifacts[PROVENANCE_FILE] == '{"bundle": "b"}' + assert gaps(expected, adapter.observe(), adapter) == [] + + +def test_a_raising_tracker_is_ended_and_not_called_again(caplog): + class RaisesOnTrial(_Recorder): + def on_trial(self, record, event): + super().on_trial(record, event) + raise RuntimeError('tracker down') + + broken = RaisesOnTrial() + replay(broken) + + assert [hook for hook, _ in broken.calls] == [ + 'start', 'on_trial', 'finish' + ] + assert broken.calls[-1] == ('finish', None) + assert caplog.text.count('off for the rest of this run') == 1 + + +def test_adapter_records_a_failed_fit(adapter): + failed = _event('fit_ended', status='failed', error='ValueError: stop') + expected = replay(adapter.make(), [*EVENTS[:-1], failed], None, + ValueError('stop')) + + assert expected[RUN].tags['kumo.error'] == 'ValueError: stop' + assert gaps(expected, adapter.observe(), adapter) == [] + + +def test_only_rank_zero_tracks(monkeypatch): + monkeypatch.setenv('LOCAL_RANK', '1') + recorder = _Recorder() + session = TrackingSession(recorder) + + session.start(lambda: CONTEXT) + session.finish(RESULT, None) + + assert recorder.calls == [] + + +def test_a_hung_tracker_does_not_stall_the_fit(monkeypatch, caplog): + import kumo.engine._tracking as tracking + + monkeypatch.setattr(tracking, 'FINISH_SECONDS', 0.5) + release = threading.Event() + + class Hung(_Recorder): + def start(self, record): + super().start(record) + release.wait(30) + + hung = Hung() + began = time.monotonic() + try: + replay(hung) + elapsed = time.monotonic() - began + finally: + release.set() + join_tracking_threads() + + assert elapsed < 1.5 + assert 'did not finish within' in caplog.text diff --git a/packages/kumo-engine/test/test_trainer_registry.py b/packages/kumo-engine/test/test_trainer_registry.py index 7ea9ba950..e20b5f387 100644 --- a/packages/kumo-engine/test/test_trainer_registry.py +++ b/packages/kumo-engine/test/test_trainer_registry.py @@ -29,6 +29,7 @@ SPEC_FILE, write_training_spec, ) +from kumo.engine.trainer.output_config import ModelOutputConfig from kumo.engine.trainer.trainer import Trainer, TrainingResult @@ -238,7 +239,11 @@ def test_finish_fit_forwards_publication_details_and_returns_assigned_version( class Sink: def publish_model(self, connector, bundle_dir, **kwargs): spec = json.loads((Path(bundle_dir) / SPEC_FILE).read_text()) - assert kwargs['metadata'] == { + assert { + key: value + for key, value in kwargs['metadata'].items() + if not key.startswith('kumo_') + } == { name: str(spec[name]) for name in ('spec_version', 'predictive_query', 'snapshot_id') } @@ -258,6 +263,7 @@ def publish_model(self, connector, bundle_dir, **kwargs): assert calls[0][2]['path'] == path assert calls[0][2]['location'] == location assert calls[0][2]['metrics'] == {'val_auroc': .9} + assert calls[0][2]['metadata']['kumo_bundle_id'] == result.bundle_id assert result.published_model is published assert result.metrics == original.metrics assert result.bundle_dir == bundle @@ -572,3 +578,182 @@ def fake_predict(**kwargs) -> None: binary_classification_threshold=.5) assert result.frame['TARGET_PRED'].tolist() == [True] + + +def _export(sink, bundle, location=None, artifact_paths=None): + return Trainer(ModelPlan())._finish_fit( + TrainingResult({}, artifact_paths or {}, + bundle), predictive_query='PREDICT COUNT(*)', + snapshot_id='snap-123', training_table_generation=_generation_plan(), + export_config=SimpleNamespace( + connector=SimpleNamespace(sink_backend=sink), + path=None if location else 'm', location=location)) + + +def _recording_sink(calls): + return SimpleNamespace(publish_model=lambda *a, **k: calls.append(k) or + PublishedModel(name='m', version='1', uri='u')) + + +def test_publish_records_only_the_bundle_id_pointer(tmp_path: Path) -> None: + calls = [] + result = _export(_recording_sink(calls), _fake_published_bundle(tmp_path)) + assert {k: v + for k, v in calls[0]['metadata'].items() + if k.startswith('kumo_')} == {'kumo_bundle_id': result.bundle_id} + + +class _AliasingSink(_FakeSink): + def __init__(self, recorded: Optional[str] = None, + fail: bool = False) -> None: + super().__init__() + self.recorded = recorded + self.fail = fail + self.aliases: dict = {} + + def set_model_alias(self, connector, location, alias) -> None: + if self.fail: + raise RuntimeError('registry unavailable') + self.aliases[alias] = location.version + + def fetch_model(self, connector, local_dir, **kwargs) -> PublishedModel: + shutil.copytree(self.published['CHURN_GNN@latest'], local_dir, + dirs_exist_ok=True) + return PublishedModel(name='CHURN_GNN', version='V1', uri='fake://', + metadata={'kumo_bundle_id': self.recorded}) + + +@pytest.mark.parametrize('sink,warning', [ + (_AliasingSink(), None), + (_AliasingSink(fail=True), 'registry unavailable'), + (_FakeSink(), 'cannot set model aliases'), +]) +def test_export_with_alias_publishes_then_sets_it_best_effort( + tmp_path: Path, sink, warning) -> None: + import warnings + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter('always') + result = _export(sink, _fake_published_bundle(tmp_path), + ModelRegistryLocation('CHURN_GNN', alias='prod')) + assert result.published_model.version == 'V1' + assert list(sink.published) == ['CHURN_GNN@latest'] + messages = [str(w.message) for w in caught if 'alias' in str(w.message)] + if warning is None: + assert sink.aliases == {'prod': 'V1'} and messages == [] + else: + assert len(messages) == 1 and warning in messages[0] + + +def test_load_by_alias_warns_on_recorded_bundle_mismatch( + tmp_path: Path) -> None: + bundle = _fake_published_bundle(tmp_path) + write_training_spec(bundle, predictive_query='PREDICT COUNT(*)', + model_plan=ModelPlan(), snapshot_id='snap-123', + training_table_generation=_generation_plan()) + sink = _AliasingSink('other') + sink.publish_model(None, str(bundle), + location=ModelRegistryLocation('CHURN_GNN')) + with pytest.warns(UserWarning, match='registry records bundle other'): + loaded = Trainer.load_from_connector( + _FakeConnector(sink), + location=ModelRegistryLocation('CHURN_GNN', alias='prod')) + assert loaded.result.id == Trainer.load(bundle).result.id + + +def test_training_result_exports_like_its_trainer(tmp_path: Path) -> None: + bundle = _fake_published_bundle(tmp_path) + write_training_spec(bundle, predictive_query='PREDICT COUNT(*)', + model_plan=ModelPlan(), snapshot_id='snap-123', + training_table_generation=_generation_plan()) + sink = _AliasingSink() + published = TrainingResult({}, {}, bundle).export( + ModelOutputConfig(_FakeConnector(sink), + location=ModelRegistryLocation('CHURN_GNN', + alias='prod'))) + assert published.version == 'V1' and sink.aliases == {'prod': 'V1'} + assert list(sink.published) == ['CHURN_GNN@latest'] + + +def test_alias_needs_an_aliasing_sink() -> None: + connector = _FakeConnector(_FakeSink()) + with pytest.raises(NotImplementedError, match='alias'): + Trainer.load_from_connector( + connector, location=ModelRegistryLocation('m', alias='prod')) + for kwargs in (dict(alias='prod', version='3'), dict(alias='')): + with pytest.raises(ValueError, match='alias'): + ModelRegistryLocation('m', **kwargs) + + +def test_invalid_export_alias_is_rejected_before_publishing( + tmp_path: Path) -> None: + sink = _AliasingSink() + trainer = Trainer(ModelPlan()) + trainer._training_result = TrainingResult({}, {}, + _fake_published_bundle(tmp_path)) + with pytest.raises(ValueError, match='model alias'): + trainer.export( + SimpleNamespace( + connector=SimpleNamespace(sink_backend=sink), path=None, + location=ModelRegistryLocation('CHURN_GNN', alias='LATEST'))) + assert sink.published == {} + + +@pytest.mark.parametrize('sink,registered', [ + (_AliasingSink(), 'V1'), + (SimpleNamespace(publish_model=lambda *a, **k: 1 / 0), None), +]) +def test_mlops_registry_registers_after_the_fit(tmp_path: Path, monkeypatch, + sink, registered) -> None: + import warnings + + import kumo.engine as kumo + + bundle = _fake_published_bundle(tmp_path) + write_training_spec(bundle, predictive_query='PREDICT COUNT(*)', + model_plan=ModelPlan(), snapshot_id='snap-123', + training_table_generation=_generation_plan()) + monkeypatch.setattr( + Trainer, '_fit', lambda self, *a, **k: TrainingResult( + {}, {}, bundle, bundle_id='b1')) + registry = kumo.ModelOutputConfig( + _FakeConnector(sink), + location=ModelRegistryLocation('CHURN_GNN', alias='prod')) + trainer = Trainer(ModelPlan()) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter('always') + result = trainer.fit( + SimpleNamespace(query='PREDICT COUNT(*)'), + snapshot=SimpleNamespace(id='s'), + mlops=kumo.MLOps(str(tmp_path / 'root'), registry=registry, + store=False)) + failures = [w for w in caught if 'Could not register' in str(w.message)] + assert getattr(result.registered_model, 'version', None) == registered + assert result.published_model is None + if registered: + assert trainer.result is result and not failures + assert sink.aliases == {'prod': 'V1'} + else: + assert len(failures) == 1 + assert 'ZeroDivisionError' in str(failures[0].message) + with pytest.raises(RuntimeError, match='ZeroDivisionError') as raised: + trainer.fit( + SimpleNamespace(query='PREDICT COUNT(*)'), + snapshot=SimpleNamespace(id='s'), + mlops=kumo.MLOps(str(tmp_path / 'root'), registry=registry, + store=False, on_error='raise')) + assert isinstance(raised.value.__cause__, ZeroDivisionError) + + +def test_mlops_registry_alias_is_checked_before_the_fit(tmp_path: Path, + monkeypatch) -> None: + import kumo.engine as kumo + + monkeypatch.setattr(Trainer, '_fit', lambda *a, **k: 1 / 0) + registry = kumo.ModelOutputConfig( + _FakeConnector(_AliasingSink()), + location=ModelRegistryLocation('CHURN_GNN', alias='LATEST')) + with pytest.raises(ValueError, match='model alias'): + Trainer(ModelPlan()).fit( + SimpleNamespace(query='q'), snapshot=SimpleNamespace(id='s'), + mlops=kumo.MLOps(str(tmp_path), registry=registry, store=False)) diff --git a/packages/kumo-snowflake/pyproject.toml b/packages/kumo-snowflake/pyproject.toml index 960687071..47547c8b9 100644 --- a/packages/kumo-snowflake/pyproject.toml +++ b/packages/kumo-snowflake/pyproject.toml @@ -27,13 +27,15 @@ dependencies = [ [project.optional-dependencies] # What the SDK plugin half needs beyond the platform reader: Snowpark for the -# session it opens, the model registry for what a sink publishes, and pyyaml -# for the manifest that carries. The platform installs this distribution +# session it opens, the model registry for what a sink publishes (and, from +# 2.0, the experiment tracking kumo-engine records fits in; capped below 3 +# because that adapter reads ExperimentTracking internals), and pyyaml for the +# manifest that carries. The platform installs this distribution # without this extra and imports `SnowflakeSourceReader` alone, so none of it # is loaded there -- `__init__` exports the plugin lazily to keep that true. sdk = [ "snowflake-snowpark-python", - "snowflake-ml-python", + "snowflake-ml-python>=2.0,<3", "pyyaml", ] test = [ diff --git a/packages/kumo-snowflake/src/kumo_snowflake/registry.py b/packages/kumo-snowflake/src/kumo_snowflake/registry.py index fd220653e..c4d88b580 100644 --- a/packages/kumo-snowflake/src/kumo_snowflake/registry.py +++ b/packages/kumo-snowflake/src/kumo_snowflake/registry.py @@ -33,11 +33,14 @@ """ import json +import logging +import re import shutil from contextlib import contextmanager from pathlib import Path from typing import ( TYPE_CHECKING, + Any, ContextManager, Dict, Iterator, @@ -53,9 +56,13 @@ resolved, ) from kumoapi.plugin import ( + BUNDLE_ID_KEY, + ENGINE_MARKER, ModelRegistryLocation, PublishedModel, SourceConnector, + check_model_alias, + registry_lineage, ) if TYPE_CHECKING: @@ -74,8 +81,11 @@ #: merely happens to be named alike. _MANIFEST_FILE = 'publication-manifest.json' -#: Engine-bundle marker stored through the registry metadata API. -_ENGINE_MARKER = 'kumo_engine' +_SYSTEM_ALIASES = frozenset({'DEFAULT', 'FIRST', 'LAST'}) +_IDENTIFIER = r'(?:[A-Za-z_][A-Za-z0-9_$]*|"(?:[^"]|"")+")' +_MODEL_NAME = re.compile(rf'(?:{_IDENTIFIER}\.){{0,2}}{_IDENTIFIER}') + +_LOG = logging.getLogger(__name__) #: Where a fetched version is unpacked before its bundle is lifted out. Named #: rather than a temporary directory so a half-finished import leaves something @@ -113,7 +123,11 @@ class _BundleModel(custom_model.CustomModel): Deliberately not runnable. Scoring a Kumo model needs a materialized DiskGraph task over the graph being scored, built where the data is, so this refuses rather than answering. The version exists to carry, - version and expose the bundle; :func:`kumoai_studio.import_model` runs it. + version and expose the bundle. Load it with + ``Trainer.load_from_connector(SnowflakeConnector(...), + location=ModelRegistryLocation(...))`` and predict against a snapshot + (kumo-engine), or with :func:`kumoai_studio.import_model` (hosted + service). """ # Unquoted, and evaluated here rather than deferred: the registry # compares this annotation against `pandas.DataFrame` itself, so a @@ -122,13 +136,36 @@ class _BundleModel(custom_model.CustomModel): def predict(self, df: pd.DataFrame) -> pd.DataFrame: raise NotImplementedError( "This model version carries a Kumo model bundle; it does not " - "score in Snowflake. Import it with kumoai_studio.import_model() and " - "predict against a graph.") + "score in Snowflake. Load it with Trainer.load_from_connector(" + "SnowflakeConnector(...), location=ModelRegistryLocation(...)) " + "and predict against a snapshot (kumo-engine), or " + "kumoai_studio.import_model() (hosted service).") return _BundleModel( custom_model.ModelContext(artifacts={'bundle': bundle_dir})) +#: Snowflake ML task types by Kumo task type; Model Monitor needs one. +_TASKS = { + 'binary_classification': 'TABULAR_BINARY_CLASSIFICATION', + 'multiclass_classification': 'TABULAR_MULTI_CLASSIFICATION', + 'regression': 'TABULAR_REGRESSION', +} + + +def _task(bundle: Path) -> Dict[str, Any]: + r"""``task=`` for ``log_model`` from the bundle's prediction schema; + empty for other tasks, or when it cannot be read.""" + try: + from snowflake.ml.model import type_hints + + schema = json.loads((bundle / 'prediction-schema.json').read_text()) + name = _TASKS.get(schema.get('task_type')) + return {} if name is None else {'task': getattr(type_hints.Task, name)} + except Exception: + return {} + + def _carrier_signature() -> Dict[str, "model_signature.ModelSignature"]: r"""Declares the carrier's signature instead of inferring one. @@ -209,6 +246,40 @@ def _refuse_path(path: Optional[str], action: str) -> None: f'a path.') +def _version_record(metadata: Mapping[str, str]) -> Dict[str, str]: + record = {ENGINE_MARKER: '1'} + if metadata.get(BUNDLE_ID_KEY): + record[BUNDLE_ID_KEY] = str(metadata[BUNDLE_ID_KEY]) + return record + + +def _recorded_lineage(version: "ModelVersion") -> Dict[str, str]: + try: + return registry_lineage(dict(version.show_metrics() or {})) + except Exception as error: + _LOG.warning('Could not read registry metadata: %s', error) + return {} + + +def _alias_holders(model: Any) -> Dict[str, str]: + holders = {} + for row in model.show_versions().to_dict('records'): + for alias in json.loads(row.get('aliases') or '[]'): + holders[alias] = str(row['name']) + return holders + + +def _restore_alias(model: Any, version: str, alias: str) -> None: + try: + model.version(version).set_alias(alias) + except Exception as error: + _LOG.warning('Alias %s is now unset: restoring it on version %s ' + 'failed: %s', alias, version, error) + else: + _LOG.warning('Could not move alias %s; it stays on version %s.', + alias, version) + + class SnowflakeModelRegistryMixin: r"""The model half of :class:`SnowflakeSinkBackend`. @@ -229,6 +300,18 @@ def _session( ) -> ContextManager[SnowflakeSession]: ... + def check_model_location(self, connector: SourceConnector, + location: ModelRegistryLocation) -> None: + r"""Raises :class:`ValueError` unless ``location.model_name`` is a + Snowflake identifier.""" + del connector + if not _MODEL_NAME.fullmatch(location.model_name): + raise ValueError( + f'Snowflake cannot register a model named ' + f'{location.model_name!r}: start with a letter or underscore ' + f'and use letters, digits, underscores and $, or double-quote ' + f"the name, e.g. 'USERS_CHURN'.") + @contextmanager def _registry(self, connector: SourceConnector) -> Iterator["Registry"]: r"""A registry statement is a session statement, so the caller keeps @@ -267,8 +350,8 @@ def publish_model( the model, and optionally the version. A version of ``None`` lets the registry assign one -- the usual case, and the reason this reports back what it created. - metadata: Returned with the publication result. Training metadata - remains in the bundle and is not recorded in registry tags. + metadata: Returned with the publication result; its + ``kumo_bundle_id`` is also recorded in the version's metrics. metrics: Returned with the publication result. Training metrics remain in the bundle's metrics.json. project_id: Ignored; a local runtime has no projects. @@ -288,14 +371,16 @@ def publish_model( f"{bundle} does not look like a published model bundle: it " f"has no {_SPEC_FILE}, which an export writes into every one.") + record = _version_record(metadata) with self._registry(connector) as registry: version = registry.log_model( _bundle_model(str(bundle)), model_name=location.model_name, version_name=location.version, comment=location.comment, - metrics={_ENGINE_MARKER: '1'}, + metrics=record, signatures=_carrier_signature(), + **_task(bundle), ) return _published(connector, version, metadata=metadata, metrics=metrics) @@ -308,10 +393,11 @@ def list_models( ) -> List[PublishedModel]: r"""Lists Kumo model versions visible in the connector's registry. - A Kumo export records only ``kumo_engine=1`` in registry metadata. + A Kumo export records ``kumo_engine=1`` in registry metadata. That marker makes discovery both cheap and precise: versions logged by another tool are omitted, and no bundle is downloaded merely to decide whether it belongs in this list. + ``metadata`` carries the recorded ``kumo_bundle_id``. """ del project_id found: List[PublishedModel] = [] @@ -319,9 +405,11 @@ def list_models( for model in registry.models(): for version in model.versions(): recorded = dict(version.show_metrics() or {}) - if recorded.get(_ENGINE_MARKER) != '1': + if recorded.get(ENGINE_MARKER) != '1': continue - found.append(_published(connector, version)) + found.append( + _published(connector, version, + metadata=registry_lineage(recorded))) return sorted(found, key=lambda item: (item.name, item.version)) def fetch_model( @@ -345,14 +433,15 @@ def fetch_model( though the caller always hands this a fresh, empty ``mkdtemp``. path: Refused, as in :meth:`publish_model`. - location: The model, and optionally the version. ``None`` takes the - registry's default version, which is what a caller who names - none means. + location: The model, and optionally the version or an alias. With + neither, takes the registry's default version, which is what a + caller who names none means. project_id: Ignored. Returns: The name, the version actually read -- resolved, when the caller - named none -- where it lives, and metrics read from metrics.json. + named none -- where it lives, metrics read from metrics.json, and + the version's recorded ``kumo_bundle_id`` as ``metadata``. """ del project_id _refuse_path(path, 'fetching') @@ -366,8 +455,18 @@ def fetch_model( # Resolved before anything on disk is touched: a name the registry # does not know must not cost the caller the bundle already there. model = registry.get_model(location.model_name) - version = (model.default if location.version is None else - model.version(location.version)) + alias = getattr(location, 'alias', None) + if alias is not None: + if resolved(alias) not in _alias_holders(model): + raise ValueError( + f'No version of {location.model_name} holds alias ' + f'{alias}.') + version = model.version(alias) + elif location.version is None: + version = model.default + else: + version = model.version(location.version) + recorded = _recorded_lineage(version) destination.parent.mkdir(parents=True, exist_ok=True) staged = destination.parent / f'{destination.name}{_EXPORT_DIR}' @@ -385,9 +484,52 @@ def fetch_model( isinstance(key, str) and isinstance(value, (int, float)) for key, value in metrics.items())): raise ValueError('Invalid training metrics in model bundle.') - published = _published(connector, version, metrics=metrics) + published = _published(connector, version, metadata=recorded, + metrics=metrics) if destination.exists(): shutil.rmtree(destination) shutil.move(str(bundle), str(destination)) shutil.rmtree(staged, ignore_errors=True) return published + + def set_model_alias( + self, + connector: SourceConnector, + location: ModelRegistryLocation, + alias: str, + ) -> None: + r"""Points :obj:`alias` at ``location.version``, moving it if held. + + Snowflake allows one alias per version. A failed move puts the alias + back on the version that held it. + """ + if location.version is None: + raise ValueError('Setting an alias needs a ModelRegistryLocation ' + 'naming an exact version.') + check_model_alias(alias) + wanted = resolved(alias) + with self._registry(connector) as registry: + model = registry.get_model(location.model_name) + target = model.version(location.version) + name = resolved(str(target.version_name)) + holders = { + held: version + for held, version in _alias_holders(model).items() + if held not in _SYSTEM_ALIASES + } + if holders.get(wanted) == name: + return + held = sorted(a for a, v in holders.items() if v == name) + if held: + raise ValueError( + f'Version {name} of {location.model_name} already holds ' + f'alias {held[0]}; unset it first.') + previous = holders.get(wanted) + if previous is not None: + model.version(previous).unset_alias(alias) + try: + target.set_alias(alias) + except Exception: + if previous is not None: + _restore_alias(model, previous, alias) + raise diff --git a/packages/kumo-snowflake/tests/test_registry.py b/packages/kumo-snowflake/tests/test_registry.py index 9d07e6c28..7c6e80dbc 100644 --- a/packages/kumo-snowflake/tests/test_registry.py +++ b/packages/kumo-snowflake/tests/test_registry.py @@ -8,6 +8,7 @@ from pathlib import Path from types import SimpleNamespace +import pandas as pd import pytest from kumo_snowflake import registry as registry_module from kumo_snowflake.registry import SnowflakeModelRegistryMixin @@ -30,7 +31,11 @@ def export(directory): version = SimpleNamespace( model_name='CHURN', version_name='V1', fully_qualified_model_name='DB.SCHEMA.CHURN', export=export, - show_metrics=lambda: pytest.fail('read registry metrics')) + show_metrics=lambda: { + 'kumo_engine': '1', + 'kumo_bundle_id': 'b1', + 'auc': 0.5 + }) model = SimpleNamespace(default=version, version=lambda _: version) @contextmanager @@ -60,6 +65,7 @@ def test_fetch_returns_bundle_metrics(fetch, requested_version, metrics): payload = json.dumps(metrics) if metrics is not None else None published = run(payload, requested_version) assert published.metrics == (metrics or {}) + assert published.metadata == {'kumo_bundle_id': 'b1'} assert published.version == 'V1' assert (destination / 'weights.bin').read_bytes() == b'model weights' assert not (destination / 'previous').exists() @@ -75,9 +81,13 @@ def test_invalid_metrics_preserve_previous_destination(fetch, payload): assert (destination / 'previous').read_bytes() == b'previous bundle' -@pytest.mark.parametrize('metadata', [{}, {'predictive_query': 'x' * 300}]) -def test_publish_records_only_engine_marker_in_catalog(tmp_path, monkeypatch, - metadata): +@pytest.mark.parametrize('metadata', [{}, { + 'predictive_query': 'x' * 300, + 'kumo_bundle_id': 'b1', + 'kumo_engine': '0' +}]) +def test_publish_records_only_pointers_in_catalog( + tmp_path, monkeypatch, metadata): bundle = tmp_path / 'bundle' bundle.mkdir() (bundle / 'training-spec.json').write_text('{}') @@ -99,16 +109,20 @@ def registry(_): monkeypatch.setattr(registry_module, '_carrier_signature', lambda: {}) published = sink.publish_model(SimpleNamespace(), str(bundle), location=ModelRegistryLocation('CHURN'), - metadata=metadata, metrics={'auc': 0.92}) - assert calls[0]['metrics'] == {'kumo_engine': '1'} + metadata=metadata, + metrics={'auc': 0.92, 'loss': float('nan')}) + assert calls[0]['metrics'] == { + **({'kumo_bundle_id': 'b1'} if metadata else {}), 'kumo_engine': '1' + } assert published.metadata == metadata - assert published.metrics == {'auc': 0.92} + assert published.metrics['auc'] == 0.92 -def test_listing_filters_engine_marker_and_returns_empty_metrics(monkeypatch): +def test_listing_filters_engine_marker_and_returns_lineage(monkeypatch): versions = [] for name, recorded in [('V1', { 'kumo_engine': '1', + 'kumo_bundle_id': 'b1', 'auc': 0.92 }), ('V2', { 'other': '1', @@ -134,4 +148,163 @@ def registry(_): found = sink.list_models(SimpleNamespace()) assert [item.version for item in found] == ['V1', 'V4'] assert [item.metrics for item in found] == [{}, {}] - assert all(item.metadata == {} for item in found) + assert [item.metadata for item in found] == [{'kumo_bundle_id': 'b1'}, {}] + + +class _Model: + def __init__(self, aliases): + self.aliases = aliases + + def show_versions(self): + return pd.DataFrame([ + dict(name=name, aliases=json.dumps(held)) + for name, held in self.aliases.items() + ]) + + def version(self, version_or_alias): + name = next((name for name, held in self.aliases.items() + if version_or_alias.upper() in held), version_or_alias) + held = self.aliases[name] + + def export(directory): + bundle = Path(directory) / 'bundle' + bundle.mkdir() + (bundle / 'training-spec.json').write_text(name) + (bundle / 'publication-manifest.json').write_text('{}') + + return SimpleNamespace( + version_name=name, model_name='CHURN', + fully_qualified_model_name='DB.SCHEMA.CHURN', export=export, + set_alias=lambda alias: held.append(alias.upper()), + unset_alias=lambda alias: held.remove(alias.upper()), + show_metrics=lambda: { + 'kumo_engine': '1', + 'kumo_bundle_id': f'id-{name}' + }) + + +def _aliasing_sink(monkeypatch, model): + @contextmanager + def registry(_): + yield SimpleNamespace(get_model=lambda _: model) + + sink = SnowflakeModelRegistryMixin() + monkeypatch.setattr(sink, '_registry', registry) + return sink + + +def test_set_alias_moves_it_and_refuses_what_snowflake_cannot(monkeypatch): + from kumo_snowflake.sink import SnowflakeSinkBackend + from kumoapi.plugin import ModelAliasingSink + + assert isinstance(SnowflakeSinkBackend(), ModelAliasingSink) + model = _Model({'V1': ['FIRST'], 'V2': ['LAST', 'DEFAULT'], 'V3': []}) + sink = _aliasing_sink(monkeypatch, model) + + def set_alias(version, alias): + sink.set_model_alias(SimpleNamespace(), + ModelRegistryLocation('CHURN', version=version), + alias) + + set_alias('V2', 'prod') + set_alias('V3', 'prod') + set_alias('V3', 'prod') + assert model.aliases == { + 'V1': ['FIRST'], + 'V2': ['LAST', 'DEFAULT'], + 'V3': ['PROD'] + } + for version, alias, match in [('V3', 'staging', 'already holds'), + ('V1', 'default', 'reserved'), + (None, 'prod', 'exact version')]: + with pytest.raises(ValueError, match=match): + set_alias(version, alias) + assert model.aliases['V3'] == ['PROD'] + + +def test_fetch_by_alias_reads_the_aliased_version(monkeypatch, tmp_path): + sink = _aliasing_sink(monkeypatch, _Model({ + 'V1': ['DEFAULT'], + 'V2': ['PROD'] + })) + published = sink.fetch_model( + SimpleNamespace(), str(tmp_path / 'out'), + location=ModelRegistryLocation('CHURN', alias='prod')) + assert published.version == 'V2' + assert published.metadata == {'kumo_bundle_id': 'id-V2'} + assert (tmp_path / 'out' / 'training-spec.json').read_text() == 'V2' + with pytest.raises(ValueError, match='holds alias'): + sink.fetch_model(SimpleNamespace(), str(tmp_path / 'other'), + location=ModelRegistryLocation('CHURN', alias='V1')) + + +def test_failed_alias_move_puts_it_back(monkeypatch): + model = _Model({'V1': ['PROD'], 'V2': []}) + + def fail(alias): + raise RuntimeError('lost connection') + + def failing_version(name): + version = _Model.version(model, name) + if name == 'V2': + version.set_alias = fail + return version + + model.version = failing_version + sink = _aliasing_sink(monkeypatch, model) + with pytest.raises(RuntimeError, match='lost connection'): + sink.set_model_alias(SimpleNamespace(), + ModelRegistryLocation('CHURN', version='V2'), + 'prod') + assert model.aliases == {'V1': ['PROD'], 'V2': []} + + +@pytest.mark.parametrize('task, expected', [ + ('binary_classification', 'TABULAR_BINARY_CLASSIFICATION'), + ('link_prediction', None), +]) +def test_publish_records_the_task_type_monitors_need(tmp_path, monkeypatch, + task, expected): + import sys + + bundle = tmp_path / 'bundle' + bundle.mkdir() + (bundle / 'training-spec.json').write_text('{}') + (bundle / 'prediction-schema.json').write_text( + json.dumps({'task_type': task})) + tasks = SimpleNamespace( + TABULAR_BINARY_CLASSIFICATION='TABULAR_BINARY_CLASSIFICATION') + monkeypatch.setitem(sys.modules, 'snowflake.ml.model', + SimpleNamespace(type_hints=SimpleNamespace(Task=tasks))) + calls = [] + + @contextmanager + def registry(_): + yield SimpleNamespace(log_model=lambda *a, **kw: calls.append(kw) or + SimpleNamespace( + model_name='CHURN', version_name='V1', + fully_qualified_model_name='DB.S.CHURN')) + + sink = SnowflakeModelRegistryMixin() + monkeypatch.setattr(sink, '_registry', registry) + monkeypatch.setattr(registry_module, '_bundle_model', lambda _: object()) + monkeypatch.setattr(registry_module, '_carrier_signature', lambda: {}) + sink.publish_model(SimpleNamespace(), str(bundle), + location=ModelRegistryLocation('CHURN')) + assert calls[0].get('task') == expected + + +@pytest.mark.parametrize('name', ['USERS_CHURN', 'users_churn$2', '_m', + 'ML.USERS_CHURN', 'DB.ML.USERS_CHURN', + '"users-churn"', '"a ""b"""']) +def test_snowflake_identifiers_are_accepted_as_model_names(name): + SnowflakeModelRegistryMixin().check_model_location( + SimpleNamespace(), ModelRegistryLocation(name)) + + +@pytest.mark.parametrize('name', ['users-churn', '1churn', 'users churn', + 'A.B.C.D', 'ML.', '""']) +def test_other_model_names_are_refused_with_the_rule(name): + with pytest.raises(ValueError, match="letter or underscore.*USERS_CHURN"): + SnowflakeModelRegistryMixin().check_model_location( + SimpleNamespace(), ModelRegistryLocation(name)) diff --git a/setup.py b/setup.py index a8446e82b..8b1f65e0e 100644 --- a/setup.py +++ b/setup.py @@ -293,6 +293,9 @@ def parse_project_version(project_file_path: Path) -> str: "respx==0.22.0", "pytorch-frame==0.3.0", "pytorch-lightning==2.6.6", + # A superset of kumo-engine's `mlflow` extra (mlflow-skinny): the tests use + # MLflow's local SQLite store. Its wheel is tested with --no-deps. + "mlflow>=3.13,<4", ] dev_requires = test_requires + [ diff --git a/uv.lock b/uv.lock index fae1d0a85..79523d2dd 100644 --- a/uv.lock +++ b/uv.lock @@ -140,7 +140,7 @@ dependencies = [ { name = "frozenlist" }, { name = "multidict" }, { name = "propcache" }, - { name = "typing-extensions", marker = "python_full_version < '3.13' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, { name = "yarl" }, ] sdist = { url = "https://files.pythonhosted.org/packages/6c/4c/bdccd81e9ee225b69c60e7766c9a5b05364f118f4d383713b89a682d772d/aiohttp-3.14.5.tar.gz", hash = "sha256:5558a7f5a05af9ecf744af91e5baefc436f93c9333e656c27ec253f9a6bbe178", size = 8078052, upload-time = "2026-10-11T01:05:12.408Z" } @@ -304,7 +304,7 @@ version = "1.4.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "frozenlist" }, - { name = "typing-extensions", marker = "python_full_version < '3.13' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/61/62/06741b579156360248d1ec624842ad0edf697050bbaf7c3e46394e106ad1/aiosignal-1.4.0.tar.gz", hash = "sha256:f47eecd9468083c2029cc99945502cb7708b082c232f9aca65da147157b251c7", size = 25007, upload-time = "2025-07-03T22:54:43.528Z" } wheels = [ @@ -376,7 +376,7 @@ version = "4.14.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "idna" }, - { name = "typing-extensions", marker = "python_full_version < '3.13' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/61/cc/a381afa6efea9f496eff839d4a6a1aed3bfafc7b3ab4b0d1b243a12573dd/anyio-4.14.2.tar.gz", hash = "sha256:cfa139f3ed1a23ee8f88a145ddb5ac7605b8bbfd8592baacd7ce3d8bb4313c7f", size = 260176, upload-time = "2026-07-12T20:29:07.082Z" } wheels = [ @@ -572,12 +572,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/88/c6/92fcd42f1ba33e1184263f25bfabf3d27c383410470f169e4b8163bf9c17/beautifulsoup4-4.15.0-py3-none-any.whl", hash = "sha256:d6f88de62e1d4e38ecb1077eb9724cd0eff29d2a08ca16a401e9b9e93f117cf9", size = 109924, upload-time = "2026-06-07T16:44:21.566Z" }, ] +[[package]] +name = "blinker" +version = "1.9.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/21/28/9b3f50ce0e048515135495f198351908d99540d69bfdc8c1d15b73dc55ce/blinker-1.9.0.tar.gz", hash = "sha256:b4ce2265a7abece45e7cc896e98dbebe6cead56bcf805a3d23136d145f5445bf", size = 22460, upload-time = "2024-11-08T17:25:47.436Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/10/cb/f2ad4230dc2eb1a74edf38f1a38b9b52277f75bef262d8908e60d957e13c/blinker-1.9.0-py3-none-any.whl", hash = "sha256:ba0efaa9080b619ff2f3459d1d500c57bddea4a6b424b60a91141db6fd2f08bc", size = 8458, upload-time = "2024-11-08T17:25:46.184Z" }, +] + [[package]] name = "blockbuster" version = "1.5.28" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "forbiddenfruit", marker = "implementation_name == 'cpython' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "forbiddenfruit", marker = "implementation_name == 'cpython'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/28/14/de0062a8e3758fe1880729c1828196e6e951b8cd3299dfdd1e68aeae6e63/blockbuster-1.5.28.tar.gz", hash = "sha256:5af2ed9de6c874339a76cc831d1d11db79800a974573580e52ad827a70803482", size = 98235, upload-time = "2026-09-18T10:51:11.263Z" } wheels = [ @@ -605,7 +614,7 @@ source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "botocore-stubs" }, { name = "types-s3transfer" }, - { name = "typing-extensions", marker = "python_full_version < '3.12' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "typing-extensions", marker = "python_full_version < '3.12'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/7b/94/9dd7426ce339c15ff6bae15b609977aad2f8f778244f9a06e99db821c974/boto3_stubs-1.42.90.tar.gz", hash = "sha256:d0ea4bc1352ba5ef8fae13cfb2fef511f372113ede0388b50cdd19811d39ce01", size = 102688, upload-time = "2026-04-16T21:25:36.455Z" } wheels = [ @@ -657,8 +666,8 @@ dependencies = [ { name = "numpy" }, { name = "pandas" }, { name = "scikit-learn" }, - { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "scipy", version = "1.18.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, + { name = "scipy", version = "1.18.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/6d/28/3325463e27238bacefce34902ce2e5883ea3bd4044ce93e744ef499a6d72/category_encoders-2.11.1.tar.gz", hash = "sha256:f21b99185800c19b562b2928a0fe67dbf5658d048cd70451bfe8811c4123fdb0", size = 76600, upload-time = "2026-09-08T15:29:09.823Z" } wheels = [ @@ -679,7 +688,7 @@ name = "cffi" version = "2.1.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "pycparser", marker = "implementation_name != 'PyPy' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "pycparser", marker = "implementation_name != 'PyPy'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/9e/ef/008a1939e372c06329a3fce4279c02f328488f3526744906eeec3da7ad5f/cffi-2.1.1.tar.gz", hash = "sha256:dd31f52ea1086513bb9df30f8fcee9b8918323ae067a3d5b78bc826a000712be", size = 530807, upload-time = "2026-08-03T21:21:18.939Z" } wheels = [ @@ -960,7 +969,7 @@ name = "colorlog" version = "6.12.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "colorama", marker = "sys_platform == 'win32' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "colorama", marker = "sys_platform == 'win32'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/8c/55/ba79756cb90c8d69d599d57785398ac87bba7b19c80e87f4e8a562197c93/colorlog-6.12.0.tar.gz", hash = "sha256:2a7924c1dadf18b22a0eb8b06d1c7b01d5341707ec1641eb6fcc4fde0c3e8e5f", size = 18151, upload-time = "2026-07-23T13:40:40.71Z" } wheels = [ @@ -981,7 +990,7 @@ resolution-markers = [ "python_full_version < '3.12' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu130' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests'", ] dependencies = [ - { name = "numpy", marker = "python_full_version < '3.12' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "numpy" }, ] sdist = { url = "https://files.pythonhosted.org/packages/58/01/1253e6698a07380cd31a736d248a3f2a50a7c88779a1813da27503cadc2a/contourpy-1.3.3.tar.gz", hash = "sha256:083e12155b210502d0bca491432bb04d56dc3432f95a979b429f2848c3dbe880", size = 13466174, upload-time = "2025-07-26T12:03:12.549Z" } wheels = [ @@ -1086,7 +1095,7 @@ resolution-markers = [ "python_full_version == '3.12.*' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu130' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests'", ] dependencies = [ - { name = "numpy", marker = "python_full_version >= '3.12' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "numpy" }, ] sdist = { url = "https://files.pythonhosted.org/packages/83/5a/a55177dd22553a277388e8a1b3220e92de91bacb28356cdc73caa240121d/contourpy-1.4.0.tar.gz", hash = "sha256:20156f5a1ac4f8ce02656e39a61e82164a3d359796dc8026f75b062783d500e1", size = 13323726, upload-time = "2026-09-11T19:05:05.808Z" } wheels = [ @@ -1277,7 +1286,7 @@ wheels = [ [package.optional-dependencies] toml = [ - { name = "tomli", marker = "python_full_version <= '3.11' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "tomli", marker = "python_full_version <= '3.11'" }, ] [[package]] @@ -1285,7 +1294,7 @@ name = "cryptography" version = "50.0.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "cffi", marker = "platform_python_implementation != 'PyPy' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "cffi", marker = "platform_python_implementation != 'PyPy'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/de/41/6cbdcf9142d00fe82836fbb51e503e58088575cf7a0fe1dbff6695bf0840/cryptography-50.0.0.tar.gz", hash = "sha256:eeac2acb5a20ed25e0ad6d1df9891a520b78b404266b6d11778f25d5d691a6c9", size = 880201, upload-time = "2026-07-31T14:25:10.11Z" } wheels = [ @@ -1341,7 +1350,7 @@ name = "cuda-bindings" version = "13.4.2" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "cuda-pathfinder", marker = "extra == 'extra-4-kumo-cu130' or extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra != 'extra-4-kumo-cpu' and extra != 'group-4-kumo-public-tests')" }, + { name = "cuda-pathfinder", marker = "extra == 'extra-4-kumo-cu130' or extra == 'extra-4-kumo-cu132' or (extra != 'extra-4-kumo-cpu' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/19/6f/e00ffcebcad6405a2326a52612eff98240a67c637dfed516e61b881e2db3/cuda_bindings-13.4.2-cp311-cp311-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:fc0a18f65b26459cd3c49c0250a44b419cdf790cffcb096ff4d46a1973f6f4b7", size = 6484069, upload-time = "2026-09-17T22:05:52.594Z" }, @@ -1390,34 +1399,34 @@ wheels = [ [package.optional-dependencies] cudart = [ - { name = "nvidia-cuda-runtime", version = "13.0.96", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cuda-runtime", version = "13.0.96", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] cufft = [ - { name = "nvidia-cufft", version = "12.0.0.61", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cufft", version = "12.0.0.61", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] cufile = [ - { name = "nvidia-cufile", version = "1.15.1.6", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cufile", version = "1.15.1.6", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] cupti = [ - { name = "nvidia-cuda-cupti", version = "13.0.85", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cuda-cupti", version = "13.0.85", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] curand = [ - { name = "nvidia-curand", version = "10.4.0.35", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-curand", version = "10.4.0.35", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] cusolver = [ - { name = "nvidia-cusolver", version = "12.0.4.66", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cusolver", version = "12.0.4.66", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] cusparse = [ - { name = "nvidia-cusparse", version = "12.6.3.3", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cusparse", version = "12.6.3.3", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] nvjitlink = [ - { name = "nvidia-nvjitlink", version = "13.0.88", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-nvjitlink", version = "13.0.88", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] nvrtc = [ - { name = "nvidia-cuda-nvrtc", version = "13.0.88", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cuda-nvrtc", version = "13.0.88", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] nvtx = [ - { name = "nvidia-nvtx", version = "13.0.85", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-nvtx", version = "13.0.85", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] [[package]] @@ -1436,43 +1445,43 @@ wheels = [ [package.optional-dependencies] cublas = [ - { name = "nvidia-cublas", version = "13.4.0.1", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "nvidia-cuda-nvrtc", version = "13.2.78", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cublas", version = "13.4.0.1", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cuda-nvrtc", version = "13.2.78", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] cudart = [ - { name = "nvidia-cuda-runtime", version = "13.2.75", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cuda-runtime", version = "13.2.75", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] cufft = [ - { name = "nvidia-cufft", version = "12.2.0.46", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "nvidia-nvjitlink", version = "13.4.92", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cufft", version = "12.2.0.46", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-nvjitlink", version = "13.4.92", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] cufile = [ - { name = "nvidia-cufile", version = "1.17.1.22", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'aarch64' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (platform_machine == 'aarch64' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'aarch64' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (platform_machine == 'aarch64' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'aarch64' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (platform_machine == 'x86_64' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (platform_machine == 'x86_64' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (platform_machine == 'x86_64' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cufile", version = "1.17.1.22", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] cupti = [ - { name = "nvidia-cuda-cupti", version = "13.2.75", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cuda-cupti", version = "13.2.75", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] curand = [ - { name = "nvidia-curand", version = "10.4.2.55", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-curand", version = "10.4.2.55", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] cusolver = [ - { name = "nvidia-cublas", version = "13.4.0.1", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "nvidia-cusolver", version = "12.2.0.1", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "nvidia-cusparse", version = "12.7.10.1", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "nvidia-nvjitlink", version = "13.4.92", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cublas", version = "13.4.0.1", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cusolver", version = "12.2.0.1", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cusparse", version = "12.7.10.1", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-nvjitlink", version = "13.4.92", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] cusparse = [ - { name = "nvidia-cusparse", version = "12.7.10.1", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "nvidia-nvjitlink", version = "13.4.92", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cusparse", version = "12.7.10.1", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-nvjitlink", version = "13.4.92", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] nvjitlink = [ - { name = "nvidia-nvjitlink", version = "13.4.92", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-nvjitlink", version = "13.4.92", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] nvrtc = [ - { name = "nvidia-cuda-nvrtc", version = "13.2.78", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cuda-nvrtc", version = "13.2.78", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] nvtx = [ - { name = "nvidia-nvtx", version = "13.2.75", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-nvtx", version = "13.2.75", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] [[package]] @@ -1594,6 +1603,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/02/08/9c41fb51ab5b43eb21674aff13df270e8ba6c4b29c8624e328dc7a9482af/distlib-0.4.3-py2.py3-none-any.whl", hash = "sha256:4b0ce306c966eb73bc3a7b6abad017c556dadd92c44701562cd528ac7fde4d5b", size = 470628, upload-time = "2026-06-12T08:04:50.506Z" }, ] +[[package]] +name = "docker" +version = "7.2.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pywin32", marker = "sys_platform == 'win32'" }, + { name = "requests" }, + { name = "urllib3" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/88/7f/731ff914b0255d3d065f45fd4e626d4b8c95dbcbaada049f337a6ac16410/docker-7.2.0.tar.gz", hash = "sha256:cebb93773d334f778e023a7ee352a8d6e13ab1bd3b863a4d4a59dec897df43ac", size = 118731, upload-time = "2026-07-09T14:53:46.39Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/75/23/529140fe1aab80fc6992f93a706deec709140a6397439139a054e1515c45/docker-7.2.0-py3-none-any.whl", hash = "sha256:a3f45fdeb9165e2d25d9a1d02ddf3bc70fb572cf5ebbf9b58558c22caf29b71f", size = 148775, upload-time = "2026-07-09T14:53:45.224Z" }, +] + [[package]] name = "docutils" version = "0.21.2" @@ -1717,6 +1740,36 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/35/f8/08d37b2cd89da306e3520bd27f8a85692122b42b56c0c2c3784ff09c022f/flake8-7.1.2-py2.py3-none-any.whl", hash = "sha256:1cbc62e65536f65e6d754dfe6f1bada7f5cf392d6f5db3c2b85892466c3e7c1a", size = 57745, upload-time = "2025-02-16T18:45:42.351Z" }, ] +[[package]] +name = "flask" +version = "3.1.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "blinker" }, + { name = "click" }, + { name = "itsdangerous" }, + { name = "jinja2" }, + { name = "markupsafe" }, + { name = "werkzeug" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/26/00/35d85dcce6c57fdc871f3867d465d780f302a175ea360f62533f12b27e2b/flask-3.1.3.tar.gz", hash = "sha256:0ef0e52b8a9cd932855379197dd8f94047b359ca0a78695144304cb45f87c9eb", size = 759004, upload-time = "2026-02-19T05:00:57.678Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7f/9c/34f6962f9b9e9c71f6e5ed806e0d0ff03c9d1b0b2340088a0cf4bce09b18/flask-3.1.3-py3-none-any.whl", hash = "sha256:f4bcbefc124291925f1a26446da31a5178f9483862233b23c0c96a20701f670c", size = 103424, upload-time = "2026-02-19T05:00:56.027Z" }, +] + +[[package]] +name = "flask-cors" +version = "6.0.5" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "flask" }, + { name = "werkzeug" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/47/03/4e464a50860f9adf08b5c1d3479cb8ea1f12af2aa69535c7042c6e628135/flask_cors-6.0.5.tar.gz", hash = "sha256:30c5031552cd59f620ac0c8211dac45b345d3b2df310e7721879e4f46ef9c601", size = 101386, upload-time = "2026-06-08T20:20:17.765Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/49/55/5bb1a2d918e9f02f131e47a59032bae70e48050e986e941511fd737a935c/flask_cors-6.0.5-py3-none-any.whl", hash = "sha256:68fcf75693e961f3af26683b23c4b9a8fb6b64de17d20d0c37b95e8de7ab2ed8", size = 16692, upload-time = "2026-06-08T20:20:16.247Z" }, +] + [[package]] name = "fonttools" version = "4.65.0" @@ -1926,6 +1979,30 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/69/43/a0096d96d3271640bbda056b9d22249be141a54509859b99002ee6149a01/gcsfs-2026.6.0-py3-none-any.whl", hash = "sha256:4a81fe40fbdef450aeaacb534b42c04059f8f3b080e3e169137247a5bb023464", size = 89629, upload-time = "2026-06-18T03:16:07.407Z" }, ] +[[package]] +name = "gitdb" +version = "4.0.12" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "smmap" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/72/94/63b0fc47eb32792c7ba1fe1b694daec9a63620db1e313033d18140c2320a/gitdb-4.0.12.tar.gz", hash = "sha256:5ef71f855d191a3326fcfbc0d5da835f26b13fbcba60c32c21091c349ffdb571", size = 394684, upload-time = "2025-01-02T07:20:46.413Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a0/61/5c78b91c3143ed5c14207f463aecfc8f9dbb5092fb2869baf37c273b2705/gitdb-4.0.12-py3-none-any.whl", hash = "sha256:67073e15955400952c6565cc3e707c554a4eea2e428946f7a4c162fab9bd9bcf", size = 62794, upload-time = "2025-01-02T07:20:43.624Z" }, +] + +[[package]] +name = "gitpython" +version = "3.2.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "gitdb" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/6e/2d/6f6e649818da44d4499604802c89329b8d9799687a124e3a5e467a643336/gitpython-3.2.0.tar.gz", hash = "sha256:fb92310af6844d96adc95ca066ed2e617c00e1dbd146a326626c81e72e18cc2e", size = 238091, upload-time = "2026-09-30T09:10:32.65Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fa/48/f36d233ce749d1bb880cb5e6f8bb171dc2effabc7435c31da604d831b3ec/gitpython-3.2.0-py3-none-any.whl", hash = "sha256:bd70c5ec05cd2b797423e7eb312147d2458d3cca92085888fba2213f85905537", size = 228537, upload-time = "2026-09-30T09:10:31.159Z" }, +] + [[package]] name = "google-api-core" version = "2.30.3" @@ -2151,6 +2228,42 @@ grpc = [ { name = "grpcio" }, ] +[[package]] +name = "graphene" +version = "3.4.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "graphql-core" }, + { name = "graphql-relay" }, + { name = "python-dateutil" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/cc/f6/bf62ff950c317ed03e77f3f6ddd7e34aaa98fe89d79ebd660c55343d8054/graphene-3.4.3.tar.gz", hash = "sha256:2a3786948ce75fe7e078443d37f609cbe5bb36ad8d6b828740ad3b95ed1a0aaa", size = 44739, upload-time = "2024-11-09T20:44:25.757Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/66/e0/61d8e98007182e6b2aca7cf65904721fb2e4bce0192272ab9cb6f69d8812/graphene-3.4.3-py2.py3-none-any.whl", hash = "sha256:820db6289754c181007a150db1f7fff544b94142b556d12e3ebc777a7bf36c71", size = 114894, upload-time = "2024-11-09T20:44:23.851Z" }, +] + +[[package]] +name = "graphql-core" +version = "3.2.13" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/4e/5e/aa0d4e701b50db0bab71b125dd19ddb0f98c638d008d6f7e3d8ce9cbc92e/graphql_core-3.2.13.tar.gz", hash = "sha256:bb81dd266d4ab7b591bd976f1b23639d97776cb9ac1a896b4a93c271e11ed618", size = 577471, upload-time = "2026-09-27T12:22:57.136Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/1f/b4/b52a324b0b297ddc49b8ea31a51d39fae63132f694a267b9381d94e2e0af/graphql_core-3.2.13-py3-none-any.whl", hash = "sha256:b0eb04f2c31556b2310a77c8fb53c74e8b56570f8fea594d89c6ef7827dbb497", size = 268191, upload-time = "2026-09-27T12:22:55.661Z" }, +] + +[[package]] +name = "graphql-relay" +version = "3.2.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "graphql-core" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/d1/13/98fbf8d67552f102488ffc16c6f559ce71ea15f6294728d33928ab5ff14d/graphql-relay-3.2.0.tar.gz", hash = "sha256:1ff1c51298356e481a0be009ccdff249832ce53f30559c1338f22a0e0d17250c", size = 50027, upload-time = "2022-04-16T11:03:45.447Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/74/16/a4cf06adbc711bd364a73ce043b0b08d8fa5aae3df11b6ee4248bcdad2e0/graphql_relay-3.2.0-py3-none-any.whl", hash = "sha256:c9b22bd28b170ba1fe674c74384a8ff30a76c8e26f88ac3aa1584dd3179953e5", size = 16940, upload-time = "2022-04-16T11:03:43.895Z" }, +] + [[package]] name = "graphviz" version = "0.20.1" @@ -2326,6 +2439,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/71/c4/3a77e4273e866b1b0c412afd80882e95941a37170032b5109d847c501124/grpcio_status-1.84.0-py3-none-any.whl", hash = "sha256:0c182ca0d6e60acbfd0e14499cf39a155e4827a1c3fd9f7638e49af15a74c30a", size = 14639, upload-time = "2026-09-14T07:10:15.175Z" }, ] +[[package]] +name = "gunicorn" +version = "26.2.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d9/8a/e4ef6ee11701b6cd64702848415ffb69eeff85cb388a3c6c7fe86f22f3f8/gunicorn-26.2.0.tar.gz", hash = "sha256:62b864895d9ebff0b2f9867ba04fe811c93121596540830c9c916d0769668447", size = 787921, upload-time = "2026-08-24T15:05:59.3Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fe/85/7522a52e5e2f42faf1a129113ab63e548c42e103e9af395b7bfe65e403e2/gunicorn-26.2.0-py3-none-any.whl", hash = "sha256:bd249d0b3f7972f7432f0a6b6ff3b3ee2d129f70cd1ff6c09a9dd9e29a2b88e3", size = 228389, upload-time = "2026-08-24T15:05:57.67Z" }, +] + [[package]] name = "h11" version = "0.16.0" @@ -2376,6 +2498,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/56/95/9377bcb415797e44274b51d46e3249eba641711cf3348050f76ee7b15ffc/httpx-0.27.2-py3-none-any.whl", hash = "sha256:7bb2708e112d8fdd7829cd4243970f0c223274051cb35ee80c03301ee29a3df0", size = 76395, upload-time = "2024-08-27T12:53:59.653Z" }, ] +[[package]] +name = "huey" +version = "3.4.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/57/25/1281d8f3febc30b5facbfc5e4163447241711591c094dac642d32c827eb0/huey-3.4.0.tar.gz", hash = "sha256:ac9b02e741f13da4092ea19bcaafc73669e97d3153914d2fed2868df144f6dfa", size = 630013, upload-time = "2026-09-04T02:04:34.521Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/57/31/cee585bd4d1e39947a426a6736f4e17d22b454952370fb9cf48eb3103085/huey-3.4.0-py3-none-any.whl", hash = "sha256:d0580762397744026f83baae5f2b97a62474a62b8565ac75bbafefe72f9c784f", size = 136493, upload-time = "2026-09-04T02:04:32.848Z" }, +] + [[package]] name = "humanize" version = "4.6.0" @@ -2473,6 +2604,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/0a/63/4036ae70eea279c63e2304b91ee0ac182f467f24f86394ecfe726092340b/isort-5.12.0-py3-none-any.whl", hash = "sha256:f84c2818376e66cf843d497486ea8fed8700b340f308f076c6fb1229dff318b6", size = 91198, upload-time = "2023-01-28T17:10:21.149Z" }, ] +[[package]] +name = "itsdangerous" +version = "2.2.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/9c/cb/8ac0172223afbccb63986cc25049b154ecfb5e85932587206f42317be31d/itsdangerous-2.2.0.tar.gz", hash = "sha256:e0050c0b7da1eea53ffaf149c0cfbb5c6e2e2b69c4bef22c81fa6eb73e5f6173", size = 54410, upload-time = "2024-04-16T21:28:15.614Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/04/96/92447566d16df59b2a776c0fb82dbc4d9e07cd95062562af01e408583fc4/itsdangerous-2.2.0-py3-none-any.whl", hash = "sha256:c6242fc49e35958c8b15141343aa660db5fc54d4f13a1db01a3f5891b98700ef", size = 16234, upload-time = "2024-04-16T21:28:14.499Z" }, +] + [[package]] name = "jinja2" version = "3.1.6" @@ -2751,8 +2891,8 @@ dependencies = [ { name = "requests" }, { name = "s3fs" }, { name = "scikit-learn" }, - { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "scipy", version = "1.18.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, + { name = "scipy", version = "1.18.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, { name = "setuptools" }, { name = "setuptools-scm" }, { name = "smart-open" }, @@ -2764,11 +2904,11 @@ dependencies = [ { name = "tenacity" }, { name = "tensorboard" }, { name = "thrift" }, - { name = "torch", version = "2.12.0", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "torch", version = "2.12.0", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, { name = "torch", version = "2.12.0", source = { registry = "https://pypi.org/simple" }, marker = "(extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu130' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests')" }, - { name = "torch", version = "2.12.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra != 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "torch", version = "2.12.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, { name = "torch", version = "2.12.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "extra == 'extra-4-kumo-cu130' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "torch", version = "2.12.0+cu132", source = { registry = "https://download.pytorch.org/whl/cu132" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "torch", version = "2.12.0+cu132", source = { registry = "https://download.pytorch.org/whl/cu132" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, { name = "types-aiobotocore", extra = ["s3"] }, { name = "typing-extensions" }, { name = "urllib3" }, @@ -2788,7 +2928,7 @@ cu130 = [ { name = "torch", version = "2.12.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } }, ] cu132 = [ - { name = "kumo-core", extra = ["faiss-gpu"], marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "kumo-core", extra = ["faiss-gpu"], marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, { name = "torch", version = "2.12.0+cu132", source = { registry = "https://download.pytorch.org/whl/cu132" } }, ] data-platform = [ @@ -2805,6 +2945,7 @@ dev = [ { name = "flake8" }, { name = "isort" }, { name = "kubernetes-stubs" }, + { name = "mlflow" }, { name = "moto" }, { name = "packaging" }, { name = "pick" }, @@ -2883,6 +3024,7 @@ rest = [ test = [ { name = "category-encoders" }, { name = "coverage" }, + { name = "mlflow" }, { name = "moto" }, { name = "pytest" }, { name = "pytest-asyncio" }, @@ -2906,7 +3048,7 @@ tracing = [ [package.dev-dependencies] public-tests = [ - { name = "torch", version = "2.12.0", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "torch", version = "2.12.0", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, { name = "torch", version = "2.12.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform != 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] @@ -2988,6 +3130,8 @@ requires-dist = [ { name = "markdown", marker = "extra == 'rest'", specifier = "==3.10" }, { name = "matplotlib", marker = "extra == 'drawing'", specifier = ">=3.10.5" }, { name = "matplotlib", marker = "extra == 'full'", specifier = ">=3.10.5" }, + { name = "mlflow", marker = "extra == 'dev'", specifier = ">=3.13,<4" }, + { name = "mlflow", marker = "extra == 'test'", specifier = ">=3.13,<4" }, { name = "moto", marker = "extra == 'dev'", specifier = "==4.2.14" }, { name = "moto", marker = "extra == 'test'", specifier = "==4.2.14" }, { name = "mypy-boto3-s3", specifier = "==1.42.94" }, @@ -3212,11 +3356,11 @@ dependencies = [ { name = "sqlglot" }, { name = "tabulate" }, { name = "tensorboard" }, - { name = "torch", version = "2.12.0", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "torch", version = "2.12.0", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, { name = "torch", version = "2.12.0", source = { registry = "https://pypi.org/simple" }, marker = "(extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu130' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests')" }, - { name = "torch", version = "2.12.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra != 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "torch", version = "2.12.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, { name = "torch", version = "2.12.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "extra == 'extra-4-kumo-cu130' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "torch", version = "2.12.0+cu132", source = { registry = "https://download.pytorch.org/whl/cu132" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "torch", version = "2.12.0+cu132", source = { registry = "https://download.pytorch.org/whl/cu132" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, { name = "torch-geometric" }, { name = "torchmetrics" }, { name = "typing-extensions" }, @@ -3346,9 +3490,19 @@ databricks = [ gcs = [ { name = "kumo-gcs" }, ] +mlflow = [ + { name = "mlflow-skinny" }, + { name = "nvidia-ml-py" }, +] s3 = [ { name = "kumo-s3" }, ] +sagemaker = [ + { name = "kumo-s3" }, + { name = "mlflow-skinny" }, + { name = "nvidia-ml-py" }, + { name = "sagemaker-mlflow" }, +] snowflake = [ { name = "kumo-snowflake", extra = ["sdk"] }, ] @@ -3359,11 +3513,11 @@ test = [ { name = "numpy" }, { name = "pytest" }, { name = "pytest-mock" }, - { name = "torch", version = "2.12.0", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "torch", version = "2.12.0", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, { name = "torch", version = "2.12.0", source = { registry = "https://pypi.org/simple" }, marker = "(extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu130' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests')" }, - { name = "torch", version = "2.12.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra != 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "torch", version = "2.12.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, { name = "torch", version = "2.12.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "extra == 'extra-4-kumo-cu130' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "torch", version = "2.12.0+cu132", source = { registry = "https://download.pytorch.org/whl/cu132" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "torch", version = "2.12.0+cu132", source = { registry = "https://download.pytorch.org/whl/cu132" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, ] text = [ { name = "kumo-glove-6b-50d" }, @@ -3382,20 +3536,24 @@ requires-dist = [ { name = "kumo-bigquery", marker = "extra == 'bigquery'", editable = "packages/kumo-bigquery" }, { name = "kumo-core", editable = "packages/kumo-core" }, { name = "kumo-databricks", marker = "extra == 'databricks'", editable = "packages/kumo-databricks" }, + { name = "kumo-engine", extras = ["mlflow", "s3"], marker = "extra == 'sagemaker'" }, { name = "kumo-gcs", marker = "extra == 'gcs'", editable = "packages/kumo-gcs" }, { name = "kumo-glove-6b-50d", marker = "extra == 'text'", editable = "tools/glove-6b-50d" }, { name = "kumo-s3", marker = "extra == 's3'", editable = "packages/kumo-s3" }, { name = "kumo-snowflake", extras = ["sdk"], marker = "extra == 'snowflake'", editable = "packages/kumo-snowflake" }, + { name = "mlflow-skinny", marker = "extra == 'mlflow'", specifier = ">=3.0.1,<4" }, { name = "numpy", marker = "extra == 'test'", specifier = ">=1.26,<3" }, + { name = "nvidia-ml-py", marker = "extra == 'mlflow'" }, { name = "pandas" }, { name = "pyarrow", specifier = ">=8.0.0" }, { name = "pydantic", specifier = ">=1.10.21" }, { name = "pytest", marker = "extra == 'test'" }, { name = "pytest-mock", marker = "extra == 'test'" }, + { name = "sagemaker-mlflow", marker = "extra == 'sagemaker'", specifier = ">=0.5.0" }, { name = "torch", marker = "extra == 'test'", specifier = ">=2.9,<2.13" }, { name = "typing-extensions", specifier = ">=4.5.0" }, ] -provides-extras = ["cpu", "test", "snowflake", "databricks", "bigquery", "s3", "text", "gcs"] +provides-extras = ["cpu", "test", "snowflake", "databricks", "bigquery", "s3", "text", "gcs", "mlflow", "sagemaker"] [[package]] name = "kumo-gcs" @@ -3541,7 +3699,7 @@ requires-dist = [ { name = "pyyaml", marker = "extra == 'sdk'" }, { name = "requests", specifier = ">=2.31" }, { name = "snowflake-connector-python", specifier = ">=4.3" }, - { name = "snowflake-ml-python", marker = "extra == 'sdk'" }, + { name = "snowflake-ml-python", marker = "extra == 'sdk'", specifier = ">=2.0,<3" }, { name = "snowflake-snowpark-python", marker = "extra == 'sdk'" }, ] provides-extras = ["sdk", "test"] @@ -3776,8 +3934,8 @@ name = "matplotlib" version = "3.11.2" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "contourpy", version = "1.3.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "contourpy", version = "1.4.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "contourpy", version = "1.3.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, + { name = "contourpy", version = "1.4.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, { name = "cycler" }, { name = "fonttools" }, { name = "kiwisolver" }, @@ -3864,6 +4022,89 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ec/9d/1832c79be2a87aece4ecbbddd3ae315b9cc90ae62c6755581afc06670867/mermaid_py-0.9.0-py3-none-any.whl", hash = "sha256:6abc53f9527118c147d59568852fe782dcdbe54d5512f23d0e3a1b5dec9a302b", size = 32952, upload-time = "2026-09-16T10:46:35.6Z" }, ] +[[package]] +name = "mlflow" +version = "3.17.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "aiohttp" }, + { name = "alembic" }, + { name = "cryptography" }, + { name = "docker" }, + { name = "flask" }, + { name = "flask-cors" }, + { name = "graphene" }, + { name = "gunicorn", marker = "sys_platform != 'win32'" }, + { name = "huey" }, + { name = "matplotlib" }, + { name = "mlflow-skinny" }, + { name = "mlflow-tracing" }, + { name = "numpy" }, + { name = "pandas" }, + { name = "pyarrow" }, + { name = "scikit-learn" }, + { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, + { name = "scipy", version = "1.18.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, + { name = "skops" }, + { name = "sqlalchemy" }, + { name = "waitress", marker = "sys_platform == 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/6e/30/953491ffbd83c90aba4a72c8a68c03ea16a8b016050557094d059268bf7d/mlflow-3.17.1.tar.gz", hash = "sha256:b41225845f200968bfeb015c48ad9e57a02f3bf4cee246fb04a75325728c21d8", size = 11158262, upload-time = "2026-10-13T09:00:18.681Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a4/64/ee5ca7500b3f4ca187b7c0c2f080a029eed1df7f22b334391587979b23dc/mlflow-3.17.1-py3-none-any.whl", hash = "sha256:726bf7feac25dbdb5b3098bf20b86574393bb75468eef721902ec025e724752f", size = 12004276, upload-time = "2026-10-13T09:00:14.749Z" }, +] + +[[package]] +name = "mlflow-skinny" +version = "3.17.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "anyio" }, + { name = "cachetools" }, + { name = "click" }, + { name = "cloudpickle" }, + { name = "databricks-sdk" }, + { name = "fastapi" }, + { name = "gitpython" }, + { name = "importlib-metadata" }, + { name = "opentelemetry-api" }, + { name = "opentelemetry-proto" }, + { name = "opentelemetry-sdk" }, + { name = "packaging" }, + { name = "protobuf" }, + { name = "pydantic" }, + { name = "python-dotenv" }, + { name = "pyyaml" }, + { name = "requests" }, + { name = "sqlparse" }, + { name = "starlette" }, + { name = "typing-extensions" }, + { name = "uvicorn" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c7/b3/a72208dd21b7b8345d3d15eafa8d5d2897d36e21420e22d290fad5775aa4/mlflow_skinny-3.17.1.tar.gz", hash = "sha256:c4587083811f638d9c2e6e8ab7caabbccc51ccc09de2c87a503913b32bbfd7bf", size = 3483742, upload-time = "2026-10-13T08:58:49.651Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0f/54/d0e5e2588013d7655f9cfb3d17fff4760c22652ce628b57ae4ac0fa913e6/mlflow_skinny-3.17.1-py3-none-any.whl", hash = "sha256:7374adf971e06f77f861dd8196c50e344557728b643192c665f0603c19a8d146", size = 4105632, upload-time = "2026-10-13T08:58:47.785Z" }, +] + +[[package]] +name = "mlflow-tracing" +version = "3.17.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cachetools" }, + { name = "databricks-sdk" }, + { name = "opentelemetry-api" }, + { name = "opentelemetry-proto" }, + { name = "opentelemetry-sdk" }, + { name = "packaging" }, + { name = "protobuf" }, + { name = "pydantic" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/4c/da/e226d475789bbb8caf68aafa3cbf2b715e5f6eda00f7f659b7eda6fba1ef/mlflow_tracing-3.17.1.tar.gz", hash = "sha256:c154492045b5f3544ab07fd9ab9d7c07febfb53065d43f1a07f33a31c46b3505", size = 1688196, upload-time = "2026-10-13T08:59:23.98Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ce/19/2ba5a327efa9a45ea46c927103a7c57b24441375d8b3435f6b27d187ac91/mlflow_tracing-3.17.1-py3-none-any.whl", hash = "sha256:4f6405e48efdf9330101aa059383103e1e15559664ec792290b9b95556f9e851", size = 1981200, upload-time = "2026-10-13T08:59:22.314Z" }, +] + [[package]] name = "more-itertools" version = "11.1.0" @@ -4091,7 +4332,7 @@ name = "mypy-boto3-s3" version = "1.42.94" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions", marker = "python_full_version < '3.12' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "typing-extensions", marker = "python_full_version < '3.12'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/6d/55/2268b22037d1c6814c2229904d43f9cb3fcd9b9e9d01a3402b58a10270ce/mypy_boto3_s3-1.42.94.tar.gz", hash = "sha256:1d92d722cf00573b8111e98493ab386e0c1b59a1530b7fee4af77f2d9a1c477d", size = 77049, upload-time = "2026-04-22T21:31:00.857Z" } wheels = [ @@ -4289,7 +4530,7 @@ resolution-markers = [ "python_full_version < '3.12'", ] dependencies = [ - { name = "nvidia-cuda-nvrtc", version = "13.2.78", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cuda-nvrtc", version = "13.2.78", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/d0/48/8571e55aaabbd112f06a39a71fbc3abf0b8a0d450921816a13a6d5e8fa0b/nvidia_cublas-13.4.0.1-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:705d7d214fbb20f134415ecadc488abf74f444c155a6555dec5687404afb18a9", size = 513051304, upload-time = "2026-04-13T09:49:32.758Z" }, @@ -4432,7 +4673,7 @@ version = "9.20.0.48" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "nvidia-cublas", version = "13.1.1.3", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-4-kumo-cu130' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests')" }, - { name = "nvidia-cublas", version = "13.4.0.1", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cublas", version = "13.4.0.1", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/56/c5/83384d846b2fd17c44bd499b36c75a45ed4f095fbbb2252294e89cea5c5c/nvidia_cudnn_cu13-9.20.0.48-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:e31454ae00094b0c55319d9d15b6fa2fc50a9e1c0f5c8c80fb75258234e731e1", size = 444574296, upload-time = "2026-03-09T19:28:27.751Z" }, @@ -4470,7 +4711,7 @@ resolution-markers = [ "python_full_version < '3.12'", ] dependencies = [ - { name = "nvidia-nvjitlink", version = "13.4.92", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-nvjitlink", version = "13.4.92", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/be/b1/6e36381739b51e692d91646298ad5685c562929b899331c2bb6e9722a176/nvidia_cufft-12.2.0.46-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:2e8f3a22c2745f95327b7639ba2868ea2cbd5d3131ebe6313394e24164054f41", size = 218244743, upload-time = "2026-04-13T09:51:24.75Z" }, @@ -4582,9 +4823,9 @@ resolution-markers = [ "python_full_version < '3.12'", ] dependencies = [ - { name = "nvidia-cublas", version = "13.4.0.1", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "nvidia-cusparse", version = "12.7.10.1", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "nvidia-nvjitlink", version = "13.4.92", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cublas", version = "13.4.0.1", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cusparse", version = "12.7.10.1", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-nvjitlink", version = "13.4.92", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/94/21/21e5c491a8285a6ed039cbb8c60ec13ff3aff60c669e1caa4adaeb71f997/nvidia_cusolver-12.2.0.1-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:c9def09fcf300b9465cdf800142abf10149df28c4290bdf9e7b4a700480d31d7", size = 249744985, upload-time = "2026-04-13T09:54:33.757Z" }, @@ -4622,7 +4863,7 @@ resolution-markers = [ "python_full_version < '3.12'", ] dependencies = [ - { name = "nvidia-nvjitlink", version = "13.4.92", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-nvjitlink", version = "13.4.92", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/b5/09/9f1fa88dd305e838b71f2ca4e50064e7cbdf2e15cd6410d9e1b74066cc41/nvidia_cusparse-12.7.10.1-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:80cc98b38fdcf054a80bd9782fe2e63c66f90e202cac269f641357ccac5fe7e3", size = 168860233, upload-time = "2026-04-13T09:55:48.891Z" }, @@ -4640,6 +4881,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/31/83/f3647ce26916c94a6ca4ff1810623e2c405cff2dea6e78d29516b2514df9/nvidia_cusparselt_cu13-0.8.1-py3-none-win_amd64.whl", hash = "sha256:dccbd362f91a7b9024d1f55ee9f548ac065027ff15d8c8b0db889ab3a8f31215", size = 156885108, upload-time = "2025-09-05T18:51:35.958Z" }, ] +[[package]] +name = "nvidia-ml-py" +version = "13.615.71" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/fd/30/b25216758be3d3e2834825d8193609e2d71c770a8bd7984438c058c90268/nvidia_ml_py-13.615.71.tar.gz", hash = "sha256:bebe4e48f51b1dc75028c0815cb7bfa14a31a5bb80be70c9d980c6036953fc3d", size = 57485, upload-time = "2026-09-25T15:15:28.226Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/53/a1/1681dfa1c904d4e3e72e51b55a0ff012d50b766843ef832d584abe2113c6/nvidia_ml_py-13.615.71-py3-none-any.whl", hash = "sha256:959bf4adf6fe1308e4bd739e722236b0d1ec8392e2cefad33ff70c311380b9b6", size = 58132, upload-time = "2026-09-25T15:15:26.54Z" }, +] + [[package]] name = "nvidia-nccl-cu13" version = "2.29.7" @@ -5093,7 +5343,7 @@ name = "pick" version = "2.6.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "windows-curses", marker = "sys_platform == 'win32' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "windows-curses", marker = "sys_platform == 'win32'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/b6/f9/450cc44b21190d99b86df3e59f7573d85f3ef2c4f7670ce4f65cb2def3b4/pick-2.6.0.tar.gz", hash = "sha256:f24434a845ae9265ebf7657fc224b9a203ac75aa576b568e061484a72e27be82", size = 6266, upload-time = "2026-02-28T13:47:06.31Z" } wheels = [ @@ -5260,6 +5510,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/45/e2/bbb7129c9e7999a6b8ee9cca3b66486c25c423ab5a75f34071798b74ce94/pre_commit-4.6.2-py2.py3-none-any.whl", hash = "sha256:e2dde9a75d3bce11bd3831c26d134df00a2803c1d818be6a0383c3dcda25dc4e", size = 226202, upload-time = "2026-08-10T22:07:16.942Z" }, ] +[[package]] +name = "prettytable" +version = "3.18.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "wcwidth" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/81/74/ba08d81e668ccfe8658d7520a307e63c19862c08eb4ccb26f356c5239a7a/prettytable-3.18.0.tar.gz", hash = "sha256:439217116152244369caf3d9f1caf2f9fe29b03bd79e88d2928c8e718c95d680", size = 76373, upload-time = "2026-06-22T16:07:50.174Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fe/be/2e6798ace5cc036f5d05d36b7b2fd85346f1a708c87060890b070d0ec607/prettytable-3.18.0-py3-none-any.whl", hash = "sha256:b3346e0e6f79180833aebaac088ae926340586cf6d7d991b9eb125b65f72313a", size = 37357, upload-time = "2026-06-22T16:07:48.595Z" }, +] + [[package]] name = "prometheus-client" version = "0.14.1" @@ -5773,7 +6035,7 @@ version = "26.4.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cryptography" }, - { name = "typing-extensions", marker = "python_full_version < '3.13' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/3f/e8/7325d258199b159eb2c03fe32107533e2832e70e63f4fb88a6aa00023201/pyopenssl-26.4.0.tar.gz", hash = "sha256:28dfcce0162b9211413e26dfbfdf1d24317fbeba18fc93c12400a1856b2a0bc7", size = 182046, upload-time = "2026-08-01T19:50:50.512Z" } wheels = [ @@ -5806,7 +6068,7 @@ name = "pytest" version = "9.0.3" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "colorama", marker = "sys_platform == 'win32' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "colorama", marker = "sys_platform == 'win32'" }, { name = "iniconfig" }, { name = "packaging" }, { name = "pluggy" }, @@ -5823,7 +6085,7 @@ version = "1.3.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "pytest" }, - { name = "typing-extensions", marker = "python_full_version < '3.13' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/90/2c/8af215c0f776415f3590cac4f9086ccefd6fd463befeae41cd4d3f193e5a/pytest_asyncio-1.3.0.tar.gz", hash = "sha256:d7f52f36d231b80ee124cd216ffb19369aa168fc10095013c6b014a34d3ee9e5", size = 50087, upload-time = "2025-11-10T16:07:47.256Z" } wheels = [ @@ -5926,6 +6188,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/16/7d/e9ffbadfbf89c93848412d04594135c4ae8c1d37d9e053b9c3ed718fabc4/python_discovery-1.6.1-py3-none-any.whl", hash = "sha256:d43fcdef879fe795352bd13ccf8d185ba5a9f86f36cfcd00529f596e737442b3", size = 38664, upload-time = "2026-09-18T01:31:52.448Z" }, ] +[[package]] +name = "python-dotenv" +version = "1.2.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/74/26/2fbeedb218a787a5eea551c7532cac4e009f83d689dd2faa0d0353473f86/python_dotenv-1.2.4.tar.gz", hash = "sha256:f0d53e69935a851c0dcc78f3ab7aaccd8cabef0b92382b576b824212902873c0", size = 60824, upload-time = "2026-10-01T05:36:10Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/60/d1/38f3a3405989a89ac18390803e70c6ad7c7760da4f9b83cbeca0c44a0c72/python_dotenv-1.2.4-py3-none-any.whl", hash = "sha256:42269a8a5b3fd54ffa6f3d84b18abed50064717576b4ecf03dc4a55d8aa04fdc", size = 23266, upload-time = "2026-10-01T05:36:08.633Z" }, +] + [[package]] name = "python-logging-loki" version = "0.3.1" @@ -5957,11 +6228,11 @@ dependencies = [ { name = "pandas" }, { name = "pillow" }, { name = "pyarrow" }, - { name = "torch", version = "2.12.0", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "torch", version = "2.12.0", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, { name = "torch", version = "2.12.0", source = { registry = "https://pypi.org/simple" }, marker = "(extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu130' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests')" }, - { name = "torch", version = "2.12.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra != 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "torch", version = "2.12.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, { name = "torch", version = "2.12.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "extra == 'extra-4-kumo-cu130' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "torch", version = "2.12.0+cu132", source = { registry = "https://download.pytorch.org/whl/cu132" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "torch", version = "2.12.0+cu132", source = { registry = "https://download.pytorch.org/whl/cu132" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, { name = "tqdm" }, ] sdist = { url = "https://files.pythonhosted.org/packages/22/64/4a4aaa2a90eab208bf9f929a3bcc3b17ceb99bcb9abecf1c2323818f2898/pytorch_frame-0.3.0.tar.gz", hash = "sha256:412b42002e1a21a09fc8b0c6e4f0d048bd4e5ebef2f7139bbbbae82d81040152", size = 104795, upload-time = "2025-11-02T23:22:07.936Z" } @@ -5978,11 +6249,11 @@ dependencies = [ { name = "lightning-utilities" }, { name = "packaging" }, { name = "pyyaml" }, - { name = "torch", version = "2.12.0", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "torch", version = "2.12.0", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, { name = "torch", version = "2.12.0", source = { registry = "https://pypi.org/simple" }, marker = "(extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu130' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests')" }, - { name = "torch", version = "2.12.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra != 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "torch", version = "2.12.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, { name = "torch", version = "2.12.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "extra == 'extra-4-kumo-cu130' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "torch", version = "2.12.0+cu132", source = { registry = "https://download.pytorch.org/whl/cu132" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "torch", version = "2.12.0+cu132", source = { registry = "https://download.pytorch.org/whl/cu132" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, { name = "torchmetrics" }, { name = "tqdm" }, { name = "typing-extensions" }, @@ -6001,6 +6272,28 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/0f/7b/39c34ca613b0b198cb866466651b26b045e2009864c5183c979a3b83f383/pytz-2026.3.post1-py2.py3-none-any.whl", hash = "sha256:dd95840dd199baea12d9cc096a1d452caa6596a1c1e4b5f3dbd1541855d5e815", size = 508283, upload-time = "2026-07-25T15:12:05.782Z" }, ] +[[package]] +name = "pywin32" +version = "312" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/1f/f5/10a6e845a00fc5e7afd0a988b744f403d4d57162a28d160a093c4d9322f0/pywin32-312-cp311-cp311-win32.whl", hash = "sha256:17948aeadbdb091f0ced6ef0841620794e68327b94ee415571c1203594b7215c", size = 6362659, upload-time = "2026-06-04T07:49:21.349Z" }, + { url = "https://files.pythonhosted.org/packages/35/c4/dcd2d62b5944b6d5db53413a5899016ccd57ffcb7278f3f81655d25d2027/pywin32-312-cp311-cp311-win_amd64.whl", hash = "sha256:d11417d84412f859b722fad0841b3614459ed0047f7542d8362e77884f6b6e8a", size = 6928825, upload-time = "2026-06-04T07:49:23.934Z" }, + { url = "https://files.pythonhosted.org/packages/b7/56/3cbb433fe4501cdba2eb9040f56a4e1a8243faa4186b25295564d1a7a79d/pywin32-312-cp311-cp311-win_arm64.whl", hash = "sha256:b2200a054ca6d6625c4842fc56a4976a4b47f96b73dbe5538c3f813a80359f47", size = 6721875, upload-time = "2026-06-04T07:49:26.416Z" }, + { url = "https://files.pythonhosted.org/packages/83/ff/32aa7d2ed0ab12b323aaa64f9b75e6ad4f8fd09f9ccfc28c79414d46838d/pywin32-312-cp312-cp312-win32.whl", hash = "sha256:dab4f65ac9c4e48400a2a0530c46c3c579cd5905ecd11b80692373915269208b", size = 6371877, upload-time = "2026-06-04T07:49:28.836Z" }, + { url = "https://files.pythonhosted.org/packages/03/d9/77040d3b43df3f3be32ea289433d660d2727f5ba327bc73be835127d9d60/pywin32-312-cp312-cp312-win_amd64.whl", hash = "sha256:b457f6d628a47e8a7346ce22acb7e1a46a4a78b52e1d17e1af56871bd19a93bc", size = 6914841, upload-time = "2026-06-04T07:49:31.85Z" }, + { url = "https://files.pythonhosted.org/packages/e3/cc/7b1ec671775756020a0ee7f4feeaf3c568f0ab86bd3900088cf986937a92/pywin32-312-cp312-cp312-win_arm64.whl", hash = "sha256:6017c58e12f6809fbb0555b75df144c2922a9ffd18e4b9b5afa863b6c1a9d950", size = 6727901, upload-time = "2026-06-04T07:49:34.244Z" }, + { url = "https://files.pythonhosted.org/packages/2d/41/12fbfd7f36ed2146d8bc9de96c2741296bf0d490b98508496cff322e274c/pywin32-312-cp313-cp313-win32.whl", hash = "sha256:7a27df850933d16a8eabfbaeb73d52b273e2da667f80d70b01a89d1f6828d02c", size = 6370184, upload-time = "2026-06-04T07:49:36.253Z" }, + { url = "https://files.pythonhosted.org/packages/ba/db/36a78e3403099d31d9746d13fdcde5accc43c1155f375a34d15983a479a7/pywin32-312-cp313-cp313-win_amd64.whl", hash = "sha256:c53e878d15a1c44788082bfe712a905433473aa38f86375b7cf8b45e3acbaaf9", size = 6914298, upload-time = "2026-06-04T07:49:38.876Z" }, + { url = "https://files.pythonhosted.org/packages/84/37/c1697194092b76de9ed47ca124323f02c57ffc8a45c06f88a3d5acaf01eb/pywin32-312-cp313-cp313-win_arm64.whl", hash = "sha256:59aba5d5940842075343a5ddc6b11f1cdf0d1567fe745290359dfbcc7c2eb831", size = 6727640, upload-time = "2026-06-04T07:49:41.083Z" }, + { url = "https://files.pythonhosted.org/packages/fc/2b/1f3cded5822fd49c02f40544cbb5f58c7cfd6b1694869fd476cb6170ee97/pywin32-312-cp314-cp314-win32.whl", hash = "sha256:a77a90fbb6881238d2ca9c6fd797b25817f3768fe78d214a90137ff055a75f5b", size = 6468928, upload-time = "2026-06-04T07:49:43.188Z" }, + { url = "https://files.pythonhosted.org/packages/21/82/3bf86d2e2808902013132e1ce905a7da0da53790f3836c64bf44d55e24f3/pywin32-312-cp314-cp314-win_amd64.whl", hash = "sha256:a4dd3a848290ef724347b19f301045831d8e802fa4464f491b98b1e0a081432e", size = 7024157, upload-time = "2026-06-04T07:49:45.34Z" }, + { url = "https://files.pythonhosted.org/packages/a4/0e/73f6d6800b4f27655abd9e9f6aaeaefcddb2b946e4674efa2bab184a7f7b/pywin32-312-cp314-cp314-win_arm64.whl", hash = "sha256:9fce94568364e0155e6dfb781ac5d95903be8baf28670632beab1b523f300daa", size = 6839598, upload-time = "2026-06-04T07:49:47.613Z" }, + { url = "https://files.pythonhosted.org/packages/eb/61/caa39686032d2ebdd04ff0ab5cbe163126c0066d98e00c9018646e42393b/pywin32-312-cp315-cp315-win32.whl", hash = "sha256:5c1fbe4a937a73ae9297384a3da38518cbc694c68ad8a809b2e19acd350f03ed", size = 6471159, upload-time = "2026-06-04T07:49:50.035Z" }, + { url = "https://files.pythonhosted.org/packages/0f/cd/7e1de64a4a6f69c04214169657ccab0d93a670ea50e35eb8f489d7378249/pywin32-312-cp315-cp315-win_amd64.whl", hash = "sha256:c2f03a0f73f804a13c2735b99392b0cd426bb4f2c4d0178e5ac966a0f21618d5", size = 7025293, upload-time = "2026-06-04T07:49:54.857Z" }, + { url = "https://files.pythonhosted.org/packages/23/ed/4532e9388e65fa16b46776ef47ad631a64eda1631884488af707666350ed/pywin32-312-cp315-cp315-win_arm64.whl", hash = "sha256:a8597d28f267b39074aef51fa593530082b39cbe5a074226096857b1fed2dfb9", size = 6840337, upload-time = "2026-06-04T07:49:57.531Z" }, +] + [[package]] name = "pyyaml" version = "6.0.2" @@ -6209,6 +6502,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/03/19/90d7d4ed51932c022d53f1d02d564b62d10e272692a1f9b76425c1ad2a02/s3transfer-0.16.1-py3-none-any.whl", hash = "sha256:61bcd00ccb83b21a0fe7e91a553fff9729d46c83b4e0106e7c314a733891f7c2", size = 86825, upload-time = "2026-04-22T20:36:04.992Z" }, ] +[[package]] +name = "sagemaker-mlflow" +version = "0.5.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "boto3" }, + { name = "mlflow-skinny" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/d5/e6/c147424bb1b1d16963037940949533dc893b511cf5919ecbc82388dfcf0d/sagemaker_mlflow-0.5.0.tar.gz", hash = "sha256:6786b9aa738bf0ec551dc99bd44069ac39ad202fa234f99b2fc914556f1af2f2", size = 22371, upload-time = "2026-06-27T00:06:19.285Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/69/50/4d0a43a18e0596f2fd1a506a4467577c45d2697b245cb5f61e8704209c29/sagemaker_mlflow-0.5.0-py3-none-any.whl", hash = "sha256:f55dc39c6ed345d4359c7d5cddd99fb59c852609c463a397d18e8bbaba672dac", size = 31159, upload-time = "2026-06-27T00:06:18.253Z" }, +] + [[package]] name = "scikit-learn" version = "1.9.1" @@ -6217,8 +6523,8 @@ dependencies = [ { name = "joblib" }, { name = "narwhals" }, { name = "numpy" }, - { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "scipy", version = "1.18.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, + { name = "scipy", version = "1.18.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, { name = "threadpoolctl" }, ] sdist = { url = "https://files.pythonhosted.org/packages/d2/eb/eaf5e07fcc0da7149b0e084f24e54edd7441b9a89ce7e034032ae97fe3a0/scikit_learn-1.9.1.tar.gz", hash = "sha256:629cada3e33e2b9bf376cdc7614a47a4140b8aedc1d836579e359736fbd82977", size = 7786908, upload-time = "2026-09-10T18:34:04.679Z" } @@ -6281,7 +6587,7 @@ resolution-markers = [ "python_full_version < '3.12' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu130' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests'", ] dependencies = [ - { name = "numpy", marker = "python_full_version < '3.12' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "numpy" }, ] sdist = { url = "https://files.pythonhosted.org/packages/7a/97/5a3609c4f8d58b039179648e62dd220f89864f56f7357f5d4f45c29eb2cc/scipy-1.17.1.tar.gz", hash = "sha256:95d8e012d8cb8816c226aef832200b1d45109ed4464303e997c5b13122b297c0", size = 30573822, upload-time = "2026-02-23T00:26:24.851Z" } wheels = [ @@ -6375,7 +6681,7 @@ resolution-markers = [ "python_full_version == '3.12.*' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu130' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests'", ] dependencies = [ - { name = "numpy", marker = "python_full_version >= '3.12' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "numpy" }, ] sdist = { url = "https://files.pythonhosted.org/packages/7e/74/66de6258867beb2ef08f35f9f2ac017a52cacd5081714d239ff1a442d458/scipy-1.18.1.tar.gz", hash = "sha256:52c4b7422442aba924d03ad4019852b08a92e64ea187b933135687bfe2747307", size = 30781235, upload-time = "2026-08-21T23:28:50.599Z" } wheels = [ @@ -6486,6 +6792,23 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274", size = 11050, upload-time = "2024-12-04T17:35:26.475Z" }, ] +[[package]] +name = "skops" +version = "0.16.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, + { name = "packaging" }, + { name = "prettytable" }, + { name = "scikit-learn" }, + { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, + { name = "scipy", version = "1.18.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/52/4b/7fef0af1efcb06c0986762645fb414ad850314a891434a2bbaea12a962ea/skops-0.16.0.tar.gz", hash = "sha256:58ddb63fed9f2a8545baa1a91c45c26d44cf65c83e2fd217541c41e411cf1d74", size = 673544, upload-time = "2026-09-24T06:29:46.96Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/44/a3/08ae9ea7f7c16e20a447d9289dce40f06f52607ed9a39f554a8426d96ba3/skops-0.16.0-py3-none-any.whl", hash = "sha256:5b2eb1bc49392c74a38db04481e2d110f97d9a72d6ae38fbd698aa7a865a9070", size = 148650, upload-time = "2026-09-24T06:29:45.491Z" }, +] + [[package]] name = "smart-open" version = "5.2.1" @@ -6495,6 +6818,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/cd/11/05f68ea934c24ade38e95ac30a38407767787c4e3db1776eae4886ad8c95/smart_open-5.2.1-py3-none-any.whl", hash = "sha256:71d14489da58b60ce12fc3ecb823facc59a8b23cd1b58edb97175640350d3a62", size = 58638, upload-time = "2021-08-28T11:06:50.07Z" }, ] +[[package]] +name = "smmap" +version = "5.0.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/1f/ea/49c993d6dfdd7338c9b1000a0f36817ed7ec84577ae2e52f890d1a4ff909/smmap-5.0.3.tar.gz", hash = "sha256:4d9debb8b99007ae47165abc08670bd74cb74b5227dda7f643eccc4e9eb5642c", size = 22506, upload-time = "2026-03-09T03:43:26.1Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c1/d4/59e74daffcb57a07668852eeeb6035af9f32cbfd7a1d2511f17d2fe6a738/smmap-5.0.3-py3-none-any.whl", hash = "sha256:c106e05d5a61449cf6ba9a1e650227ecfb141590d2a98412103ff35d89fc7b2f", size = 24390, upload-time = "2026-03-09T03:43:24.361Z" }, +] + [[package]] name = "sniffio" version = "1.3.1" @@ -6592,8 +6924,8 @@ dependencies = [ { name = "pytimeparse" }, { name = "pyyaml" }, { name = "retrying" }, - { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "scipy", version = "1.18.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, + { name = "scipy", version = "1.18.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, { name = "snowflake-connector-python", extra = ["pandas"] }, { name = "snowflake-snowpark-python" }, { name = "sqlparse" }, @@ -6650,7 +6982,7 @@ source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "alabaster" }, { name = "babel" }, - { name = "colorama", marker = "sys_platform == 'win32' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "colorama", marker = "sys_platform == 'win32'" }, { name = "docutils" }, { name = "imagesize" }, { name = "jinja2" }, @@ -6804,7 +7136,7 @@ name = "sqlalchemy" version = "2.0.54" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "greenlet", marker = "platform_machine == 'AMD64' or platform_machine == 'WIN32' or platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'ppc64le' or platform_machine == 'win32' or platform_machine == 'x86_64' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "greenlet", marker = "platform_machine == 'AMD64' or platform_machine == 'WIN32' or platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'ppc64le' or platform_machine == 'win32' or platform_machine == 'x86_64'" }, { name = "typing-extensions" }, ] sdist = { url = "https://files.pythonhosted.org/packages/29/9c/271aa905cf2964f841371a97f3e63ab692bf51b4423d0491e67bc7f64037/sqlalchemy-2.0.54.tar.gz", hash = "sha256:baa8521e8ee9f24e75dfc7aaabc08020e551ef0d48d7c3e3536f5cddf277586b", size = 9969559, upload-time = "2026-09-15T21:06:57.337Z" } @@ -6871,7 +7203,7 @@ version = "1.3.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "anyio" }, - { name = "typing-extensions", marker = "python_full_version < '3.13' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/eb/e3/7c1dc7381d9f8ab7d854328ebfa884e62cb3f3d8549ddfd37c7814f42afa/starlette-1.3.1.tar.gz", hash = "sha256:05d0213193f2fbaae60e2ecb593b4add4262ad4e46536b54abe36f11a71724e0", size = 2703240, upload-time = "2026-06-12T09:23:11.602Z" } wheels = [ @@ -7077,14 +7409,14 @@ resolution-markers = [ "python_full_version < '3.12' and sys_platform == 'darwin'", ] dependencies = [ - { name = "filelock", marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "fsspec", extra = ["http"], marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "jinja2", marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "networkx", version = "3.6.1", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.12' and sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (python_full_version < '3.12' and sys_platform == 'darwin' and extra != 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (python_full_version >= '3.12' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (python_full_version >= '3.12' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (python_full_version >= '3.12' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (python_full_version >= '3.12' and extra != 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'darwin' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'darwin' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'darwin' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra != 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra != 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "networkx", version = "3.7", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.12' and sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (python_full_version >= '3.12' and sys_platform == 'darwin' and extra != 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (python_full_version < '3.12' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (python_full_version < '3.12' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (python_full_version < '3.12' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (python_full_version < '3.12' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (python_full_version < '3.12' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (python_full_version < '3.12' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'darwin' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'darwin' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'darwin' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra != 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra != 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "setuptools", marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "sympy", marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "typing-extensions", marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "filelock", marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "fsspec", extra = ["http"], marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "jinja2", marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "networkx", version = "3.6.1", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.12' and sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (python_full_version < '3.12' and sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "networkx", version = "3.7", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.12' and sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (python_full_version >= '3.12' and sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "setuptools", marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "sympy", marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "typing-extensions", marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] wheels = [ { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.12.0-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:10802fd383bbfed646212e765a72c37d2185205d4f26eb197a254e8ac7ddcb25", upload-time = "2026-05-12T16:20:07Z" }, @@ -7161,14 +7493,14 @@ resolution-markers = [ "python_full_version < '3.12' and sys_platform != 'darwin'", ] dependencies = [ - { name = "filelock", marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra != 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "fsspec", extra = ["http"], marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra != 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "jinja2", marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra != 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "networkx", version = "3.6.1", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.12' and sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (python_full_version < '3.12' and sys_platform != 'darwin' and extra != 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (python_full_version >= '3.12' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (python_full_version >= '3.12' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (python_full_version >= '3.12' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (python_full_version >= '3.12' and extra != 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'darwin' and extra != 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra != 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra != 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "networkx", version = "3.7", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.12' and sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (python_full_version >= '3.12' and sys_platform != 'darwin' and extra != 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (python_full_version < '3.12' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (python_full_version < '3.12' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (python_full_version < '3.12' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (python_full_version < '3.12' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (python_full_version < '3.12' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (python_full_version < '3.12' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'darwin' and extra != 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra != 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra != 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "setuptools", marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra != 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "sympy", marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra != 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "typing-extensions", marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra != 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "filelock", marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "fsspec", extra = ["http"], marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "jinja2", marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "networkx", version = "3.6.1", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.12' and sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (python_full_version < '3.12' and sys_platform != 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "networkx", version = "3.7", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.12' and sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (python_full_version >= '3.12' and sys_platform != 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "setuptools", marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "sympy", marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "typing-extensions", marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] wheels = [ { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.12.0%2Bcpu-cp311-cp311-linux_s390x.whl", hash = "sha256:aaa9c1f5e8c518d7be1e3c3e1b090ca7d63b6e353a1abd6cfdaf902405093467", upload-time = "2026-05-12T23:16:01Z" }, @@ -7211,21 +7543,21 @@ resolution-markers = [ "python_full_version < '3.12'", ] dependencies = [ - { name = "cuda-bindings", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "cuda-toolkit", version = "13.0.2", source = { registry = "https://pypi.org/simple" }, extra = ["cudart", "cufft", "cufile", "cupti", "curand", "cusolver", "cusparse", "nvjitlink", "nvrtc", "nvtx"], marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "cuda-bindings", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "cuda-toolkit", version = "13.0.2", source = { registry = "https://pypi.org/simple" }, extra = ["cudart", "cufft", "cufile", "cupti", "curand", "cusolver", "cusparse", "nvjitlink", "nvrtc", "nvtx"], marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, { name = "filelock", marker = "extra == 'extra-4-kumo-cu130' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, { name = "fsspec", extra = ["http"], marker = "extra == 'extra-4-kumo-cu130' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, { name = "jinja2", marker = "extra == 'extra-4-kumo-cu130' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, { name = "networkx", version = "3.6.1", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.12' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "networkx", version = "3.7", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.12' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "nvidia-cublas", version = "13.1.1.3", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "nvidia-cudnn-cu13", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "nvidia-cusparselt-cu13", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "nvidia-nccl-cu13", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "nvidia-nvshmem-cu13", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "networkx", version = "3.7", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.12' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cublas", version = "13.1.1.3", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cudnn-cu13", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cusparselt-cu13", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-nccl-cu13", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-nvshmem-cu13", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, { name = "setuptools", marker = "extra == 'extra-4-kumo-cu130' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, { name = "sympy", marker = "extra == 'extra-4-kumo-cu130' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "triton", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "triton", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, { name = "typing-extensions", marker = "extra == 'extra-4-kumo-cu130' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] wheels = [ @@ -7260,21 +7592,21 @@ resolution-markers = [ "python_full_version < '3.12'", ] dependencies = [ - { name = "cuda-bindings", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "cuda-toolkit", version = "13.2.1", source = { registry = "https://pypi.org/simple" }, extra = ["cublas", "cudart", "cufft", "cufile", "cupti", "curand", "cusolver", "cusparse", "nvjitlink", "nvrtc", "nvtx"], marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "filelock", marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "fsspec", extra = ["http"], marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "jinja2", marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "cuda-bindings", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "cuda-toolkit", version = "13.2.1", source = { registry = "https://pypi.org/simple" }, extra = ["cublas", "cudart", "cufft", "cufile", "cupti", "curand", "cusolver", "cusparse", "nvjitlink", "nvrtc", "nvtx"], marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "filelock", marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "fsspec", extra = ["http"], marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "jinja2", marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, { name = "networkx", version = "3.6.1", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.12' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "networkx", version = "3.7", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.12' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "nvidia-cudnn-cu13", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "nvidia-cusparselt-cu13", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "nvidia-nccl-cu13", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "nvidia-nvshmem-cu13", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "setuptools", marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "sympy", marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "triton", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "typing-extensions", marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "networkx", version = "3.7", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.12' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cudnn-cu13", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cusparselt-cu13", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-nccl-cu13", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-nvshmem-cu13", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "setuptools", marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "sympy", marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "triton", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "typing-extensions", marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, ] wheels = [ { url = "https://download-r2.pytorch.org/whl/cu132/torch-2.12.0%2Bcu132-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:c8b75ec50e9a6bfc6db0f1e0bd297a1b94c9803c378c706c4e856c7cad331398", upload-time = "2026-05-13T00:04:15Z" }, @@ -7325,11 +7657,11 @@ dependencies = [ { name = "lightning-utilities" }, { name = "numpy" }, { name = "packaging" }, - { name = "torch", version = "2.12.0", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "torch", version = "2.12.0", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, { name = "torch", version = "2.12.0", source = { registry = "https://pypi.org/simple" }, marker = "(extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu130' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests')" }, - { name = "torch", version = "2.12.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra != 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "torch", version = "2.12.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, { name = "torch", version = "2.12.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "extra == 'extra-4-kumo-cu130' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "torch", version = "2.12.0+cu132", source = { registry = "https://download.pytorch.org/whl/cu132" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "torch", version = "2.12.0+cu132", source = { registry = "https://download.pytorch.org/whl/cu132" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/78/1f/2cd9eb8f3390c3ec4693ac0871913d4b468964b3833638e4091a70817e0a/torchmetrics-1.8.1.tar.gz", hash = "sha256:04ca021105871637c5d34d0a286b3ab665a1e3d2b395e561f14188a96e862fdb", size = 580373, upload-time = "2025-08-07T20:44:44.631Z" } wheels = [ @@ -7341,7 +7673,7 @@ name = "tqdm" version = "4.70.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "colorama", marker = "sys_platform == 'win32' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "colorama", marker = "sys_platform == 'win32'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/0d/ea/b2a5bd54b28a324dae8211928b2d730b6547500342c7e6c6dea08bd0a485/tqdm-4.70.1.tar.gz", hash = "sha256:cefd0eca11b2a37a3aee776544d4f4ae913f02688135b5556b8788dfa474afc4", size = 171846, upload-time = "2026-09-11T07:25:16.601Z" } wheels = [ @@ -7409,7 +7741,7 @@ version = "3.7.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "botocore-stubs" }, - { name = "typing-extensions", marker = "python_full_version < '3.12' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "typing-extensions", marker = "python_full_version < '3.12'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/76/e8/ef1fcb876937dbdddc0f01b5df4ed53f33b166a6367d80a9014d5e5f091d/types_aiobotocore-3.7.0.tar.gz", hash = "sha256:fe35de52c12e5fdb89ca60b3989766e7fe827e3d2e95fcf4583e91581945205c", size = 87992, upload-time = "2026-05-10T03:19:32.353Z" } wheels = [ @@ -7426,7 +7758,7 @@ name = "types-aiobotocore-s3" version = "3.7.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions", marker = "python_full_version < '3.12' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "typing-extensions", marker = "python_full_version < '3.12'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/94/ce/5b7baa06cde79ab5f083f56f060ccd2b60f6249127d70bbb8a37aafbdcd4/types_aiobotocore_s3-3.7.0.tar.gz", hash = "sha256:6ec738853dbba9133707991b98ea4dab19f7c62e02b3cca016e6cd8e3d684576", size = 77662, upload-time = "2026-05-10T03:17:35.088Z" } wheels = [ @@ -7494,7 +7826,7 @@ name = "tzlocal" version = "5.4.4" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "tzdata", marker = "sys_platform == 'win32' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "tzdata", marker = "sys_platform == 'win32'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/81/5b/879b2f932adfa7a053c360d50bc896c977fa6426109185f7c12ebdd0cb9d/tzlocal-5.4.4.tar.gz", hash = "sha256:8dbb8660838688a7b6ba4fed31d18dedf842afb4d47ca050d6d891c2c15f3be4", size = 31170, upload-time = "2026-06-29T08:03:40.026Z" } wheels = [ @@ -7538,6 +7870,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/9d/e5/9446b7a2a3d73de5f9fcf8f825efe78b071b111729e89f4ac07649b9448b/virtualenv-21.9.0-py3-none-any.whl", hash = "sha256:8c7cde05f630fbedf61758258014a92e6d6d5b4b9c6e6be7279d233f89e09aac", size = 5477625, upload-time = "2026-09-20T00:29:18.065Z" }, ] +[[package]] +name = "waitress" +version = "3.0.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/bf/cb/04ddb054f45faa306a230769e868c28b8065ea196891f09004ebace5b184/waitress-3.0.2.tar.gz", hash = "sha256:682aaaf2af0c44ada4abfb70ded36393f0e307f4ab9456a215ce0020baefc31f", size = 179901, upload-time = "2024-11-16T20:02:35.195Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8d/57/a27182528c90ef38d82b636a11f606b0cbb0e17588ed205435f8affe3368/waitress-3.0.2-py3-none-any.whl", hash = "sha256:c56d67fd6e87c2ee598b76abdd4e96cfad1f24cacdea5078d382b1f9d7b5ed2e", size = 56232, upload-time = "2024-11-16T20:02:33.858Z" }, +] + [[package]] name = "watchdog" version = "6.0.0" @@ -7582,6 +7923,42 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/23/41/36c7246c1c273fd316d9b8b81581074becc3137135f9ccd6307960a7879c/watchtower-3.0.1-py3-none-any.whl", hash = "sha256:566874a9dee34b90517b9a9b7c2d5e6b15fb9fd00dacb03bff8a7bf4c6295fe0", size = 17331, upload-time = "2023-01-29T15:33:52.494Z" }, ] +[[package]] +name = "wcwidth" +version = "0.9.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f0/b4/7830542634bb2d3e62aa3b586a72d5b3b6c91c3168929e7000ef3fed041d/wcwidth-0.9.2.tar.gz", hash = "sha256:ae0ef90b90f6af38b54f1fe6d58662ec33b3cb4b8391958a62416d654231727b", size = 955039, upload-time = "2026-10-05T00:24:05.521Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/59/1e/4532a81fb9dfbf4114a816775e0a36c3a64ee1d1f4bba2094e2da50be5dc/wcwidth-0.9.2-cp310-abi3-macosx_10_9_x86_64.whl", hash = "sha256:7ef5a940bd5e30bac6e721f1a48fce0cd7bb3ece19e9c5d139e72c76c35cfd07", size = 605343, upload-time = "2026-10-05T00:23:22.649Z" }, + { url = "https://files.pythonhosted.org/packages/a0/07/cb6940e81134b7ed25fa312ee9ab536a63db0793b149f88a90e603ceace9/wcwidth-0.9.2-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:ae0800c5339423cc53d33a266ad264b42ba8aaa16d4464f6e6b1bee607f50b17", size = 609960, upload-time = "2026-10-05T00:23:27.049Z" }, + { url = "https://files.pythonhosted.org/packages/a4/80/15ad05d40bfa99155639fb9e13b3d77083aa0fab893c816db2543d29005c/wcwidth-0.9.2-cp310-abi3-manylinux1_i686.manylinux2014_i686.manylinux_2_17_i686.manylinux_2_5_i686.whl", hash = "sha256:9e542f1f8475b78452a295495d7a5bc3ead565112e9446a64dc93462a41c2a79", size = 762566, upload-time = "2026-10-05T00:23:38.322Z" }, + { url = "https://files.pythonhosted.org/packages/bc/f0/b8ef7758003d66b60f093695831a86dcc726aac01ee6446ffcbda27b61e3/wcwidth-0.9.2-cp310-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:674b518af28d38ee645ff97b74f5760abee5fad4bac74413bfc4b881ef2ce724", size = 771632, upload-time = "2026-10-05T00:23:32.448Z" }, + { url = "https://files.pythonhosted.org/packages/db/6c/f940133c71427c208575910e981942bd78c98b1f7cd0d1425ca4b7457c04/wcwidth-0.9.2-cp310-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:751bef0ab404b6a1dc028b56b4b85d46486be1c55833f80da533e42dc691f389", size = 769763, upload-time = "2026-10-05T00:23:40.175Z" }, + { url = "https://files.pythonhosted.org/packages/92/8f/285f862826f721964ec7c42f81dc53d23afbd723a0f4cd989651f8218e25/wcwidth-0.9.2-cp310-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:c3d80f39ba4653a595edae9aa46a509d14883790a8fc23c5db221ceb207f64b7", size = 782134, upload-time = "2026-10-05T00:23:33.926Z" }, + { url = "https://files.pythonhosted.org/packages/c2/2d/64aa54882a5d556d3654c1f926d9118b797461033e23a158409941a37c8f/wcwidth-0.9.2-cp310-abi3-musllinux_1_2_i686.whl", hash = "sha256:0a47e03d8293590ecce66c45dc20ff7b4b885e3c78093722239585eca0d77ab2", size = 782851, upload-time = "2026-10-05T00:23:41.974Z" }, + { url = "https://files.pythonhosted.org/packages/59/39/52389f6de7fe2e9c14ceb8253dd99034bd86e1c87847ea3c100a97dded9a/wcwidth-0.9.2-cp310-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:67d901a4ad99249eb775b4ee4769ca97fa405d35a75f46e83166910a47003f04", size = 784621, upload-time = "2026-10-05T00:23:43.449Z" }, + { url = "https://files.pythonhosted.org/packages/b3/8b/20225500a076ace27bbcc8a6fd7c55125133c57a618816c7b7b8b73070b1/wcwidth-0.9.2-cp310-abi3-win32.whl", hash = "sha256:ee1fd0db9d9fd711a70f3e7765e0e04c05d26982fa05361456163062549d7da4", size = 593123, upload-time = "2026-10-05T00:23:55.953Z" }, + { url = "https://files.pythonhosted.org/packages/5a/d6/b0690f55ea0483530a18bac917fbadbf54f35122510446fc370f5f1c2453/wcwidth-0.9.2-cp310-abi3-win_amd64.whl", hash = "sha256:2a9746de704242bd4fdaabb31dd46b82f694a56a8d21081ad89b679a89da9fec", size = 597827, upload-time = "2026-10-05T00:23:57.489Z" }, + { url = "https://files.pythonhosted.org/packages/e5/11/6ecf4e9e268ab1a4ec617ffcccc2ee4a71301625f5490912dbaba462fa9c/wcwidth-0.9.2-cp310-abi3-win_arm64.whl", hash = "sha256:b9c6ab615e03723b7f8760ea2f27758d656e7e13b51515c9dca5c3e8b04612fa", size = 598415, upload-time = "2026-10-05T00:23:51.517Z" }, + { url = "https://files.pythonhosted.org/packages/4e/41/549eef1ab767032bdbdc1f0ab655d404b082b1e9a1dab1361dbba90f64ed/wcwidth-0.9.2-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:eda88ffdc97c0fbf193d407114f2c7a54b379f67f6e52a7531ee3b9fe749eca7", size = 606555, upload-time = "2026-10-05T00:23:24.188Z" }, + { url = "https://files.pythonhosted.org/packages/9b/64/a875ed7ea71cacadc0ae11b5fd3fac3486efd58bb25e67a7344248dceadd/wcwidth-0.9.2-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:1bf361c8705576760623b4724ae564666d73b016f9a778bcfd1c7345378ef4ec", size = 610829, upload-time = "2026-10-05T00:23:28.563Z" }, + { url = "https://files.pythonhosted.org/packages/c6/98/513095e484fe79b6f2613d6a72f855f5d56b65e15c215c2a6746fbc638f5/wcwidth-0.9.2-cp314-cp314t-manylinux1_i686.manylinux2014_i686.manylinux_2_17_i686.manylinux_2_5_i686.whl", hash = "sha256:97b878d1e158da5ed9ac5aac53fa3a55e282103af6a09ec353865613d1a31a76", size = 775819, upload-time = "2026-10-05T00:23:45.116Z" }, + { url = "https://files.pythonhosted.org/packages/22/fc/c02f3eec57224731e78f84b68e272250f784b6205acc7e0dcef6a7c23a0e/wcwidth-0.9.2-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:59dab4049cbd982b478bca098528df2c79a9160636a3a163ffebffcbd7d1b892", size = 786915, upload-time = "2026-10-05T00:23:35.323Z" }, + { url = "https://files.pythonhosted.org/packages/3d/6f/b0529a79bac3fe8d94f32b4237a13dbc3f955508753f6a6f06c73d679dc2/wcwidth-0.9.2-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:bb08ceb501d6aaf94066c3ee122dd825b152df40ff0bd0df4dc27126233b948e", size = 784186, upload-time = "2026-10-05T00:23:46.366Z" }, + { url = "https://files.pythonhosted.org/packages/d5/bd/6357c84ca9a734bfc735b7c48dbe21336b3777fab8a4101d14976dfe49a7/wcwidth-0.9.2-cp314-cp314t-win32.whl", hash = "sha256:8b4e381590b9b7390e07e22b2c0c1bb96ce50e1d2243c866d9387600362d51ed", size = 600791, upload-time = "2026-10-05T00:23:59.398Z" }, + { url = "https://files.pythonhosted.org/packages/98/de/037591ca18d897cc2179559dde72e6efc6ce0c90e9cd1e6bca4e87c38b4b/wcwidth-0.9.2-cp314-cp314t-win_amd64.whl", hash = "sha256:f2f7b3bba5a5d5f31fc350fd36ce5b84b693c83b7eb95ee630b720da5a5ce06f", size = 605463, upload-time = "2026-10-05T00:24:01.049Z" }, + { url = "https://files.pythonhosted.org/packages/d0/07/c9d96e106d938d26f7ab639bc80b8199359a1645ba6e3498413313ab6f38/wcwidth-0.9.2-cp314-cp314t-win_arm64.whl", hash = "sha256:734aa9405b321d1042301aa19c943c4731ee9e3460e4f8feea3299c064c97a14", size = 605936, upload-time = "2026-10-05T00:23:52.765Z" }, + { url = "https://files.pythonhosted.org/packages/82/8a/a28d61d910005ac93dfe48be3a0ebaa49352d88cebd25323e69e6ff2f4a8/wcwidth-0.9.2-cp315-cp315t-macosx_10_15_x86_64.whl", hash = "sha256:42dbcb76ce8af39e2c9db410ac3f9bdf4e47eb41d6f44525952f172d3d98f724", size = 606575, upload-time = "2026-10-05T00:23:25.663Z" }, + { url = "https://files.pythonhosted.org/packages/01/c2/a3c66bd32766c8f4d6dc47d572532ba014fe5be30489f2576aff7cada363/wcwidth-0.9.2-cp315-cp315t-macosx_11_0_arm64.whl", hash = "sha256:138e1f8898e431b2f2d7881f8ca8d75591c1d3c21aa53f54e989bd6b39811da2", size = 611045, upload-time = "2026-10-05T00:23:30.421Z" }, + { url = "https://files.pythonhosted.org/packages/ec/8a/d39964f8f8c019d7d439b9b501d3e7bb42fee69f00354040ba0b27b5824c/wcwidth-0.9.2-cp315-cp315t-manylinux1_i686.manylinux2014_i686.manylinux_2_17_i686.manylinux_2_5_i686.whl", hash = "sha256:5175609bf8cc7398a5f48aa35207bd64ebf9f45e4c70df65f7fdc7a988041a3c", size = 778143, upload-time = "2026-10-05T00:23:47.7Z" }, + { url = "https://files.pythonhosted.org/packages/2f/53/525da13e8f9ff7b5b4e74ec6f8d68bdee63905796972e086c6b1b96670d2/wcwidth-0.9.2-cp315-cp315t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:e5f669ae8c3d969c72032f9cdee019674b666e522d45e1e2099a2e9dda4a341d", size = 786630, upload-time = "2026-10-05T00:23:36.967Z" }, + { url = "https://files.pythonhosted.org/packages/ef/9f/d6a0c6df354b9d93466548a65cbf4ffcb48c719bbd307504cf3e76740837/wcwidth-0.9.2-cp315-cp315t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:196b47cf32f9df27ccda6dc513237f3c2429c4c659db428d60a5bc443d10f270", size = 784839, upload-time = "2026-10-05T00:23:49.88Z" }, + { url = "https://files.pythonhosted.org/packages/bf/d7/3021feed1ed7926021ec134943ad3b24a2f7ea742cc9976461171482ed77/wcwidth-0.9.2-cp315-cp315t-win32.whl", hash = "sha256:0cd4f7f2e53905dcb110d213a4c8529b6733fa3d232d8c717f946cc69a10349b", size = 600845, upload-time = "2026-10-05T00:24:02.497Z" }, + { url = "https://files.pythonhosted.org/packages/63/80/6a03356d8ee38261e3a78cf89ee03d8e7f12c572d969237be00869e2dc73/wcwidth-0.9.2-cp315-cp315t-win_amd64.whl", hash = "sha256:33df042f96c61ed3cd5fb3742fba427553a635bc578799857a48aa79f774a0b9", size = 605441, upload-time = "2026-10-05T00:24:04.052Z" }, + { url = "https://files.pythonhosted.org/packages/0c/48/1a308a86a833fd12ff7a08d0d2491ff4a72c8a92d12f5ead8317630f771e/wcwidth-0.9.2-cp315-cp315t-win_arm64.whl", hash = "sha256:48719a9bc76c2f84238693fe5013571fa5beffa3621cf228f1f3a9e30dae84b8", size = 605885, upload-time = "2026-10-05T00:23:54.274Z" }, + { url = "https://files.pythonhosted.org/packages/9c/b4/0bfa065af506540d9d558e3e5548cff00bc1f9b24e6e2a8512498e8628de/wcwidth-0.9.2-py3-none-any.whl", hash = "sha256:89ca642c5bf0101157a09366be69fad0379db1f700ae39a920e103234573670e", size = 301667, upload-time = "2026-10-05T00:23:21.097Z" }, +] + [[package]] name = "websocket-client" version = "1.9.2" From 996b2d261a6360190c8b0a7d87e70513801df7ac Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 28 Oct 2026 18:41:42 +0000 Subject: [PATCH 019/278] Track, store, register and serve any bundle through internal seams Tracking builds each run record from a neutral RunSpec, and Trainer.fit adapts its context and result into one. One _store_and_register stores, logs and registers a bundle from its directory and ModelManifest. The manifest lists the model's dependencies: the predictor entry point that the MLflow pyfunc loads, the kumo-engine package, and content-addressed assets such as text-encoder word vectors, stored once under /assets/// and resolved with a sha256 check when the model loads. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/docs/mlops.md | 29 +- packages/kumo-engine/kumo/engine/_deps.py | 301 ++++++++++++++++++ .../kumo-engine/kumo/engine/_embeddings.py | 90 +++++- packages/kumo-engine/kumo/engine/_manifest.py | 35 +- .../kumo-engine/kumo/engine/_predictor.py | 65 ++++ packages/kumo-engine/kumo/engine/_tracking.py | 182 +++++------ .../kumo/engine/integrations/_mlflow_model.py | 8 +- .../kumo/engine/integrations/mlflow.py | 65 ++-- .../kumo/engine/trainer/_provenance.py | 6 +- .../kumo/engine/trainer/_run_spec.py | 56 ++++ .../kumo/engine/trainer/trainer.py | 177 ++++++---- .../test/test_bundle_provenance.py | 9 +- .../kumo-engine/test/test_mlflow_model.py | 42 +-- .../kumo-engine/test/test_mlflow_tracking.py | 7 +- .../test/test_model_agnostic_mlops.py | 269 ++++++++++++++++ .../test/test_snowflake_tracking.py | 8 +- .../test/test_tracking_contract.py | 26 +- 17 files changed, 1144 insertions(+), 231 deletions(-) create mode 100644 packages/kumo-engine/kumo/engine/_deps.py create mode 100644 packages/kumo-engine/kumo/engine/_predictor.py create mode 100644 packages/kumo-engine/kumo/engine/trainer/_run_spec.py create mode 100644 packages/kumo-engine/test/test_model_agnostic_mlops.py diff --git a/packages/kumo-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index 0a79573eb..061a9a253 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -51,6 +51,33 @@ With MLOps on, each bundle is stored under `/models//`. A compl The bundle also carries the published model's holdout predictions (`holdout.parquet`, with its truncation flag), so `result.holdout_df()` works on a model loaded with `Trainer.load` or `Trainer.load_from_connector` (a path or a registry), not only on the result of `fit`. Bundles fitted before this have no copy, and `holdout_df()` raises `FileNotFoundError` for them. +### Artifacts, bundles and dependencies + +- An **artifact** is anything with its own id and `_provenance.json`: a snapshot, a training table, a model, a prediction run, or an asset. +- A model's **bundle** holds only the model's own state: weights, encoders, the training spec and the holdout. +- Its manifest's **`parents`** are lineage: the training table and the warm-start model it came from. +- Its **`dependencies`** are what `predict` needs besides the bundle and a snapshot. The snapshot is an input to each call, never a dependency. Each entry has a `kind`: + - `predictor`: the entry point that scores the bundle; + - `package`: the `kumo-engine` version and its extras, with `[text]` only when the model encodes a text column; + - `asset`: a content-addressed file set, by `asset_kind`, `name`, `version` and `sha256`. +- Bundles fitted before dependencies were recorded list none. They load as before, with the built-in predictor. +- **Assets** are not copied into bundles. With MLOps on, each asset is stored once under `/assets///` with its own `_provenance.json`, and every model that uses it shares that copy. The dependency records where it was stored. +- **Resolving assets:** `Trainer.load`, the MLflow pyfunc and scoring jobs look for each asset at the location the dependency records, then under the artifact root of the enclosing `kumo.mlops` block, then in the embedding-model store. They check its sha256 and cache non-local copies under `~/.cache/kumo/assets//`. If an asset is stored but cannot be found, loading fails with `MissingDependencyError`, which names the asset kind, name and sha256. A bundle without stored assets resolves its models from the store, as before. +- **Text encoders** are the first assets (`asset_kind="text_encoder"`): the word vectors of each text column's embedding model. Word-vector models (`GloVeVectors`) are stored. Other `TokenEmbeddingModel`s are recorded by fingerprint and must be registered wherever the model is loaded. + +To train with your own word vectors and serve the model without registering them again, register them under a name, fit with MLOps on, and load the stored model anywhere that can reach the artifact root: + +```python +store = kumo.EmbeddingModelStore() +store.add("colors", kumo.GloVeVectors("/data/vectors")) # a .zip release or a directory of converted files +plan.column_processing.encoder_overrides = {"users.bio": TextEmbedding("colors")} +with kumo.mlops(kumo.MLOps(artifact_root="s3://bucket/kumo")): + result = kumo.Trainer(plan, embedding_models=store).fit(table, snapshot=snapshot) +# Elsewhere, with no store: the vectors come from /assets/text_encoder//. +trainer = kumo.Trainer.load_from_connector(kumo.S3Connector(root_dir="s3://bucket/kumo"), + path=f"models/{result.id}") +``` + ### Scoring defaults A model can carry the prediction arguments `predict` uses when a call omits them: @@ -290,7 +317,7 @@ scores = model.predict(pd.DataFrame({"ENTITY": [1, 2, 3]}), params={ - `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 once (`Could not log the model: ...` or `Could not register the model: ...`, with the full error) and still returns its result with `registered_model=None` (with `on_error="raise"`, it raises instead), the run ends `FAILED` with the error in `kumo.error`, 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, the registered version's tags and `result.registered_model.metadata` hold `kumo_bundle_id`. The model's `provenance` artifact is the bundle's `_provenance.json`. Its pip requirement is `kumo-engine[s3,gcs]` of the same version, plus `[text]` when the model encodes a text column. A model trained with a custom embedding model also needs that model's package in the serving environment. +- The model's metadata, the registered version's tags and `result.registered_model.metadata` hold `kumo_bundle_id`. The model's `provenance` artifact is the bundle's `_provenance.json`. Its pip requirement is `kumo-engine[s3,gcs]` of the same version, plus `[text]` when the model encodes a text column. The pyfunc scores with the predictor the bundle's manifest names. It resolves the model's assets, such as non-default word vectors, as described in [Artifacts, bundles and dependencies](#artifacts-bundles-and-dependencies). - 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=` (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. - Input rows have an `ENTITY` column and, optionally, a `TIMESTAMP` column. Scores come back in the input's row order, one per input row (repeated rows included); rows whose entity is not in the snapshot are left out. A temporal query scores entities at a time: pass `params={"prediction_time": ...}` or a `TIMESTAMP` column; without either, `predict` fails with `temporal query: pass params={'prediction_time': ...} or a TIMESTAMP column`. Binary classification needs `params={"binary_classification_threshold": ...}` unless the model's `scoring_defaults` set one. - Serving needs an image with kumo-engine already installed; MLflow's stock serving containers cannot install it. Serve with `mlflow models serve -m --env-manager local` (on SageMaker endpoints, `-m /opt/ml/model -p 8080`). SageMaker batch transform takes JSON input (`dataframe_split` plus `params`), not CSV, which carries no params. diff --git a/packages/kumo-engine/kumo/engine/_deps.py b/packages/kumo-engine/kumo/engine/_deps.py new file mode 100644 index 000000000..189614856 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/_deps.py @@ -0,0 +1,301 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""What a model needs to predict besides its bundle and a snapshot: its +dependencies, recorded at fit, and the assets among them, resolved at +predict time.""" + +import hashlib +import importlib +import os +import shutil +import tempfile +import warnings +from pathlib import Path +from typing import ( + Any, + List, + Mapping, + Optional, + Protocol, + Sequence, + Tuple, + Union, +) + +from kumo_core.artifacts import PROVENANCE_FILE + +from kumo.engine._manifest import ( + AssetManifest, + Dependency, + ModelManifest, + manifest_bytes, + read_manifest, + stamp, +) + +#: A place assets are stored: a URI or path, or a ``(connector, path)``. +Location = Union[str, Tuple[Any, str]] + + +class Asset(Protocol): + r"""A file or directory a model needs at predict time, addressed by its + content. ``files`` maps names to local files; the class's + ``load(directory, dependency)`` restores one from a directory holding + them, and must give the same ``fingerprint()``.""" + + asset_kind: str + name: str + version: Optional[str] + + def fingerprint(self) -> str: + ... + + def files(self) -> Mapping[str, Path]: + ... + + +class MissingDependencyError(RuntimeError): + r"""A model's dependency is in none of the places looked in.""" + + +class FileAsset: + r"""A plain file, or a directory of files.""" + def __init__(self, asset_kind: str, name: str, path: Union[str, Path], + version: Optional[str] = None) -> None: + self.asset_kind, self.name, self.version = asset_kind, name, version + self.path = Path(path) + + def files(self) -> Mapping[str, Path]: + if self.path.is_file(): + return {self.path.name: self.path} + return { + path.relative_to(self.path).as_posix(): path + for path in sorted(self.path.rglob('*')) + if path.is_file() and path.name != PROVENANCE_FILE + } + + def fingerprint(self) -> str: + digest = hashlib.sha256() + for name, path in sorted(self.files().items()): + digest.update(f'{name}\0{path.stat().st_size}\0'.encode()) + with open(path, 'rb') as handle: + for block in iter(lambda: handle.read(1 << 20), b''): + digest.update(block) + return digest.hexdigest() + + @classmethod + def load(cls, directory: Path, dependency: Dependency) -> 'FileAsset': + files = [ + path for path in directory.iterdir() + if path.is_file() and path.name != PROVENANCE_FILE + ] + return cls(str(dependency.asset_kind), dependency.name, + files[0] if len(files) == 1 else directory, + dependency.version) + + +def predictor(entry_point: str) -> Dependency: + return Dependency(kind='predictor', name=entry_point) + + +def package(name: str, version: str, extras: Sequence[str]) -> Dependency: + return Dependency(kind='package', name=name, version=version, + extras=list(extras)) + + +def entry_point(value: Any) -> str: + cls = type(value) + return f'{cls.__module__}:{cls.__qualname__}' + + +def imported(entry: str) -> Any: + module, _, name = entry.partition(':') + value: Any = importlib.import_module(module) + for part in name.split('.'): + value = getattr(value, part) + return value + + +class Dependencies: + r"""Records a fit's dependencies (``require``) and resolves a fitted + model's assets (``resolve``). With MLOps ``settings`` whose export + includes models, each required asset is stored once under the + artifact root; ``locations`` are asset roots looked in first.""" + def __init__(self, records: Sequence[Dependency] = (), + settings: Any = None, + locations: Sequence[Location] = ()) -> None: + self.records: List[Dependency] = list(records) + self._settings = settings + self._locations = list(locations) + + @classmethod + def of(cls, manifest: Optional[ModelManifest], settings: Any = None, + locations: Sequence[Location] = ()) -> 'Dependencies': + return cls(getattr(manifest, 'dependencies', None) or (), settings, + locations) + + def add(self, dependency: Dependency) -> Dependency: + self.records.append(dependency) + return dependency + + def require(self, asset: Asset) -> Dependency: + sha256 = asset.fingerprint() + location = None + settings = self._settings + if settings is not None and settings.wants('model'): + location = _store(asset, sha256, settings) + return self.add( + Dependency(kind='asset', name=asset.name, version=asset.version, + asset_kind=asset.asset_kind, sha256=sha256, + entry_point=entry_point(asset), packing='external', + location=location)) + + def assets(self, asset_kind: str) -> List[Dependency]: + return [ + record for record in self.records + if record.kind == 'asset' and record.asset_kind == asset_kind + ] + + def resolve(self, asset_kind: str, name: str) -> Path: + r"""A local directory holding the asset, checked against its + sha256: from the given locations, then where it was stored, then + the current artifact root.""" + record = next((record for record in self.assets(asset_kind) + if record.name == name), None) + if record is None: + raise MissingDependencyError( + f'The model has no {asset_kind} dependency named {name!r}.') + tried = [] + for connector, path in self.candidates(record): + tried.append(str(path)) + try: + directory = _fetch(connector, path, str(record.sha256)) + except (FileNotFoundError, OSError, ValueError): + continue + if directory is not None and _verified(directory, record): + return directory + raise MissingDependencyError( + f"The model needs {asset_kind} {name!r} (sha256 " + f"{record.sha256}), which is in none of " + f"{tried or 'no configured locations'}. Store it by fitting " + f"with MLOps on (it is kept under /assets/), " + f"predict inside a kumo.mlops block whose artifact root holds " + f"it, or register it where the model looks it up.") + + def candidates(self, record: Dependency) -> List[Tuple[Any, str]]: + from kumo.engine.connector import connector_for_root + + relative = f'{record.asset_kind}/{record.sha256}' + found: List[Tuple[Any, str]] = [] + for location in self._locations: + if isinstance(location, str): + found.append((connector_for_root(location), relative)) + else: + connector, path = location + found.append((connector, f"{path.rstrip('/')}/{relative}")) + if record.location: + parent, _, leaf = record.location.rstrip('/').rpartition('/') + try: + found.append((connector_for_root(parent or '/'), leaf)) + except ValueError: + pass + settings = self._settings + if settings is not None: + found.append((settings.root.connector, + settings.root.location(f'assets/{relative}'))) + return found + + +def asset_cache() -> Path: + base = os.environ.get('XDG_CACHE_HOME') or Path.home() / '.cache' + return Path(base) / 'kumo' / 'assets' + + +def _store(asset: Asset, sha256: str, settings: Any) -> Optional[str]: + r"""Writes ``asset`` once to ``/assets///``, + its ``_provenance.json`` last; returns where, or None (with a + warning) if it failed.""" + from kumo.engine.connector import file_store + + relative = f'assets/{asset.asset_kind}/{sha256}' + try: + filesystem, root = file_store(settings.root.connector, + settings.root.location(relative)) + root = root.rstrip('/') + marker = f'{root}/{PROVENANCE_FILE}' + if filesystem.get_file_info(marker).is_file: + return settings.root.uri(relative) + filesystem.create_dir(root, recursive=True) + for name, path in asset.files().items(): + with open(path, 'rb') as reader, \ + filesystem.open_output_stream(f'{root}/{name}') as writer: + shutil.copyfileobj(reader, writer, 1 << 20) + with filesystem.open_output_stream(marker) as writer: + writer.write(manifest_bytes(_manifest(asset, sha256))) + return settings.root.uri(relative) + except Exception as error: + warnings.warn( + f"Could not store {asset.asset_kind} {asset.name!r} under the " + f"artifact root: {error}", UserWarning, stacklevel=4) + return None + + +def _manifest(asset: Asset, sha256: str) -> AssetManifest: + return AssetManifest(id=sha256, asset_kind=asset.asset_kind, + name=asset.name, version=asset.version, + format=entry_point(asset), + source=str(getattr(asset, 'path', '') or '') or None, + **stamp()) + + +def _fetch(connector: Any, path: str, sha256: str) -> Optional[Path]: + from pyarrow.fs import FileSelector, FileType, LocalFileSystem + + from kumo.engine.connector import file_store + + store = file_store(connector, path) + if store is None: + return None + filesystem, root = store + root = root.rstrip('/') + if filesystem.get_file_info( + f'{root}/{PROVENANCE_FILE}').type != FileType.File: + return None + if isinstance(filesystem, LocalFileSystem): + return Path(root) + target = asset_cache() / sha256 + if (target / PROVENANCE_FILE).is_file(): + return target + target.parent.mkdir(parents=True, exist_ok=True) + partial = Path(tempfile.mkdtemp(prefix=f'.{sha256}-', dir=target.parent)) + try: + for info in filesystem.get_file_info(FileSelector(root, + recursive=True)): + if info.type != FileType.File: + continue + local = partial / info.path[len(root) + 1:] + local.parent.mkdir(parents=True, exist_ok=True) + with filesystem.open_input_stream(info.path) as reader, \ + open(local, 'wb') as writer: + shutil.copyfileobj(reader, writer, 1 << 20) + try: + partial.rename(target) + except OSError: + if not (target / PROVENANCE_FILE).is_file(): + raise + finally: + shutil.rmtree(partial, ignore_errors=True) + return target + + +def _verified(directory: Path, record: Dependency) -> bool: + stored = read_manifest(directory / PROVENANCE_FILE, AssetManifest) + if stored is None or stored.id != record.sha256: + return False + loader = imported(record.entry_point or entry_point(FileAsset)) + return bool(loader.load(directory, record).fingerprint() == record.sha256) + + +def loaded(directory: Path, record: Dependency) -> Any: + return imported(record.entry_point or '').load(directory, record) diff --git a/packages/kumo-engine/kumo/engine/_embeddings.py b/packages/kumo-engine/kumo/engine/_embeddings.py index 4916590d6..1a51a42c8 100644 --- a/packages/kumo-engine/kumo/engine/_embeddings.py +++ b/packages/kumo-engine/kumo/engine/_embeddings.py @@ -4,7 +4,8 @@ import logging from contextlib import contextmanager -from typing import Any, Iterator, Optional +from pathlib import Path +from typing import Any, Dict, Iterator, List, Optional from kumo_core import embedding_models from kumo_core.embedding_models import EmbeddingModelStore @@ -48,3 +49,90 @@ def require_trained_embedding_models(store: EmbeddingModelStore, # A trained model is matched by the identity of its models, not names. if encoders is not None: embedding_models.require_encoder_models(encoders, store, trained=True) + + +class TextEncoderAsset: + r"""The word vectors of a text encoder, as a model dependency.""" + asset_kind = 'text_encoder' + + def __init__(self, name: str, model: Any) -> None: + self.name, self.model = name, model + self.version = model.label + self.path = getattr(model, 'path', None) + + def fingerprint(self) -> str: + return str(self.model.fingerprint()) + + def files(self) -> Dict[str, Path]: + vocab, vectors = self.model._files() + return {vocab.name: vocab, vectors.name: vectors} + + @classmethod + def load(cls, directory: Path, dependency: Any) -> 'TextEncoderAsset': + from kumo_core.embedding_models import GloVeVectors + + return cls(dependency.name, GloVeVectors.in_directory(directory)[0]) + + +def _text_encoders(encoders: Any) -> List[Any]: + from kumoapi.encoder import TextEmbedding + + found: Dict[str, Any] = {} + for table in (encoders.encoders.values() if encoders else ()): + for encoder in table.values(): + if isinstance(encoder, TextEmbedding) and encoder._identity: + found.setdefault(encoder._identity['fingerprint'], encoder) + return list(found.values()) + + +def require_text_encoders(deps: Any, store: EmbeddingModelStore, + encoders: Any) -> bool: + r"""Records the word vectors of a fitted model's text encoders in + ``deps``; returns whether it has any.""" + from kumo_core.embedding_models import GloVeVectors, resolve_encoder + + from kumo.engine._manifest import Dependency + + text = _text_encoders(encoders) + for encoder in text: + model = resolve_encoder(encoder, store) + if isinstance(model, GloVeVectors): + deps.require(TextEncoderAsset(encoder.model, model)) + else: + deps.add( + Dependency(kind='asset', name=encoder.model, + version=model.label, asset_kind='text_encoder', + sha256=encoder._identity['fingerprint'], + packing='external')) + return bool(text) + + +def with_text_encoders(store: EmbeddingModelStore, encoders: Any, + deps: Any) -> EmbeddingModelStore: + r"""``store`` plus the stored word vectors of the model's text encoders, + resolved where ``deps`` looks for assets. Unless a place was given to + look, ``store`` alone serves, as before; a missing model it does not + have either is an error.""" + from kumo.engine._deps import MissingDependencyError, loaded + + added: Dict[str, Any] = {} + for encoder in _text_encoders(encoders): + known = encoder._identity + record = next((record for record in deps.assets('text_encoder') + if record.sha256 == known['fingerprint']), None) + if record is None or not record.entry_point or not deps.candidates( + record): + continue + try: + directory = deps.resolve('text_encoder', record.name) + except MissingDependencyError: + if store._find(known, encoder.model) is not None: + continue + raise + added[f"{record.name}@{record.sha256[:12]}"] = loaded( + directory, record).model + if not added: + return store + extended = EmbeddingModelStore() + extended._models = {**store._models, **added} + return extended diff --git a/packages/kumo-engine/kumo/engine/_manifest.py b/packages/kumo-engine/kumo/engine/_manifest.py index b0ee0b015..8741da702 100644 --- a/packages/kumo-engine/kumo/engine/_manifest.py +++ b/packages/kumo-engine/kumo/engine/_manifest.py @@ -318,11 +318,30 @@ class TrackingRef(Record): run_id: str +@dataclass(frozen=True, kw_only=True) +class Dependency(Record): + r"""One thing a model needs to predict: its ``predictor`` entry point, + a ``package`` requirement, or an ``asset`` (content-addressed by + ``sha256``, loaded by ``entry_point``, stored at ``location``).""" + + kind: str + name: str + version: Optional[str] = omit_none() + extras: Optional[List[str]] = omit_none() + asset_kind: Optional[str] = omit_none() + sha256: Optional[str] = omit_none() + entry_point: Optional[str] = omit_none() + packing: Optional[str] = omit_none() + location: Optional[str] = omit_none() + + @dataclass(frozen=True, kw_only=True) class ModelManifest(Manifest): """How a model bundle was fitted: from which training table (embedded, with its snapshot) and warm-start model, when, and with which fit - mode. Its ``id`` is the bundle id.""" + mode. Its ``id`` is the bundle id; ``dependencies`` is what predicting + needs besides the bundle and a snapshot (none listed: the built-in + `Trainer` predictor).""" kind: str = 'model' URI_TAG: ClassVar[Optional[str]] = 'kumo.bundle_uri' @@ -330,10 +349,24 @@ class ModelManifest(Manifest): fit: Optional[FitTiming] = None fit_mode: Optional[str] = None tracking: Optional[TrackingRef] = omit_none() + dependencies: Optional[List[Dependency]] = omit_none() EMBED_OMITS = ('tracking', ) +@dataclass(frozen=True, kw_only=True) +class AssetManifest(Manifest): + """An asset a model depends on, stored once under + ``/assets///``; its ``id`` is the sha256.""" + + kind: str = 'asset' + asset_kind: Optional[str] = None + name: Optional[str] = None + version: Optional[str] = None + format: Optional[str] = None + source: Optional[str] = None + + def manifest_bytes(manifest: Manifest) -> bytes: return (json.dumps(manifest.to_dict(), indent=2, sort_keys=True) + '\n').encode('utf-8') diff --git a/packages/kumo-engine/kumo/engine/_predictor.py b/packages/kumo-engine/kumo/engine/_predictor.py new file mode 100644 index 000000000..b65ebb906 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/_predictor.py @@ -0,0 +1,65 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""What scores a model bundle: the ``predictor`` dependency its manifest +names.""" + +from pathlib import Path +from typing import Any, Mapping, Protocol, Tuple, Union + +import pandas as pd + +from kumo.engine._deps import imported + +#: The predictor of bundles whose manifest names none. +TRAINER_PREDICTOR = 'kumo.engine._predictor:TrainerPredictor' + + +class Predictor(Protocol): + r"""A loaded bundle: scores ``rows`` (an ``ENTITY`` frame, or a + `PredictionTable` for files) against ``snapshot``.""" + def predict(self, rows: Any, snapshot: Any, + params: Mapping[str, Any]) -> Union[pd.DataFrame, + Tuple[str, ...]]: + ... + + +def load_predictor(bundle_dir: Union[str, Path]) -> Predictor: + r"""The predictor the bundle at ``bundle_dir`` names, loaded with + ``.load(bundle_dir)``.""" + from kumo.engine.trainer._provenance import read_provenance + + manifest = read_provenance(Path(bundle_dir)) + entry = next((record.name + for record in getattr(manifest, 'dependencies', None) or () + if record.kind == 'predictor'), TRAINER_PREDICTOR) + try: + predictor = imported(entry) + except (ImportError, AttributeError, ValueError) as error: + raise ValueError( + f"The model bundle at {bundle_dir} names predictor {entry!r}, " + f"which cannot be imported here: {error}") from error + return predictor.load(bundle_dir) + + +class TrainerPredictor: + r"""Scores with the `Trainer` the bundle restores.""" + def __init__(self, trainer: Any) -> None: + self.trainer = trainer + + @classmethod + def load(cls, bundle_dir: Union[str, Path]) -> 'TrainerPredictor': + from kumo.engine.trainer.trainer import Trainer + + return cls(Trainer.load(bundle_dir)) + + def predict(self, rows: Any, snapshot: Any, + params: Mapping[str, Any]) -> Union[pd.DataFrame, + Tuple[str, ...]]: + from kumo.engine.integrations._mlflow_model import predict + from kumo.engine.pquery.prediction_table import PredictionTable + + if isinstance(rows, PredictionTable): + return self.trainer.predict(snapshot, prediction_table=rows, + **params).files + return predict(self.trainer, snapshot, rows, params) diff --git a/packages/kumo-engine/kumo/engine/_tracking.py b/packages/kumo-engine/kumo/engine/_tracking.py index 4160e97aa..562059341 100644 --- a/packages/kumo-engine/kumo/engine/_tracking.py +++ b/packages/kumo-engine/kumo/engine/_tracking.py @@ -27,7 +27,6 @@ Tuple, ) -from kumo_core.artifacts import PROVENANCE_FILE from kumo_core.instrumentation import format_error from kumo.engine._events import Event, _finite_floats @@ -36,8 +35,7 @@ if TYPE_CHECKING: from kumoapi.plugin import ModelRegistryLocation, PublishedModel - from kumo.engine.graph import GraphSnapshot - from kumo.engine.trainer.trainer import TrainingResult + from kumo.engine._manifest import ModelManifest logger = logging.getLogger(__name__) @@ -50,17 +48,6 @@ MODEL_SECONDS = 3600.0 -@dataclass(frozen=True) -class FitContext: - r"""What a fit trains, known once its model plan is resolved.""" - - run_id: str - query: str - model_plan: Mapping[str, Any] - inputs: Mapping[str, Mapping[str, Optional[str]]] = field( - default_factory=dict) - - @dataclass(frozen=True) class RunRecord: r"""One tracked run; ``steps`` are per-epoch ``(epoch, value)`` @@ -79,6 +66,28 @@ class RunRecord: default_factory=dict) +@dataclass(frozen=True) +class RunSpec: + r"""What a run records, whatever trained or scored it: ``params`` are + flat, ``parents`` become ``kumo.`` pointer tags and ``tags`` are + the user's own.""" + + kind: str + run_id: str + name: Optional[str] = None + params: Mapping[str, Any] = field(default_factory=dict) + metrics: Mapping[str, Any] = field(default_factory=dict) + steps: Mapping[str, Sequence[Tuple[int, float]]] = field( + default_factory=dict) + tags: Mapping[str, Any] = field(default_factory=dict) + inputs: Mapping[str, Mapping[str, Optional[str]]] = field( + default_factory=dict) + parents: Mapping[str, Optional[str]] = field(default_factory=dict) + artifacts: Mapping[str, str] = field(default_factory=dict) + status: Optional[str] = None + error: Optional[str] = None + + INPUT_FIELDS = ('connector', 'location', 'version', 'version_kind') @@ -116,43 +125,25 @@ def _text(value: Any) -> Optional[str]: return None if value is None else str(value) -def run_record(context: FitContext, result: Optional['TrainingResult'] = None, - fit_ended: Optional[Event] = None, *, - tags: Optional[Mapping[str, str]] = None) -> RunRecord: - r"""The fit run's record; ``tags`` are the user's extra tags.""" - metrics: Dict[str, float] = {} - pointers: Dict[str, str] = {} - artifacts = { - 'query.pql': context.query, - 'model_plan.json': json.dumps(dict(context.model_plan), indent=2, - default=str), - } - if result is not None: - metrics.update(result.metrics) - if fit_ended is not None and 'seconds' in fit_ended.metrics: - metrics['seconds'] = fit_ended.metrics['seconds'] - if result.id: - pointers['kumo.bundle_id'] = str(result.id) - provenance = Path(result.bundle_dir) / PROVENANCE_FILE - if provenance.is_file(): - artifacts[PROVENANCE_FILE] = provenance.read_text( - encoding='utf-8') - params = { - key: _param_value(value) - for key, value in _flatten(context.model_plan) - } +def run_record(spec: RunSpec) -> RunRecord: + params = {key: _param_value(value) for key, value in spec.params.items()} return RunRecord( + name=spec.name, params=params, - metrics=_finite_floats(metrics), + metrics=_finite_floats(spec.metrics), + steps={name: tuple(history) + for name, history in spec.steps.items()}, tags={ - **_user_tags(tags, params), - 'kumo.kind': 'fit', - 'kumo.run_id': context.run_id, - **pointers, - **_error_tag(fit_ended), + **_user_tags(spec.tags, params), + 'kumo.kind': spec.kind, + 'kumo.run_id': spec.run_id, + **{f'kumo.{name}': str(value) + for name, value in spec.parents.items() if value}, + **({} if spec.status is None else {'kumo.status': spec.status}), + **({} if spec.error is None else {'kumo.error': spec.error}), }, - artifacts=artifacts, - inputs=dict(context.inputs), + artifacts=dict(spec.artifacts), + inputs=dict(spec.inputs), ) @@ -170,50 +161,34 @@ def _user_tags(tags: Optional[Mapping[str, Any]], } -def _error_tag(ended: Optional[Event]) -> Dict[str, str]: - if ended is None or 'error' not in ended.params: - return {} - return {'kumo.error': str(ended.params['error'])} - - -def trial_record(context: FitContext, events: Sequence[Event]) -> RunRecord: - r"""One trial attempt's record from its events so far.""" +def trial_spec(run_id: str, events: Sequence[Event]) -> RunSpec: + r"""One trial attempt of run ``run_id``, from its events so far.""" started = next(event for event in events if event.name == 'trial_started') trial, attempt = started.trial, int(started.params['attempt']) - params: Dict[str, str] = { - key: _param_value(value) - for key, value in started.params.items() - } metrics: Dict[str, float] = {} steps: Dict[str, List[Tuple[int, float]]] = {} - ended: Dict[str, str] = {} + ended: Dict[str, Any] = {} for event in events: if event.name == 'epoch_ended' and event.epoch is not None: for name, value in event.metrics.items(): steps.setdefault(name, []).append((event.epoch, value)) elif event.name == 'trial_ended': - ended = { - 'kumo.status': str(event.params.get('status')), - **_error_tag(event), - } + ended = dict(status=str(event.params.get('status')), + error=None if 'error' not in event.params else str( + event.params['error'])) metrics.update({ name if name == 'seconds' else f'best_{name}': value for name, value in event.metrics.items() }) - return RunRecord( + return RunSpec( + kind='trial', + run_id=f'{run_id}-trial-{trial}-attempt-{attempt}', name=f'trial-{trial}-attempt-{attempt}', - params=params, + params=dict(started.params), metrics=metrics, - steps={ - name: tuple(history) - for name, history in steps.items() - }, - tags={ - 'kumo.kind': 'trial', - 'kumo.run_id': f'{context.run_id}-trial-{trial}-attempt-{attempt}', - 'kumo.parent_run': context.run_id, - **ended, - }, + steps=steps, + parents={'parent_run': run_id}, + **ended, ) @@ -225,11 +200,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); ``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``. + ``tags`` (extra fit-run tags); ``log_model(bundle_dir, manifest, + 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``. """ def start(self, record: RunRecord) -> None: ... @@ -252,9 +228,8 @@ def __init__(self, tracking: Optional[Tracking], warn: bool = True) -> None: self._tracker = tracking if _rank() == 0 else None self._warn = warn - self._context: Optional[FitContext] = None + self.spec: Optional[RunSpec] = None self._trials: Dict[int, List[Event]] = {} - self._fit_ended: Optional[Event] = None self._jobs: 'queue.Queue[Optional[Callable[[], Any]]]' = queue.Queue( QUEUE_SIZE) self._worker: Optional[threading.Thread] = None @@ -265,31 +240,29 @@ def __init__(self, tracking: Optional[Tracking], def enabled(self) -> bool: return self._tracker is not None - def start(self, context: Callable[[], FitContext]) -> None: + def start(self, spec: Callable[[], RunSpec]) -> None: if self._tracker is None: return try: - value = context() + value = spec() except Exception as error: logger.warning('Experiment tracking failed: %s; the fit continues', full_error(error), exc_info=_debug()) self._fail(error) self._tracker = None return - self._context = value + self.spec = value self._call(lambda tracker: getattr(tracker, 'prepare', _none)()) self._worker = threading.Thread(target=self._work, name='kumo-tracking', daemon=True) self._worker.start() self._submit(lambda: self._call(lambda tracker: tracker.start( - run_record(value, tags=getattr(tracker, 'tags', None))))) + run_record(dataclasses.replace(value, tags={ + **value.tags, **(getattr(tracker, 'tags', None) or {}) + }))))) def on_event(self, event: Event) -> None: - if self._context is None: - return - if event.name == 'fit_ended': - self._fit_ended = event - if event.trial is None: + if self.spec is None or event.trial is None: return if event.name == 'trial_started': self._trials[event.trial] = [] @@ -299,14 +272,15 @@ def on_event(self, event: Event) -> None: **event.params, 'attempt': events[0].params['attempt'] }) events.append(event) - record = trial_record(self._context, events) + record = run_record(trial_spec(self.spec.run_id, events)) self._submit(lambda: self._call( lambda tracker: tracker.on_trial(record, event))) def log_model( self, - result: 'TrainingResult', - snapshot: 'GraphSnapshot', + bundle_dir: Path, + manifest: Optional['ModelManifest'], + snapshot: Any, registry: Optional['ModelRegistryLocation'] = None, ) -> Tuple[Optional['PublishedModel'], Optional[BaseException]]: r"""Runs the tracker's ``log_model`` on the tracking thread and waits @@ -327,7 +301,8 @@ def log() -> None: '' if self.error is None else f' ({full_error(self.error)})')) registered.append( - getattr(tracker, 'log_model')(result, snapshot, registry)) + getattr(tracker, 'log_model')(bundle_dir, manifest, + snapshot, registry)) except BaseException as error: failed.append(error) finally: @@ -363,17 +338,18 @@ def log() -> None: return None, error return (registered or [None])[0], None - def finish(self, result: Optional['TrainingResult'], + def finish(self, spec: Optional[Callable[[], RunSpec]], error: Optional[BaseException]) -> None: + r"""Ends the run with ``spec()``, its final spec (`None`: the started + one), failed if ``error``.""" worker = self._worker - if self._context is None or worker is None: + if self.spec is None or worker is None: return try: - record: Optional[RunRecord] = run_record(self._context, result, - self._fit_ended) - if error is not None and 'kumo.error' not in record.tags: - record = dataclasses.replace(record, tags={ - **record.tags, 'kumo.error': format_error(error)}) + final = self.spec if spec is None else spec() + if error is not None and final.error is None: + final = dataclasses.replace(final, error=format_error(error)) + record: Optional[RunRecord] = run_record(final) except Exception as failure: logger.warning('Could not build the tracking record: %s', full_error(failure), exc_info=_debug()) diff --git a/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py b/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py index da763012f..a1be11e94 100644 --- a/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py +++ b/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py @@ -29,9 +29,9 @@ class KumoModel(_pyfunc.PythonModel): def load_context(self, context: Any) -> None: - from kumo.engine.trainer.trainer import Trainer + from kumo.engine._predictor import load_predictor - self._trainer = Trainer.load(context.artifacts['bundle']) + self._predictor = load_predictor(context.artifacts['bundle']) self._columns = list(context.model_config['columns']) self._snapshots: Dict[str, GraphSnapshot] = {} self._default = context.model_config.get('snapshot') @@ -47,8 +47,8 @@ def predict(self, context: Any, model_input: pd.DataFrame, if not uri: raise ValueError("Pass params={'snapshot': }.") - return predict(self._trainer, self._snapshot(uri), model_input, - params) + return self._predictor.predict(model_input, self._snapshot(uri), + params) def _snapshot(self, uri: str) -> GraphSnapshot: if uri not in self._snapshots: diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index 0eee9e20b..e5020d97c 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -50,8 +50,9 @@ if TYPE_CHECKING: from kumo_core.embedding_models import EmbeddingModelStore + from kumo.engine._manifest import ModelManifest from kumo.engine.graph import GraphSnapshot - from kumo.engine.trainer.trainer import Trainer, TrainingResult + from kumo.engine.trainer.trainer import Trainer logger = logging.getLogger(__name__) @@ -262,7 +263,8 @@ def on_trial(self, record: RunRecord, event: Event) -> None: def log_model( self, - result: 'TrainingResult', + bundle_dir: Path, + manifest: Optional['ModelManifest'], snapshot: 'GraphSnapshot', registry: Optional[ModelRegistryLocation] = None, ) -> Optional[PublishedModel]: @@ -277,8 +279,8 @@ def log_model( } os.environ.update(upload) try: - info = _log_model(self.target, self._client, result, snapshot, - name='model', run_id=run_id, + info = _log_model(self.target, self._client, bundle_dir, manifest, + snapshot, name='model', run_id=run_id, experiment_id=self._experiment_id, serving_snapshot=self.serving_snapshot) finally: @@ -293,7 +295,7 @@ def log_model( return None try: return _register(self.target, self._client, info, registry, - registry_pointers(result)) + registry_pointers(_id(manifest))) except Exception as error: raise RegistrationError(full_error(error)) from error @@ -508,6 +510,7 @@ def log_model( ``serving_snapshot`` (an exported snapshot URI), each serving replica downloads that snapshot once at startup and scores against it when no ``snapshot`` param is given.""" + from kumo.engine._manifest import stub from kumo.engine._tracking import fit_inputs mlflow = import_provider('mlflow', 'mlflow') @@ -532,8 +535,10 @@ def log_model( tracking._default_experiment or '0').info.run_id experiment_id = client.get_run(run_id).info.experiment_id _log_inputs(mlflow, client, run_id, fit_inputs(result.manifest)) - info = _log_model(target, client, result, snapshot, name=name, - run_id=run_id, experiment_id=experiment_id, + info = _log_model(target, client, result.bundle_dir, result.manifest or + (None if result.id is None else stub('model', result.id)), + snapshot, name=name, run_id=run_id, + experiment_id=experiment_id, serving_snapshot=serving_snapshot) with target.scope(): if result.metrics: @@ -594,25 +599,35 @@ def _download(mlflow: Any, target: Any, model_uri: str, tracking_uri=target.tracking_uri, registry_uri=target.registry_uri) -def _log_model(target: Any, client: Any, result: 'TrainingResult', +def _log_model(target: Any, client: Any, bundle_dir: Path, + manifest: Optional['ModelManifest'], snapshot: Optional['GraphSnapshot'], *, name: str, run_id: str, experiment_id: Optional[str], serving_snapshot: Optional[str]) -> Any: r"""What ``mlflow.pyfunc.log_model`` does, through ``client``: it only - logs to the process-wide tracking URI.""" + logs to the process-wide tracking URI. A bundle without a training spec + is logged without a signature.""" from kumo_core.artifacts import PROVENANCE_FILE from kumo.engine.integrations import _mlflow_model, _mlflow_model_code - from kumo.engine.trainer._spec import read_training_spec + from kumo.engine.trainer._spec import ( + SPEC_FILE, + listed_file, + read_training_spec, + ) from kumo.engine.trainer.trainer import registry_pointers mlflow = import_provider('mlflow', 'mlflow') if snapshot is None: raise ValueError('log_model needs the snapshot the model trained on') - spec = read_training_spec(result.bundle_dir) - signature, example = _mlflow_model.signature(snapshot, - spec.predictive_query) - provenance = result.bundle_dir / PROVENANCE_FILE + bundle_dir = Path(bundle_dir) + signature, example = None, None + if (bundle_dir / SPEC_FILE).is_file(): + signature, example = _mlflow_model.signature( + snapshot, + read_training_spec(bundle_dir).predictive_query) + pointers = registry_pointers(_id(manifest)) + provenance = bundle_dir / PROVENANCE_FILE with target.scope(): model = client.create_logged_model( experiment_id or client.get_run(run_id).info.experiment_id, @@ -620,8 +635,7 @@ def _log_model(target: Any, client: Any, result: 'TrainingResult', params=dict(client.get_run(run_id).data.params)) mlflow_model = mlflow.models.Model( artifact_path=model.artifact_location, model_uuid=model.model_id, - run_id=run_id, metadata=registry_pointers(result), - model_id=model.model_id) + run_id=run_id, metadata=pointers, model_id=model.model_id) model_logger = logging.getLogger('mlflow.models.model') model_logger.addFilter(_no_serving_check) try: @@ -631,18 +645,21 @@ def _log_model(target: Any, client: Any, result: 'TrainingResult', local, mlflow_model=mlflow_model, python_model=_mlflow_model_code.__file__, model_config={ - 'columns': signature.outputs.input_names(), + 'columns': [] if signature is None else + signature.outputs.input_names(), 'snapshot': serving_snapshot, }, artifacts={ - 'bundle': str(result.bundle_dir), + 'bundle': str(bundle_dir), **({'provenance': str(provenance)} if provenance.is_file() else {}), }, signature=signature, input_example=example, - metadata=registry_pointers(result), - pip_requirements=[_requirement(result)], + metadata=pointers, + pip_requirements=[ + _requirement(listed_file(bundle_dir, 'encoders')) + ], ) with target.scope(): repository = target.artifact_repository( @@ -665,13 +682,17 @@ def _log_model(target: Any, client: Any, result: 'TrainingResult', return mlflow_model.get_model_info(logged) -def _requirement(result: 'TrainingResult') -> str: +def _id(manifest: Optional['ModelManifest']) -> Optional[str]: + return None if manifest is None else manifest.id + + +def _requirement(encoders: Optional[Path]) -> str: from kumo_core.encoders.config import EncodersConfig from kumo_core.io.json import load_json from kumoapi.encoder import TextEmbedding extras = ['s3', 'gcs'] - path = result.artifact_paths.get('encoders') + path = None if encoders is None else str(encoders) config = (load_json(path, EncodersConfig) if path and Path(path).is_file() else None) if config is not None and any( diff --git a/packages/kumo-engine/kumo/engine/trainer/_provenance.py b/packages/kumo-engine/kumo/engine/trainer/_provenance.py index dfd6f151e..73e54e499 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_provenance.py +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -6,9 +6,10 @@ from dataclasses import replace from datetime import datetime, timezone from pathlib import Path -from typing import Optional +from typing import List, Optional from kumo.engine._manifest import ( + Dependency, FitTiming, ModelManifest, ModelParents, @@ -32,6 +33,7 @@ def write_provenance( started_at: datetime, fit_mode: Optional[str], tracked_run_id: Optional[str] = None, + dependencies: Optional[List[Dependency]] = None, ) -> None: from kumo_core.artifacts import PROVENANCE_FILE, add_bundle_file @@ -54,6 +56,7 @@ def write_provenance( fit_mode=fit_mode, tracking=None if tracked_run_id is None else TrackingRef( run_id=tracked_run_id), + dependencies=dependencies, **stamp(), )) add_bundle_file(bundle_dir, PROVENANCE_ROLE, PROVENANCE_FILE) @@ -109,6 +112,7 @@ def record_derivation(bundle_dir: Path, parent_dir: Path, parents=ModelParents( derived_from=None if parent is None else embedded_model( parent)), + dependencies=None if parent is None else parent.dependencies, **stamp(), )) add_bundle_file(bundle_dir, PROVENANCE_ROLE, PROVENANCE_FILE) diff --git a/packages/kumo-engine/kumo/engine/trainer/_run_spec.py b/packages/kumo-engine/kumo/engine/trainer/_run_spec.py new file mode 100644 index 000000000..7819797e2 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/trainer/_run_spec.py @@ -0,0 +1,56 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""A `Trainer.fit` as a tracked run's `RunSpec`.""" + +import json +from dataclasses import dataclass, field, replace +from pathlib import Path +from typing import Any, Mapping, Optional + +from kumo_core.artifacts import PROVENANCE_FILE + +from kumo.engine._tracking import RunSpec, _flatten + + +@dataclass(frozen=True) +class FitContext: + r"""What a fit trains, known once its model plan is resolved.""" + + run_id: str + query: str + model_plan: Mapping[str, Any] + inputs: Mapping[str, Mapping[str, Optional[str]]] = field( + default_factory=dict) + + +def fit_spec(context: FitContext, *, + tags: Optional[Mapping[str, Any]] = None) -> RunSpec: + return RunSpec( + kind='fit', + run_id=context.run_id, + params=dict(_flatten(context.model_plan)), + tags=dict(tags or {}), + inputs=dict(context.inputs), + artifacts={ + 'query.pql': context.query, + 'model_plan.json': json.dumps(dict(context.model_plan), indent=2, + default=str), + }, + ) + + +def fitted_spec(spec: RunSpec, result: Any, + seconds: Optional[float] = None) -> RunSpec: + r"""``spec`` once the fit returned ``result``, a `TrainingResult`, + after ``seconds``.""" + artifacts = dict(spec.artifacts) + provenance = Path(result.bundle_dir) / PROVENANCE_FILE + if provenance.is_file(): + artifacts[PROVENANCE_FILE] = provenance.read_text(encoding='utf-8') + return replace( + spec, metrics={ + **result.metrics, **({} if seconds is None else { + 'seconds': seconds + }) + }, parents={'bundle_id': result.id}, artifacts=artifacts) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 386247183..594303214 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -56,22 +56,28 @@ from kumoapi.train import TrainingStage from typing_extensions import Self +from kumo.engine import _deps as dependency +from kumo.engine._deps import Dependencies from kumo.engine._embeddings import ( no_encoder_inference_warnings, require_embedding_models, + require_text_encoders, require_trained_embedding_models, + with_text_encoders, ) from kumo.engine._events import Event, EventEmitter, query_sha256 -from kumo.engine._manifest import ModelManifest +from kumo.engine._manifest import Dependency, ModelManifest, stub from kumo.engine._mlops import full_error, warn_failure +from kumo.engine._predictor import TRAINER_PREDICTOR from kumo.engine._scratch import new_run_id, trainer_scratch_dir -from kumo.engine._tracking import FitContext, TrackingSession, fit_inputs +from kumo.engine._tracking import RunSpec, TrackingSession, fit_inputs from kumo.engine.connector import model_sink from kumo.engine.connector.base import Connector from kumo.engine.connector.source_table import SourceTable from kumo.engine.graph import GraphSnapshot, report_other_graph from kumo.engine.pquery.prediction_table import PredictionTable from kumo.engine.pquery.training_table import TrainingTable +from kumo.engine.trainer._run_spec import FitContext, fit_spec, fitted_spec from kumo.engine.trainer._scoring import ScoringDefaults, read_scoring_defaults from kumo.engine.trainer._spec import ( HOLDOUT_META_ROLE, @@ -146,7 +152,7 @@ def predict(self, *args: Any, **kwargs: Any) -> 'PredictionResult': def export(self, config: ModelOutputConfig) -> PublishedModel: """Publish this model's bundle, as :meth:`Trainer.export`.""" _require_publisher(config) - return _publish(self, config) + return _publish(self.bundle_dir, self.bundle_id, self.metrics, config) def with_scoring_defaults(self, defaults: ScoringDefaults) -> 'TrainingResult': @@ -328,7 +334,41 @@ def _listed_holdout(bundle_dir: Path) -> tuple[Optional[Path], Optional[bool]]: return path, truncated -def _store_model(result: TrainingResult, settings: Any) -> None: +def _store_and_register( + bundle_dir: Path, + manifest: Optional[ModelManifest], + settings: Any, + *, + session: Optional[TrackingSession] = None, + snapshot: Any = None, + metrics: Mapping[str, float] = {}, + warn: bool = True, +) -> Tuple[Optional[PublishedModel], Optional[BaseException]]: + r"""Store a fitted bundle under the artifact root, log it in the + session's run and register it, as MLOps ``settings`` say; returns the + registered version and the first failure.""" + if settings is not None and settings.wants('model'): + _store_model(bundle_dir, manifest, settings) + registered, error = (None, None) if session is None else session.log_model( + bundle_dir, manifest, snapshot, _mlflow_registry(settings)) + registry = _connector_registry(settings) + if registry is None: + return registered, error + errors: List[BaseException] = [] + try: + registered = _publish(bundle_dir, None if manifest is None else + manifest.id, metrics, registry, errors, warn) + except Exception as failure: + if warn: + warn_failure( + f"Could not register the model: {full_error(failure)}", + failure, 3) + return registered, error or failure + return registered, error or (errors or [None])[0] + + +def _store_model(bundle_dir: Path, manifest: Optional[ModelManifest], + settings: Any) -> None: r"""Copy the bundle to ``/models//`` unless a complete copy is there; failures warn.""" from kumo_core.object_store_sink import holds_bundle, upload_bundle @@ -336,21 +376,21 @@ def _store_model(result: TrainingResult, settings: Any) -> None: from kumo.engine.connector import file_store - if result.bundle_id is None: + if manifest is None: return try: filesystem, root = file_store( settings.root.connector, - settings.root.location(f'models/{result.bundle_id}')) - if holds_bundle(filesystem, root, str(result.bundle_dir)): + settings.root.location(f'models/{manifest.id}')) + if holds_bundle(filesystem, root, str(bundle_dir)): return if filesystem.get_file_info( FileSelector(root, allow_not_found=True)): filesystem.delete_dir_contents(root, missing_dir_ok=True) - upload_bundle(filesystem, root, str(result.bundle_dir)) + upload_bundle(filesystem, root, str(bundle_dir)) except Exception as error: warnings.warn( - f"Could not store model {result.bundle_id} under the artifact " + f"Could not store model {manifest.id} under the artifact " f"root: {error}", UserWarning, stacklevel=4) @@ -409,16 +449,18 @@ def _connector_registry(settings: Any) -> Optional[ModelOutputConfig]: return registry if isinstance(registry, ModelOutputConfig) else None -def _publish(training_result: TrainingResult, - export_config: ModelOutputConfig, +def _publish(bundle_dir: Path, bundle_id: Optional[str], + metrics: Mapping[str, float], export_config: ModelOutputConfig, errors: Optional[List[BaseException]] = None, warn: bool = True) -> PublishedModel: - spec = read_training_spec(training_result.bundle_dir) - metadata = { - name: str(getattr(spec, name)) - for name in ('spec_version', 'predictive_query', 'snapshot_id') - } - metadata.update(registry_pointers(training_result)) + metadata = {} + if (bundle_dir / SPEC_FILE).is_file(): + spec = read_training_spec(bundle_dir) + metadata = { + name: str(getattr(spec, name)) + for name in ('spec_version', 'predictive_query', 'snapshot_id') + } + metadata.update(registry_pointers(bundle_id)) location = export_config.location alias = location.alias if location else None if location and alias: @@ -426,11 +468,11 @@ def _publish(training_result: TrainingResult, sink = model_sink(export_config.connector, export_config.path) published = sink.publish_model( export_config.connector, - str(training_result.bundle_dir), + str(bundle_dir), path=export_config.path, location=location, metadata=metadata, - metrics=training_result.metrics, + metrics=metrics, ) if location and alias: try: @@ -547,7 +589,8 @@ def export(self, config: ModelOutputConfig) -> PublishedModel: """Publish the fitted or loaded model bundle.""" training_result = self._fitted('export') _require_publisher(config) - return _publish(training_result, config) + return _publish(training_result.bundle_dir, training_result.bundle_id, + training_result.metrics, config) @_without_torch_scatter_hint def fit( @@ -625,8 +668,8 @@ def fit( result = self._fit(train_table, snapshot=snapshot, export_config=export_config, warm_start=warm_start, run_id=run_id, - events=events, settings=settings, - tracking=session) + events=events, tracking=session, + deps=Dependencies(settings=settings)) except BaseException as error: events.emit('fit_ended', params={ 'status': 'failed', @@ -634,38 +677,25 @@ def fit( }) session.finish(None, error) raise + seconds = time.monotonic() - fit_start events.emit('fit_ended', metrics={ - **result.metrics, 'seconds': time.monotonic() - fit_start + **result.metrics, 'seconds': seconds }, params={'status': 'succeeded'}) - registered, error = session.log_model(result, snapshot, - _mlflow_registry(settings)) + registered, error = _store_and_register( + result.bundle_dir, result.manifest or ( + None if result.bundle_id is None else stub( + 'model', result.bundle_id)), settings, session=session, + snapshot=snapshot, metrics=result.metrics, warn=warn) if registered is not None: result = replace(result, registered_model=registered) self._training_result = result - if registry is not None: - result, failure = self._register(result, registry, warn) - error = error or failure - session.finish(result, error) + session.finish( + lambda: fitted_spec(cast(RunSpec, session.spec), result, seconds), + error) raise_on_error(settings, session.error or error, 'Trainer.fit (its model is in trainer.result)') return result - def _register( - self, result: TrainingResult, registry: ModelOutputConfig, warn: bool - ) -> Tuple[TrainingResult, Optional[BaseException]]: - errors: List[BaseException] = [] - try: - registered = _publish(result, registry, errors, warn) - except Exception as error: - if warn: - warn_failure( - f"Could not register the model: {full_error(error)}", - error, 3) - return result, error - result = replace(result, registered_model=registered) - self._training_result = result - return result, (errors or [None])[0] - def _fit( self, train_table: TrainingTable, @@ -675,8 +705,8 @@ def _fit( warm_start: Optional[Self], run_id: str, events: EventEmitter, - settings: Optional[Any], tracking: TrackingSession, + deps: Dependencies, ) -> TrainingResult: from kumo_core.automl.runtime import ( inherit_warm_start_model_plan, @@ -811,12 +841,12 @@ def _fit( and not materialized.task.task_type.is_link_pred): raise ValueError('The test split is empty: shorten the query ' 'window or add data.') - tracking.start(lambda: FitContext( + tracking.start(lambda: fit_spec(FitContext( run_id=run_id, query=train_table.query, model_plan=to_json_dict(materialized.model_plan), inputs=fit_inputs(train_table.manifest), - )) + ))) training_input = ResolvedTrainingInput( task=materialized.task, model_plan=materialized.model_plan, @@ -883,7 +913,7 @@ def _fit( warm_start_bundle_id), started_at=started_at, fit_mode=fit_mode, tracked_run_id=run_id if tracking.enabled else None), - settings=settings, + deps=deps, ) def _finish_fit( @@ -897,7 +927,7 @@ def _finish_fit( graph_config: Optional[Dict[str, Any]] = None, model_plan: Optional[ModelPlan] = None, provenance: Optional[Mapping[str, Any]] = None, - settings: Optional[Any] = None, + deps: Optional[Dependencies] = None, ) -> TrainingResult: """Commit durable fit outputs before making this trainer usable.""" from kumo_core.artifacts import bundle_id_from_manifest @@ -923,19 +953,21 @@ def _finish_fit( UserWarning, stacklevel=3) if provenance is not None and training_result.bundle_id is not None: try: - write_provenance(training_result.bundle_dir, - bundle_id=training_result.bundle_id, - **provenance) + write_provenance( + training_result.bundle_dir, + bundle_id=training_result.bundle_id, + dependencies=self._dependencies(training_result, deps + or Dependencies()), + **provenance) except Exception as error: warnings.warn( f"Publishing the model bundle without _provenance.json: " f"{error}", UserWarning, stacklevel=3) if export_config is not None: training_result = replace( - training_result, - published_model=_publish(training_result, export_config)) - if settings is not None and settings.wants('model'): - _store_model(training_result, settings) + training_result, published_model=_publish( + training_result.bundle_dir, training_result.bundle_id, + training_result.metrics, export_config)) training_result = replace(training_result, _trainer=self) self._training_result = training_result @@ -946,6 +978,18 @@ def _finish_fit( return training_result + def _dependencies(self, result: TrainingResult, + deps: Dependencies) -> List[Dependency]: + from kumo.engine._version import __version__ + + deps.add(dependency.predictor(TRAINER_PREDICTOR)) + text = require_text_encoders(deps, self._embedding_models, + _encoders(result.bundle_dir)) + deps.add( + dependency.package('kumo-engine', __version__, + ['text'] if text else [])) + return deps.records + @_without_torch_scatter_hint def predict( self, @@ -1314,6 +1358,12 @@ def _load_bundle( trainer._training_query = spec.predictive_query trainer._training_table_generation_plan = from_json_dict( spec.training_table_generation, TrainingTableGenerationPlan) + from kumo.engine._mlops import resolve + from kumo.engine.trainer._provenance import read_provenance + + trainer._embedding_models = with_text_encoders( + trainer._embedding_models, _encoders(bundle_dir), + Dependencies.of(read_provenance(bundle_dir), resolve())) holdout_path, holdout_truncated = _listed_holdout(bundle_dir) trainer._training_result = TrainingResult( metrics, paths, bundle_dir, holdout_path=holdout_path, @@ -1323,6 +1373,15 @@ def _load_bundle( return trainer +def _encoders(bundle_dir: Path) -> Any: + from kumo_core.encoders.config import EncodersConfig + from kumo_core.io.json import load_json + + path = listed_file(bundle_dir, 'encoders') + return (load_json(str(path), EncodersConfig) + if path is not None and path.is_file() else None) + + def _set_alias(connector: Connector, location: ModelRegistryLocation, alias: str) -> None: sink = connector.sink_backend @@ -1343,5 +1402,5 @@ def warn_if_other_bundle(metadata: Optional[Mapping[str, Any]], stacklevel=3) -def registry_pointers(result: TrainingResult) -> Dict[str, str]: - return {BUNDLE_ID_KEY: result.bundle_id} if result.bundle_id else {} +def registry_pointers(bundle_id: Optional[str]) -> Dict[str, str]: + return {BUNDLE_ID_KEY: bundle_id} if bundle_id else {} diff --git a/packages/kumo-engine/test/test_bundle_provenance.py b/packages/kumo-engine/test/test_bundle_provenance.py index 854e6e0dc..f07b93a10 100644 --- a/packages/kumo-engine/test/test_bundle_provenance.py +++ b/packages/kumo-engine/test/test_bundle_provenance.py @@ -22,7 +22,7 @@ ) from kumo.engine.trainer._provenance import embedded_model from kumo.engine.trainer._spec import MANIFEST_FILE -from kumo.engine.trainer.trainer import TrainingResult +from kumo.engine.trainer.trainer import TrainingResult, _store_and_register def test_fit_records_a_self_contained_manifest(fitted): @@ -103,13 +103,14 @@ def _finish_stub_fit(bundle, trainer, warm_start=None, settings=None): (bundle / MANIFEST_FILE).write_text(json.dumps({'artifacts': []})) add_bundle_file(bundle, 'trainer_config', 'model/config.json') table = SimpleNamespace(tt_id='tt-raw', id='tt-1', manifest=None) - return trainer._finish_fit( + result = trainer._finish_fit( TrainingResult({}, {}, bundle), predictive_query='PREDICT x', snapshot_id='snap', export_config=None, training_table_generation=TrainingTableGenerationPlan(), provenance=dict(train_table=table, warm_start=warm_start, - started_at=datetime.now(timezone.utc), fit_mode=None), - settings=settings) + started_at=datetime.now(timezone.utc), fit_mode=None)) + _store_and_register(result.bundle_dir, result.manifest, settings) + return result def test_missing_parents_are_stubs_and_retrain_chains_stay_flat(tmp_path): diff --git a/packages/kumo-engine/test/test_mlflow_model.py b/packages/kumo-engine/test/test_mlflow_model.py index 8937403b3..efd932dec 100644 --- a/packages/kumo-engine/test/test_mlflow_model.py +++ b/packages/kumo-engine/test/test_mlflow_model.py @@ -166,7 +166,7 @@ def predict_beside_another_call(self, *args, **kwargs): pyfunc.predict(rows, params={**params, 'population': population})), expected) - scratch = pyfunc.unwrap_python_model()._trainer.scratch_dir + scratch = pyfunc.unwrap_python_model()._predictor.trainer.scratch_dir assert list(scratch.iterdir()) == [] assert set(tables.glob('*')) == before | {tables / 'other-call'} @@ -261,7 +261,8 @@ def test_fit_waits_for_model_logging_and_fails_the_run_if_it_fails( import time from kumo.engine import _tracking - from kumo.engine._tracking import FitContext, TrackingSession + from kumo.engine._tracking import TrackingSession + from kumo.engine.trainer._run_spec import FitContext, fit_spec monkeypatch.setattr(_tracking, 'FINISH_SECONDS', 0.01) @@ -275,7 +276,7 @@ def start(self, record): def on_trial(self, record, event): pass - def log_model(self, result, snapshot, registry): + def log_model(self, bundle_dir, manifest, snapshot, registry): time.sleep(0.2) self.calls.append('log_model') if self.fail: @@ -288,11 +289,11 @@ def finish(self, record, error): for fail in (False, True): tracker = Tracker(fail) session = TrackingSession(tracker) - session.start(lambda: FitContext('r', 'q', {})) + session.start(lambda: fit_spec(FitContext('r', 'q', {}))) with warnings.catch_warnings(record=True) as caught: warnings.simplefilter('always') - _, error = session.log_model(result, None) - session.finish(result, error) + _, error = session.log_model(result.bundle_dir, None, None) + session.finish(None, error) session._worker.join(5) assert tracker.calls[:2] == ['start', 'log_model'] assert tracker.calls[2] == ('finish', error) @@ -305,7 +306,8 @@ def test_model_logging_wait_is_bounded(monkeypatch, stall): import threading from kumo.engine import _tracking - from kumo.engine._tracking import FitContext, TrackingSession + from kumo.engine._tracking import TrackingSession + from kumo.engine.trainer._run_spec import FitContext, fit_spec monkeypatch.setattr(_tracking, 'MODEL_SECONDS', 0.5) release = threading.Event() @@ -319,7 +321,7 @@ def start(self, record): def on_trial(self, record, event): pass - def log_model(self, result, snapshot, registry): + def log_model(self, bundle_dir, manifest, snapshot, registry): release.wait(30) def finish(self, record, error): @@ -327,12 +329,12 @@ def finish(self, record, error): tracker = Tracker() session = TrackingSession(tracker) - session.start(lambda: FitContext('r', 'q', {})) + session.start(lambda: fit_spec(FitContext('r', 'q', {}))) if stall == 'dead thread': session._submit(lambda: (_ for _ in ()).throw(SystemExit())) session._worker.join(5) with pytest.warns(UserWarning, match='Could not log the model'): - _, error = session.log_model(SimpleNamespace(), None) + _, error = session.log_model(None, None, None) release.set() session.finish(None, error) @@ -349,6 +351,7 @@ def test_model_upload_gets_a_longer_http_timeout_and_a_late_one_is_dropped( from kumo.engine._tracking import TrackingSession from kumo.engine.integrations import mlflow as integration + from kumo.engine.trainer._run_spec import fit_spec for name in integration._HTTP_LIMITS: monkeypatch.delenv(name, raising=False) @@ -356,7 +359,7 @@ def test_model_upload_gets_a_longer_http_timeout_and_a_late_one_is_dropped( location = kumo.ModelRegistryLocation('m') seen = [] - def log_model(target, client, result, snapshot, **kwargs): + def log_model(target, client, bundle_dir, manifest, snapshot, **kwargs): seen.append(os.environ.get('MLFLOW_HTTP_REQUEST_TIMEOUT')) if len(seen) == 2: tracker.finish(None, None) @@ -366,9 +369,8 @@ def log_model(target, client, result, snapshot, **kwargs): monkeypatch.setattr(integration, '_register', lambda *args: 'registered') for _ in range(2): session = TrackingSession(tracker) - session.start(lambda: CONTEXT) - registered, error = session.log_model( - SimpleNamespace(bundle_id=None), None, location) + session.start(lambda: fit_spec(CONTEXT)) + registered, error = session.log_model(None, None, None, location) assert error is None during = os.environ.get('MLFLOW_HTTP_REQUEST_TIMEOUT') session.finish(None, None) @@ -391,8 +393,7 @@ def test_a_model_with_text_columns_requires_the_text_extra(tmp_path): def requirement(encoders): path = tmp_path / 'encoders.json' save_json(EncodersConfig(encoders={'users': encoders}), str(path)) - return _requirement( - SimpleNamespace(artifact_paths={'encoders': str(path)})) + return _requirement(path) assert requirement({'age': Numerical()}) == ( f'kumo-engine[s3,gcs]=={kumo.__version__}') @@ -403,27 +404,28 @@ def requirement(encoders): def test_a_tracker_dropped_earlier_says_the_model_was_not_logged(): import time - from kumo.engine._tracking import FitContext, TrackingSession + from kumo.engine._tracking import TrackingSession + from kumo.engine.trainer._run_spec import FitContext, fit_spec class Tracker: def start(self, record): raise RuntimeError('server down') - def log_model(self, result, snapshot, registry): + def log_model(self, bundle_dir, manifest, snapshot, registry): raise AssertionError('not called') def finish(self, record, error): pass session = TrackingSession(Tracker()) - session.start(lambda: FitContext('r', 'q', {})) + session.start(lambda: fit_spec(FitContext('r', 'q', {}))) deadline = time.monotonic() + 10 while session._tracker is not None and time.monotonic() < deadline: time.sleep(0.01) with pytest.warns(UserWarning, match=r'Could not log the model: ' r'RuntimeError: tracking stopped after an earlier ' r'failure \(RuntimeError: server down\)'): - registered, error = session.log_model(SimpleNamespace(), None) + registered, error = session.log_model(None, None, None) session.finish(None, error) assert registered is None and isinstance(error, RuntimeError) assert str(session.error) == 'server down' diff --git a/packages/kumo-engine/test/test_mlflow_tracking.py b/packages/kumo-engine/test/test_mlflow_tracking.py index b28f30b4c..9b459e264 100644 --- a/packages/kumo-engine/test/test_mlflow_tracking.py +++ b/packages/kumo-engine/test/test_mlflow_tracking.py @@ -22,6 +22,7 @@ import kumo.engine as kumo from kumo.engine._scratch import get_scratch_dir from kumo.engine._tracking import run_record +from kumo.engine.trainer._run_spec import fit_spec mlflow = pytest.importorskip('mlflow') autologging_utils = pytest.importorskip('mlflow.utils.autologging_utils') @@ -249,11 +250,11 @@ def test_a_tracked_fit_on_sagemaker_says_how_to_open_the_ui( from kumo.engine.integrations.mlflow import MlflowTracking tracker = MlflowTracking(log_system_metrics=False) - tracker.start(run_record(CONTEXT)) + tracker.start(run_record(fit_spec(CONTEXT))) tracker.target.tracking_uri = uri with warnings.catch_warnings(record=True) as caught: warnings.simplefilter('always') - tracker.finish(run_record(CONTEXT), None) + tracker.finish(run_record(fit_spec(CONTEXT)), None) hints = [str(w.message) for w in caught if 'Studio' in str(w.message)] assert hints == ([] if command is None else [ "MLflow's run links do not open on SageMaker: open the MLflow UI " @@ -265,7 +266,7 @@ def test_user_disabled_autologging_stays_disabled(tracking_uri): tracker = MlflowTracking(log_system_metrics=False) with autologging_utils.disable_autologging(): - tracker.start(run_record(CONTEXT)) + tracker.start(run_record(fit_spec(CONTEXT))) tracker.finish(None, None) assert autologging_utils._AUTOLOGGING_GLOBALLY_DISABLED is True diff --git a/packages/kumo-engine/test/test_model_agnostic_mlops.py b/packages/kumo-engine/test/test_model_agnostic_mlops.py new file mode 100644 index 000000000..648c1b66d --- /dev/null +++ b/packages/kumo-engine/test/test_model_agnostic_mlops.py @@ -0,0 +1,269 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The MLOps seams take any bundle, run and predictor, not only the +Trainer's.""" + +import json +import os +import subprocess +import sys +import textwrap +from dataclasses import replace +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pandas as pd +import pyarrow as pa +import pyarrow.parquet as pq +import pytest +from kumo_core.artifacts import ( + PROVENANCE_FILE, + add_bundle_file, + bundle_id_from_manifest, +) +from kumoapi.encoder import TextEmbedding +from kumoapi.plugin import ModelRegistryLocation, PublishedModel +from kumoapi.typing import Stype + +import kumo.engine as kumo +from kumo.engine import _deps +from kumo.engine._deps import Dependencies, FileAsset, MissingDependencyError +from kumo.engine._manifest import ModelManifest, read_manifest, write_manifest +from kumo.engine._predictor import load_predictor +from kumo.engine._tracking import RunSpec, TrackingSession, run_record +from kumo.engine.tracking import Event +from kumo.engine.trainer._spec import MANIFEST_FILE, PROVENANCE_ROLE +from kumo.engine.trainer.trainer import _store_and_register + + +class ToyPredictor: + def __init__(self, weight): + self.weight = weight + + @classmethod + def load(cls, bundle_dir): + return cls(float((Path(bundle_dir) / 'weight.txt').read_text())) + + def predict(self, rows, snapshot, params): + return pd.DataFrame({'ENTITY': rows['ENTITY'], + 'SCORE': rows['ENTITY'] * self.weight}) + + +def _toy_bundle(root, predictor=f'{__name__}:ToyPredictor'): + root.mkdir(parents=True) + (root / 'weight.txt').write_text('2.0') + (root / MANIFEST_FILE).write_text(json.dumps({'artifacts': []})) + add_bundle_file(root, 'weights', 'weight.txt') + manifest = ModelManifest(id=bundle_id_from_manifest(root), + dependencies=[_deps.predictor(predictor)]) + write_manifest(root / PROVENANCE_FILE, manifest) + add_bundle_file(root, PROVENANCE_ROLE, PROVENANCE_FILE) + return root, manifest + + +class _Tracker: + def __init__(self): + self.calls = [] + + def start(self, record): + self.calls.append(('start', record)) + + def on_trial(self, record, event): + self.calls.append(('on_trial', record)) + + def log_model(self, bundle_dir, manifest, snapshot, registry): + self.calls.append(('log_model', (bundle_dir, manifest.id, snapshot))) + + def finish(self, record, error): + self.calls.append(('finish', record)) + + +def test_any_run_is_tracked_from_its_spec(): + tracker = _Tracker() + session = TrackingSession(tracker) + spec = RunSpec(kind='custom', run_id='loop-1', params={'lr': 0.1}, + inputs={'users': {'connector': 'S3'}}) + session.start(lambda: spec) + session.on_event(Event('trial_started', 'loop-1', 0.0, trial=1, + params={'attempt': 1})) + session.finish(lambda: replace(spec, metrics={'auc': 0.9}, + parents={'bundle_id': 'b'}), None) + session._worker.join(5) + started, trial, finished = (record for _, record in tracker.calls) + assert started == run_record(spec) + assert trial.tags['kumo.parent_run'] == 'loop-1' + assert finished.metrics == {'auc': 0.9} + assert finished.tags == {'kumo.kind': 'custom', 'kumo.run_id': 'loop-1', + 'kumo.bundle_id': 'b'} + + +def test_any_bundle_is_stored_logged_and_registered(tmp_path): + bundle, manifest = _toy_bundle(tmp_path / 'bundle') + published = [] + sink = SimpleNamespace( + publish_model=lambda connector, bundle_dir, **kw: published.append( + (bundle_dir, kw)) or PublishedModel(name='TOY', version='1', + uri='fake://TOY/1')) + registry = kumo.ModelOutputConfig(SimpleNamespace(sink_backend=sink), + location=ModelRegistryLocation('TOY')) + settings = kumo.MLOps(str(tmp_path / 'root'), registry=registry) + tracker = _Tracker() + session = TrackingSession(tracker) + session.start(lambda: RunSpec(kind='custom', run_id='loop-1')) + + registered, error = _store_and_register( + bundle, manifest, settings, session=session, snapshot='snap', + metrics={'auc': 0.9}) + + assert error is None and registered.version == '1' + stored = tmp_path / 'root' / 'models' / manifest.id + assert (stored / 'weight.txt').read_text() == '2.0' + assert ('log_model', (bundle, manifest.id, 'snap')) in tracker.calls + assert published == [(str(bundle), { + 'path': None, 'location': ModelRegistryLocation('TOY'), + 'metadata': {'kumo_bundle_id': manifest.id}, 'metrics': {'auc': 0.9} + })] + session.finish(None, None) + + +def test_a_bundle_is_scored_by_the_predictor_it_names(tmp_path): + pytest.importorskip('mlflow') + from kumo.engine.integrations._mlflow_model import KumoModel + + bundle, _ = _toy_bundle(tmp_path / 'toy') + rows = pd.DataFrame({'ENTITY': [1, 2]}) + assert load_predictor(bundle).predict(rows, None, {})['SCORE'].tolist( + ) == [2.0, 4.0] + model = KumoModel() + model.load_context(SimpleNamespace(artifacts={'bundle': str(bundle)}, + model_config={'columns': []})) + model._snapshots['s'] = 'snapshot' + assert model.predict(None, rows, {'snapshot': 's'})['SCORE'].tolist() == [ + 2.0, 4.0 + ] + + unknown, _ = _toy_bundle(tmp_path / 'unknown', 'nowhere.at_all:Model') + with pytest.raises(ValueError, match="names predictor 'nowhere.at_all:" + "Model', which cannot be imported here"): + load_predictor(unknown) + + +def test_a_file_asset_is_stored_once_and_resolved_by_its_hash(tmp_path): + (tmp_path / 'vocab').mkdir() + (tmp_path / 'vocab' / 'words.txt').write_text('a b c') + asset = FileAsset('toy', 'vocab', tmp_path / 'vocab', version='1') + settings = kumo.MLOps(str(tmp_path / 'root')) + recorded = Dependencies(settings=settings) + record = recorded.require(asset) + stored = tmp_path / 'root' / 'assets' / 'toy' / asset.fingerprint() + assert record.location == str(stored) + assert read_manifest(stored / PROVENANCE_FILE, _deps.AssetManifest).name \ + == 'vocab' + + resolved = Dependencies(recorded.records).resolve('toy', 'vocab') + assert (resolved / 'words.txt').read_text() == 'a b c' + + (stored / 'words.txt').write_text('changed') + with pytest.raises(MissingDependencyError, + match=rf"needs toy 'vocab' \(sha256 {record.sha256}\)"): + Dependencies(recorded.records).resolve('toy', 'vocab') + + +def test_a_manifest_without_dependencies_uses_the_trainer_predictor(): + old = ModelManifest.from_dict({'kind': 'model', 'schema_version': 1, + 'id': 'b'}) + assert old.dependencies is None + assert Dependencies.of(old).records == [] + + +WORDS = ['red', 'green', 'blue', 'cyan', 'gold', 'gray', 'pink'] + + +def _vectors(root): + root.mkdir() + rng = np.random.default_rng(0) + pd.DataFrame({'index': pd.Series(WORDS, dtype=object)}).to_parquet( + root / 'glove_tiny_index.parquet') + np.save(root / 'glove_tiny_8d_emb.npy', + np.asfortranarray(rng.normal(size=(len(WORDS), 8)).astype( + np.float32))) + return root + + +def _users(root): + rng = np.random.default_rng(1) + root.mkdir() + pq.write_table( + pa.table({ + 'id': pa.array(range(120), type=pa.int64()), + 'bio': [' '.join(rng.choice(WORDS, 6)) for _ in range(120)], + 'label': [bool(i % 2) if i < 100 else None for i in range(120)], + }), root / 'users.parquet') + return kumo.Table.from_source_table( + kumo.FileConnector(root_dir=str(root))['users'], + primary_key='id').with_column('bio', stype=Stype.text) + + +def test_a_model_serves_from_its_bundle_snapshot_and_assets_alone( + tmp_path, monkeypatch): + monkeypatch.setattr('kumo.engine._scratch._root', tmp_path / 'scratch') + graph = kumo.Graph([_users(tmp_path / 'source')]) + snapshot = graph.snapshot() + query = kumo.PredictiveQuery(graph, + 'PREDICT users.label FOR EACH users.id') + plan = query.suggest_model_plan(snapshot) + plan.training_job.num_experiments = 1 + plan.training_job.disable_explain = True + plan.optimization.max_epochs = 1 + plan.column_processing.encoder_overrides = { + 'users.bio': TextEmbedding('colors') + } + store = kumo.EmbeddingModelStore() + store.add('colors', kumo.GloVeVectors(_vectors(tmp_path / 'vectors'))) + trainer = kumo.Trainer(plan, embedding_models=store, accelerator='cpu') + root = tmp_path / 'root' + result = trainer.fit(query.generate_training_table(snapshot), + snapshot=snapshot, mlops=kumo.MLOps(str(root))) + location = snapshot.export( + kumo.SnapshotOutputConfig( + kumo.FileConnector(root_dir=str(tmp_path / 'exports')), 'snap')) + expected = trainer.predict( + snapshot, binary_classification_threshold=.5).frame.sort_values( + 'ENTITY').reset_index(drop=True) + expected.to_parquet(tmp_path / 'expected.parquet') + dependency, = Dependencies.of(result.manifest).assets('text_encoder') + assert dependency.name == 'colors' and dependency.packing == 'external' + assert Path(dependency.location).parent == root / 'assets' / 'text_encoder' + + serve = tmp_path / 'serve' + (serve / 'home').mkdir(parents=True) + script = textwrap.dedent(f""" + import pandas as pd + import kumo.engine as kumo + snapshot = kumo.GraphSnapshot.load_from_connector( + kumo.FileConnector(root_dir={str(tmp_path / 'exports')!r}), + path={location!r}) + trainer = kumo.Trainer.load({str(root / 'models' / result.id)!r}, + accelerator='cpu') + scored = trainer.predict( + snapshot, binary_classification_threshold=.5).frame.sort_values( + 'ENTITY').reset_index(drop=True) + pd.testing.assert_frame_equal( + scored, pd.read_parquet({str(tmp_path / 'expected.parquet')!r})) + """) + env = {**os.environ, 'HOME': str(serve / 'home'), 'CUDA_VISIBLE_DEVICES': + '', 'HTTPS_PROXY': 'http://127.0.0.1:9', + 'HTTP_PROXY': 'http://127.0.0.1:9'} + for name in ('XDG_CACHE_HOME', 'NO_PROXY', 'no_proxy'): + env.pop(name, None) + done = subprocess.run([sys.executable, '-c', script], cwd=serve, env=env, + capture_output=True, text=True, timeout=600) + assert done.returncode == 0, done.stderr[-3000:] + + (root / 'assets').rename(tmp_path / 'moved') + done = subprocess.run([sys.executable, '-c', script], cwd=serve, env=env, + capture_output=True, text=True, timeout=600) + assert "MissingDependencyError: The model needs text_encoder 'colors'" \ + in done.stderr diff --git a/packages/kumo-engine/test/test_snowflake_tracking.py b/packages/kumo-engine/test/test_snowflake_tracking.py index 3ca74f771..aebe648c5 100644 --- a/packages/kumo-engine/test/test_snowflake_tracking.py +++ b/packages/kumo-engine/test/test_snowflake_tracking.py @@ -18,6 +18,7 @@ ) from kumo.engine._tracking import TrackingSession +from kumo.engine.trainer._run_spec import fit_spec, fitted_spec FAKED = ('__init__', 'set_experiment', 'start_run', 'end_run', 'log_params', 'log_metrics', 'log_artifact') @@ -63,7 +64,7 @@ def _tracking(session=FakeSession, **kwargs): def _replay(tracking, result=RESULT, error=None, during=None): session = TrackingSession(tracking) - session.start(lambda: CONTEXT) + session.start(lambda: fit_spec(CONTEXT)) for event in EVENTS: session.on_event(event) if during is not None: @@ -71,7 +72,10 @@ def _replay(tracking, result=RESULT, error=None, during=None): session._submit(caught_up.set) caught_up.wait(30) during() - session.finish(result, error) + session.finish( + None if result is None else + lambda: fitted_spec(session.spec, result, EVENTS[-1].metrics['seconds']), + error) def test_native_model_lineage_holds_the_run_open_during_the_fit(fake): diff --git a/packages/kumo-engine/test/test_tracking_contract.py b/packages/kumo-engine/test/test_tracking_contract.py index 2efe3cb35..bcf9cf3f0 100644 --- a/packages/kumo-engine/test/test_tracking_contract.py +++ b/packages/kumo-engine/test/test_tracking_contract.py @@ -6,7 +6,7 @@ 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 @@ -17,11 +17,11 @@ from kumo.engine._tracking import ( DROPPED_METRICS, DROPPED_PARAMS, - FitContext, TrackingSession, run_record, ) from kumo.engine.tracking import Event, RunRecord +from kumo.engine.trainer._run_spec import FitContext, fit_spec, fitted_spec RUN = 'trainer-fit-1' USER_TAGS = {'team': 'ents', 'kumo.run_id': 'overridden'} @@ -118,15 +118,22 @@ def replay(tracker=None, events=EVENTS, result=RESULT, of it.""" recorder = _Recorder(tracker) session = TrackingSession(recorder) - session.start(lambda: CONTEXT) + session.start(lambda: fit_spec(CONTEXT)) for event in events: session.on_event(event) - session.finish(result, error) + session.finish( + None if result is None else lambda: fitted_spec( + session.spec, result, events[-1].metrics.get('seconds')), error) records = { record.tags['kumo.run_id']: record for hook, record in recorder.calls if hook == 'on_trial' } - records[RUN] = run_record(CONTEXT, result, events[-1], tags=USER_TAGS) + spec = fit_spec(CONTEXT, tags=USER_TAGS) + if result is not None: + spec = fitted_spec(spec, result, events[-1].metrics.get('seconds')) + if 'error' in events[-1].params: + spec = replace(spec, error=str(events[-1].params['error'])) + records[RUN] = run_record(spec) return records @@ -304,9 +311,8 @@ def test_records_of_the_recorded_fit(): } with pytest.warns(UserWarning, match="'kumo.bundle_id', 'training_job.tune_metric'"): - spoofed = run_record(CONTEXT, tags={'kumo.bundle_id': 'x', - 'training_job.tune_metric': 'y', - 'a': 1}) + spoofed = run_record(fit_spec(CONTEXT, tags={ + 'kumo.bundle_id': 'x', 'training_job.tune_metric': 'y', 'a': 1})) assert spoofed.tags.keys() - fit.tags.keys() == {'a'} assert fit.metrics == {'val_mae': 0.3, 'test_mae': 0.32, 'seconds': 9.0} assert fit.artifacts['query.pql'] == QUERY @@ -370,8 +376,8 @@ def test_only_rank_zero_tracks(monkeypatch): recorder = _Recorder() session = TrackingSession(recorder) - session.start(lambda: CONTEXT) - session.finish(RESULT, None) + session.start(lambda: fit_spec(CONTEXT)) + session.finish(lambda: fitted_spec(session.spec, RESULT), None) assert recorder.calls == [] From d2565d15bab0a8132c73aba2ee447413adad49d4 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 28 Oct 2026 21:13:47 +0000 Subject: [PATCH 020/278] Register through the tracker, record its own run, and set scoring defaults at fit kumo.MLOps(registry=) takes a ModelRegistryLocation only: MlflowTracking registers an MLflow pyfunc (Unity Catalog, SageMaker MLflow) and SnowflakeTracking registers in the Snowflake Model Registry through its connector. Trainer.load_from_connector(cfg, location=...) loads a registered version back the same way. Runs carry no Kumo run id: the model manifest records the tracker's own run as tracking (platform, run id, URL), read as TrainingResult.tracking, and trial runs are found by name. fit(scoring_defaults=...) stores the defaults in the bundle before it is stored and registered. MlflowTracking(target=) also takes a tracking URI and log_model/load_model take target=; SnowflakeTracking takes the connector or session first and schema_name=. Co-Authored-By: Claude Opus 5.5 --- examples/mlops/smoke.py | 23 ++-- packages/kumo-engine/kumo/engine/_events.py | 10 +- packages/kumo-engine/kumo/engine/_manifest.py | 6 ++ packages/kumo-engine/kumo/engine/_mlops.py | 71 +++++++++---- packages/kumo-engine/kumo/engine/_tracking.py | 48 ++++++--- .../engine/integrations/_mlflow_target.py | 55 ++++++++-- .../kumo/engine/integrations/mlflow.py | 100 ++++++++++++------ .../kumo/engine/integrations/snowflake.py | 58 ++++++---- .../kumo/engine/trainer/_provenance.py | 5 +- .../kumo/engine/trainer/_run_spec.py | 2 - .../kumo/engine/trainer/trainer.py | 73 ++++++++++--- packages/kumo-engine/test/test_fit_events.py | 3 +- .../kumo-engine/test/test_mlflow_model.py | 10 +- .../kumo-engine/test/test_mlflow_target.py | 37 ++++--- .../kumo-engine/test/test_mlflow_tracking.py | 7 +- packages/kumo-engine/test/test_mlops.py | 29 ++--- .../test/test_model_agnostic_mlops.py | 50 ++++++--- .../kumo-engine/test/test_scoring_defaults.py | 24 +++++ .../test/test_snowflake_tracking.py | 41 +++++-- .../test/test_tracking_contract.py | 37 ++++--- .../kumo-engine/test/test_trainer_registry.py | 36 +++++-- 21 files changed, 497 insertions(+), 228 deletions(-) diff --git a/examples/mlops/smoke.py b/examples/mlops/smoke.py index 47e152564..0d4b1cfcf 100644 --- a/examples/mlops/smoke.py +++ b/examples/mlops/smoke.py @@ -88,8 +88,8 @@ def tracking(args: argparse.Namespace, suffix: str): schema_name=args.schema, identity=SnowflakeIdentity( account=args.account, username=args.user, private_key=Path(args.private_key_file).read_text())) - return SnowflakeTracking(connector=connector, - experiment_name='KUMO_SMOKE'), None + return (SnowflakeTracking(connector, experiment_name='KUMO_SMOKE'), + f'KUMO_SMOKE_{suffix}') def main(argv: Optional[Sequence[str]] = None) -> None: @@ -109,10 +109,9 @@ def main(argv: Optional[Sequence[str]] = None) -> None: kumo.configure(scratch_dir=work / 'scratch') graph = data(work / 'data') tracker, model_name = tracking(args, time.strftime('%Y%m%d%H%M%S')) - cfg = kumo.MLOps( - args.artifact_root or str(work / 'root'), tracking=tracker, - registry=None if model_name is None else ModelRegistryLocation( - model_name, alias='candidate'), on_error='raise') + 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() @@ -122,15 +121,11 @@ def main(argv: Optional[Sequence[str]] = None) -> None: plan.training_job.num_experiments = 1 plan.optimization.max_epochs = 1 result = kumo.Trainer(plan).fit(table, snapshot=snapshot) - print(f'fit run {tracker.run_id}, bundle {result.id}, ' + print(f'fit run {result.tracking}, bundle {result.id}, ' f'registered {result.registered_model}') - if model_name is not None: - from kumo.engine.integrations.mlflow import load_model - - loaded = load_model(f'models:/{model_name}@candidate', - tracking=tracker) - assert loaded.result.id == result.id, (loaded.result.id, result.id) - print('loaded the registered model back') + 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') diff --git a/packages/kumo-engine/kumo/engine/_events.py b/packages/kumo-engine/kumo/engine/_events.py index 7309aab19..98e8c2394 100644 --- a/packages/kumo-engine/kumo/engine/_events.py +++ b/packages/kumo-engine/kumo/engine/_events.py @@ -32,7 +32,6 @@ class Event: r"""One progress event. ``trial`` and ``epoch`` count from 1.""" name: EventName - run_id: str time: float trial: Optional[int] = None epoch: Optional[int] = None @@ -60,17 +59,14 @@ def _finite_floats(metrics: Mapping[str, Any]) -> Dict[str, float]: class EventEmitter(NoOpTrainingInstrumentation): - def __init__(self, run_id: str, - callbacks: Sequence[Callable[[Event], None]]) -> None: - self.run_id = run_id + def __init__(self, callbacks: Sequence[Callable[[Event], None]]) -> None: self._callbacks = list(callbacks) def emit(self, name: EventName, *, trial: Optional[int] = None, epoch: Optional[int] = None, metrics: Optional[Mapping[str, float]] = None, params: Optional[Mapping[str, Any]] = None) -> None: - event = Event(name=name, run_id=self.run_id, time=time.time(), - trial=trial, epoch=epoch, + event = Event(name=name, time=time.time(), trial=trial, epoch=epoch, metrics=_finite_floats(metrics or {}), params=dict(params or {})) for callback in list(self._callbacks): @@ -79,7 +75,7 @@ def emit(self, name: EventName, *, trial: Optional[int] = None, except Exception: logger.warning( 'Progress callback %r failed; it is disabled for the ' - 'rest of run %s', callback, self.run_id, exc_info=True) + 'rest of this fit', callback, exc_info=True) self._callbacks.remove(callback) def event(self, name: str, **fields: Any) -> None: diff --git a/packages/kumo-engine/kumo/engine/_manifest.py b/packages/kumo-engine/kumo/engine/_manifest.py index 8741da702..df1fd4582 100644 --- a/packages/kumo-engine/kumo/engine/_manifest.py +++ b/packages/kumo-engine/kumo/engine/_manifest.py @@ -315,7 +315,13 @@ class FitTiming(Record): @dataclass(frozen=True, kw_only=True) class TrackingRef(Record): + r"""A run in an experiment tracker: its ``platform`` (``mlflow``, + ``databricks``, ``sagemaker``, ``snowflake``), the platform's own + ``run_id`` and, where the platform has one, the run's ``url``.""" + + platform: str run_id: str + url: Optional[str] = omit_none() @dataclass(frozen=True, kw_only=True) diff --git a/packages/kumo-engine/kumo/engine/_mlops.py b/packages/kumo-engine/kumo/engine/_mlops.py index ac4480fab..bf3ef384a 100644 --- a/packages/kumo-engine/kumo/engine/_mlops.py +++ b/packages/kumo-engine/kumo/engine/_mlops.py @@ -7,6 +7,7 @@ import contextvars import os import posixpath +import re import warnings from dataclasses import dataclass, field from typing import ( @@ -55,9 +56,10 @@ class MLOps: ``kumo.engine.tracking.Tracking``. Its ``check(registry)``, if any, runs here, so a server or registry that is unreachable or refuses these credentials fails before any work. - registry: where fits register models: a ``ModelOutputConfig`` with a - registry ``location``, or a ``ModelRegistryLocation`` registered - through the MLflow tracker. + registry: a ``ModelRegistryLocation`` each fit registers its model + at, in the tracker's registry: MLflow's (Unity Catalog, SageMaker + MLflow) for ``MlflowTracking``, the Snowflake Model Registry for + ``SnowflakeTracking``. store: the artifact kinds stored under the root when produced (``snapshot``, ``training_table``, ``model``); ``True`` for all, ``False`` for none. @@ -112,30 +114,39 @@ def __post_init__(self) -> None: f'{list(KINDS)}.') if kinds: _require_writable(root) + if self.registry is not None: + self._check_registry() + check = getattr(self.tracking, 'check', None) + if callable(check): + check(self.registry) + object.__setattr__(self, 'root', root) + object.__setattr__(self, 'kinds', kinds) + + def _check_registry(self) -> None: from kumo.engine.trainer.output_config import ( ModelOutputConfig, check_model_location, ) - if isinstance(self.registry, ModelOutputConfig): - check_model_location(self.registry) - 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: + registry = self.registry + if not isinstance(registry, ModelRegistryLocation): + raise TypeError( + f'registry takes a ModelRegistryLocation, not ' + f'{type(registry).__name__}; the tracker registers it.') + if registry.alias is not None: + check_model_alias(registry.alias) + if callable(getattr(self.tracking, 'log_model', None)): + if 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) - check = getattr(self.tracking, 'check', None) - if callable(check): - check(self.registry if isinstance( - self.registry, ModelRegistryLocation) else None) - object.__setattr__(self, 'root', root) - object.__setattr__(self, 'kinds', kinds) + return + connector = getattr(self.tracking, 'connector', None) + if connector is None: + raise ValueError( + 'registry= registers through the tracker: pass ' + 'tracking=MlflowTracking(...) or SnowflakeTracking(connector, ' + '...).') + check_model_location(ModelOutputConfig(connector, location=registry)) def wants(self, kind: str) -> bool: r"""Whether artifacts of ``kind`` are stored under the root.""" @@ -269,3 +280,23 @@ def snapshot_path(snapshot: Any) -> str: r"""Where a snapshot, or a snapshot id, is stored under a root.""" snapshot_id = snapshot if isinstance(snapshot, str) else snapshot.id return 'snapshots/' + snapshot_id.replace('@', '/') + + +_SNAPSHOT_ID = re.compile(r'[0-9a-f]+@[0-9a-f]+') + + +def snapshot_uri(snapshot: Any, + settings: Optional[MLOps] = None) -> Optional[str]: + r"""A snapshot to serve against, as a location: a URI or stage path as + given; a `GraphSnapshot` or snapshot id as stored under the artifact + root of ``settings`` (else of the enclosing block).""" + if snapshot is None or (isinstance(snapshot, str) + and not _SNAPSHOT_ID.fullmatch(snapshot)): + return snapshot + settings = resolve(settings) + if settings is None: + raise ValueError( + 'A snapshot given by object or id is read from the artifact ' + 'root: call this in a kumo.mlops block, or pass the location ' + 'snapshot.export(...) returned.') + return settings.root.uri(snapshot_path(snapshot)) diff --git a/packages/kumo-engine/kumo/engine/_tracking.py b/packages/kumo-engine/kumo/engine/_tracking.py index 562059341..159c9a83a 100644 --- a/packages/kumo-engine/kumo/engine/_tracking.py +++ b/packages/kumo-engine/kumo/engine/_tracking.py @@ -35,7 +35,7 @@ if TYPE_CHECKING: from kumoapi.plugin import ModelRegistryLocation, PublishedModel - from kumo.engine._manifest import ModelManifest + from kumo.engine._manifest import ModelManifest, TrackingRef logger = logging.getLogger(__name__) @@ -73,7 +73,6 @@ class RunSpec: the user's own.""" kind: str - run_id: str name: Optional[str] = None params: Mapping[str, Any] = field(default_factory=dict) metrics: Mapping[str, Any] = field(default_factory=dict) @@ -136,7 +135,6 @@ def run_record(spec: RunSpec) -> RunRecord: tags={ **_user_tags(spec.tags, params), 'kumo.kind': spec.kind, - 'kumo.run_id': spec.run_id, **{f'kumo.{name}': str(value) for name, value in spec.parents.items() if value}, **({} if spec.status is None else {'kumo.status': spec.status}), @@ -161,8 +159,8 @@ def _user_tags(tags: Optional[Mapping[str, Any]], } -def trial_spec(run_id: str, events: Sequence[Event]) -> RunSpec: - r"""One trial attempt of run ``run_id``, from its events so far.""" +def trial_spec(events: Sequence[Event]) -> RunSpec: + r"""One trial attempt of a fit, from its events so far.""" started = next(event for event in events if event.name == 'trial_started') trial, attempt = started.trial, int(started.params['attempt']) metrics: Dict[str, float] = {} @@ -182,12 +180,10 @@ def trial_spec(run_id: str, events: Sequence[Event]) -> RunSpec: }) return RunSpec( kind='trial', - run_id=f'{run_id}-trial-{trial}-attempt-{attempt}', name=f'trial-{trial}-attempt-{attempt}', params=dict(started.params), metrics=metrics, steps=steps, - parents={'parent_run': run_id}, **ended, ) @@ -205,7 +201,10 @@ class Tracking(Protocol): `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``. + ``ModelRegistryLocation``; ``connector``, a connector whose model + registry ``registry`` registers in when there is no ``log_model``; and + ``platform``, ``run_id`` and ``run_url``, the tracker's own run once + `start` returned, recorded as the model manifest's ``tracking``. """ def start(self, record: RunRecord) -> None: ... @@ -234,6 +233,7 @@ def __init__(self, tracking: Optional[Tracking], QUEUE_SIZE) self._worker: Optional[threading.Thread] = None self._behind = False + self._started = threading.Event() self.error: Optional[BaseException] = None @property @@ -256,10 +256,30 @@ def start(self, spec: Callable[[], RunSpec]) -> None: self._worker = threading.Thread(target=self._work, name='kumo-tracking', daemon=True) self._worker.start() - self._submit(lambda: self._call(lambda tracker: tracker.start( - run_record(dataclasses.replace(value, tags={ + + def begin(tracker: Tracking) -> None: + tracker.start(run_record(dataclasses.replace(value, tags={ **value.tags, **(getattr(tracker, 'tags', None) or {}) - }))))) + }))) + + started = self._submit(lambda: self._call(begin)) + if not (started and self._submit(self._started.set)): + self._started.set() + + def pointer(self) -> Optional['TrackingRef']: + r"""The tracker's own run, once it started: its platform, run id + and URL.""" + from kumo.engine._manifest import TrackingRef + + if self.spec is None or not self._started.wait(FINISH_SECONDS): + return None + tracker = self._tracker + run_id = getattr(tracker, 'run_id', None) + if tracker is None or not run_id: + return None + return TrackingRef( + platform=str(getattr(tracker, 'platform', type(tracker).__name__)), + run_id=str(run_id), url=getattr(tracker, 'run_url', None)) def on_event(self, event: Event) -> None: if self.spec is None or event.trial is None: @@ -272,7 +292,7 @@ def on_event(self, event: Event) -> None: **event.params, 'attempt': events[0].params['attempt'] }) events.append(event) - record = run_record(trial_spec(self.spec.run_id, events)) + record = run_record(trial_spec(events)) self._submit(lambda: self._call( lambda tracker: tracker.on_trial(record, event))) @@ -285,8 +305,8 @@ def log_model( ) -> 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 (self._tracker is not None and not hasattr( - self._tracker, 'log_model')): + if self._worker is None or (self._tracker is not None and not callable( + getattr(self._tracker, 'log_model', None))): return None, None done = threading.Event() failed: List[BaseException] = [] diff --git a/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py b/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py index 95348b466..8372eec35 100644 --- a/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py +++ b/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py @@ -56,11 +56,34 @@ class Target: def __init__(self, tracking_uri: str, registry_uri: Optional[str], swap: Optional[Callable[[], ContextManager[Any]]] = None, - aws_session: Any = None) -> None: + aws_session: Any = None, host: Optional[str] = None) -> None: self.tracking_uri = tracking_uri self.registry_uri = registry_uri self._swap = swap self._aws_session = aws_session + self._host = host + + @property + def platform(self) -> str: + if SAGEMAKER_ARN.fullmatch(self.tracking_uri or ''): + return 'sagemaker' + if (self.tracking_uri or '').partition(':')[0] == 'databricks': + return 'databricks' + return 'mlflow' + + def run_url(self, experiment_id: str, run_id: str) -> Optional[str]: + r"""The run's page in the MLflow UI; None where no link opens it + (SageMaker, local stores).""" + if self.platform == 'databricks': + host = self._host or _workspace_host() + return None if host is None else ( + f"{host.rstrip('/')}/ml/experiments/{experiment_id}/runs/" + f'{run_id}') + if self.platform == 'mlflow' and self.tracking_uri.startswith( + ('http://', 'https://')): + return (f"{self.tracking_uri.rstrip('/')}/#/experiments/" + f'{experiment_id}/runs/{run_id}') + return None @property def explicit(self) -> bool: @@ -93,22 +116,33 @@ def artifact_repository(self, uri: str) -> Any: session_token=frozen.token) -def resolve(target: Any, tracking_uri: Optional[str]) -> Target: +def resolve(target: Any) -> Target: mlflow = import_provider('mlflow', 'mlflow') - if target is not None and tracking_uri is not None: - raise ValueError('Pass target or tracking_uri, not both.') + if isinstance(target, Target): + return target if target is None: - if tracking_uri is None: - return Target(mlflow.get_tracking_uri(), mlflow.get_registry_uri()) - return Target(tracking_uri, registry_for(tracking_uri)) + return Target(mlflow.get_tracking_uri(), mlflow.get_registry_uri()) + if isinstance(target, str): + if SAGEMAKER_ARN.fullmatch(target): + return _sagemaker(SageMakerMLflow(target)) + return Target(target, registry_for(target)) if isinstance(target, SageMakerMLflow): return _sagemaker(target) from kumo.engine.connector.databricks import DatabricksConnector if isinstance(target, DatabricksConnector): return _databricks(target) - raise TypeError(f'MlflowTracking(target=...) takes a DatabricksConnector ' - f'or a SageMakerMLflow, not {type(target).__name__}.') + raise TypeError(f'MlflowTracking(target=...) takes a DatabricksConnector, ' + f'a SageMakerMLflow or a tracking URI, not ' + f'{type(target).__name__}.') + + +def _workspace_host() -> Optional[str]: + try: + return import_provider('databricks.sdk.config', + 'databricks').Config().host + except Exception: + return None def registry_for(tracking_uri: str) -> str: @@ -135,7 +169,8 @@ def _databricks(connector: Any) -> Target: client = import_provider('kumo_databricks.session', 'databricks')._client_for(principal) return Target('databricks', 'databricks-uc', - swap=lambda: _databricks_token(client)) + swap=lambda: _databricks_token(client), + host=client.config.host) @contextmanager diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index e5020d97c..49d8a4b4e 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -24,6 +24,7 @@ FrozenSet, Iterator, List, + Literal, Mapping, Optional, Set, @@ -83,8 +84,8 @@ class MlflowTracking: Args: experiment_name: experiment for the fit run, created if missing. - `None` uses MLflow's current experiment; required with - ``target`` or ``tracking_uri``. + `None` uses MLflow's current experiment; required to record + runs at a ``target``. run_name: name of the fit run. `None` names it after its UTC start time, as ``fit-20261027T061838Z``. log_system_metrics: sample CPU, memory, disk, network and (with @@ -94,13 +95,11 @@ class MlflowTracking: before the fit returns. ``kumo.MLOps(registry= ModelRegistryLocation(...))`` also registers it, which implies this. - serving_snapshot: exported snapshot URI the served model loads at - startup (see `log_model`). target: where to record, with the credentials that reach it: a ``DatabricksConnector`` (Databricks MLflow and Unity Catalog, - as its identity) or a `SageMakerMLflow`. `None` uses MLflow's - own tracking and registry URIs, read once here. - tracking_uri: another MLflow server, instead of ``target``. + as its identity), a `SageMakerMLflow`, or another MLflow + server's tracking URI. `None` uses MLflow's own tracking and + registry URIs, read once here. With explicit credentials (a Databricks service principal, an AWS access key or role), each MLflow call sets the state MLflow reads them @@ -111,7 +110,8 @@ class MlflowTracking: except that a SageMaker target still sets the process-wide tracking URI for each call, since the plugin signs and routes with it. - ``run_id`` is the MLflow id of the latest fit run, once it has started. + ``run_id`` is the MLflow id of the latest fit run, once it has started, + and ``run_url`` its link in the MLflow UI, where the server has one. """ def __init__( @@ -121,27 +121,20 @@ def __init__( log_system_metrics: bool = True, tags: Optional[Mapping[str, str]] = None, log_model: bool = False, - serving_snapshot: Optional[str] = None, *, target: Any = None, - tracking_uri: Optional[str] = None, ) -> None: self._mlflow = import_provider('mlflow', 'mlflow') - self.target = resolve_target(target, tracking_uri) - if experiment_name is None and (target is not None - or tracking_uri is not None): - raise ValueError('MlflowTracking(target=...) and ' - 'MlflowTracking(tracking_uri=...) need an ' - 'experiment_name.') + self.target = resolve_target(target) + self._named = target is not None self._default_experiment: Optional[str] = ( - None if experiment_name is not None else + None if experiment_name is not None or self._named else import_provider('mlflow.tracking.fluent', 'mlflow')._get_experiment_id()) self.experiment_name = experiment_name self.run_name = run_name self.log_system_metrics = log_system_metrics self.tags = dict(tags or {}) - self.serving_snapshot = serving_snapshot self.logs_model = log_model self._outer: Optional[str] = None self.run_id: Optional[str] = None @@ -157,6 +150,16 @@ def _reset(self) -> None: self._http_owned: FrozenSet[str] = frozenset() self._guard = ExitStack() + @property + def platform(self) -> str: + return self.target.platform + + @property + def run_url(self) -> Optional[str]: + if self.run_id is None or self._experiment_id is None: + return None + return self.target.run_url(self._experiment_id, self.run_id) + def check_model_name(self, name: str) -> None: r"""Raises :class:`ValueError` unless this tracker's registry accepts ``name``.""" @@ -167,6 +170,7 @@ def check(self, registry: Optional[ModelRegistryLocation] = None) -> None: credentials and, given ``registry``, can register its model: one experiment lookup, one registry listing and, on Unity Catalog, the model's schema.""" + self._require_experiment() if registry is not None: self.check_model_name(registry.model_name) uri = self.target.tracking_uri @@ -195,7 +199,26 @@ def prepare(self) -> None: self._outer = (active.info.run_id if active is not None and same else None) + def _require_experiment(self) -> None: + if self._named and self.experiment_name is None: + raise ValueError('MlflowTracking(target=...) records runs in an ' + 'experiment: pass experiment_name.') + + def load_model(self, location: ModelRegistryLocation, + **kwargs: Any) -> 'Trainer': + r"""The `Trainer` of the version of ``location`` its alias or + version names, from this tracker's registry.""" + if (location.alias is None) == (location.version is None): + raise ValueError('Loading from an MLflow registry takes a ' + 'ModelRegistryLocation with an alias or a ' + 'version.') + reference = (f'@{location.alias}' + if location.alias else f'/{location.version}') + return load_model(f'models:/{location.model_name}{reference}', + target=self.target, **kwargs) + def start(self, record: RunRecord) -> None: + self._require_experiment() self._reset() self.run_id = None self._http_owned = frozenset(name for name in _HTTP_LIMITS @@ -282,7 +305,7 @@ def log_model( info = _log_model(self.target, self._client, bundle_dir, manifest, snapshot, name='model', run_id=run_id, experiment_id=self._experiment_id, - serving_snapshot=self.serving_snapshot) + serving_snapshot=None) finally: if self._parent == run_id: os.environ.update({name: _HTTP_LIMITS[name] for name in upload}) @@ -499,26 +522,32 @@ def log_model( snapshot: 'GraphSnapshot', name: str = 'model', run_id: Optional[str] = None, - serving_snapshot: Optional[str] = None, - tracking: Optional[MlflowTracking] = None, + serving_snapshot: Any = None, + target: Any = None, + experiment_name: Optional[str] = None, ) -> Any: r"""Log a trainer's bundle as an MLflow pyfunc model in ``run_id`` (or - the active run, else a new one) with its source tables as run inputs; - returns its `ModelInfo`. ``tracking`` says where, with what credentials - (its experiment holds a new run); `None` uses MLflow's own tracking - URI. Scores with ``predict(entities, params={'snapshot': uri})``; with - ``serving_snapshot`` (an exported snapshot URI), each serving replica - downloads that snapshot once at startup and scores against it when no - ``snapshot`` param is given.""" + the active run, else a new one in ``experiment_name``) with its source + tables as run inputs; returns its `ModelInfo`. ``target`` says where, + with what credentials, as in `MlflowTracking`; `None` uses MLflow's own + tracking URI. Scores with ``predict(entities, params={'snapshot': + uri})``; with ``serving_snapshot`` (an exported snapshot URI, or a + `GraphSnapshot` or snapshot id stored under the artifact root of the + enclosing `kumo.mlops` block), each serving replica downloads that + snapshot once at startup and scores against it when no ``snapshot`` + param is given.""" from kumo.engine._manifest import stub + from kumo.engine._mlops import snapshot_uri from kumo.engine._tracking import fit_inputs mlflow = import_provider('mlflow', 'mlflow') result = trainer.result if result is None: raise ValueError('log_model requires a fitted or loaded Trainer') - tracking = tracking or MlflowTracking() + tracking = MlflowTracking(experiment_name, target=target) target = tracking.target + if run_id is None and tracking._named and experiment_name is None: + tracking._require_experiment() client = target.client() active = mlflow.active_run() if run_id is None and active is not None and ( @@ -539,7 +568,7 @@ def log_model( (None if result.id is None else stub('model', result.id)), snapshot, name=name, run_id=run_id, experiment_id=experiment_id, - serving_snapshot=serving_snapshot) + serving_snapshot=snapshot_uri(serving_snapshot)) with target.scope(): if result.metrics: client.log_batch(run_id, metrics=[ @@ -557,22 +586,25 @@ def load_model( model_uri: str, embedding_models: Optional['EmbeddingModelStore'] = None, *, - tracking: Optional[MlflowTracking] = None, + target: Any = None, + accelerator: Literal['cpu', 'gpu', 'mps', 'auto'] = 'auto', ) -> 'Trainer': r"""Restore the `Trainer` of a model `log_model` logged, from the - server ``tracking`` reaches (`None`: MLflow's own tracking URI).""" + server ``target`` reaches, as in `MlflowTracking` (`None`: MLflow's own + tracking URI).""" from kumo.engine._scratch import new_run_id, trainer_scratch_dir from kumo.engine.trainer.trainer import Trainer, warn_if_other_bundle mlflow = import_provider('mlflow', 'mlflow') - target = (tracking or MlflowTracking()).target + target = resolve_target(target) destination = trainer_scratch_dir(new_run_id()) / 'mlflow-model' destination.mkdir(parents=True) with target.scope(): local = Path(_download(mlflow, target, model_uri, destination)) model = mlflow.models.Model.load(str(local)) bundle = model.flavors['python_function']['artifacts']['bundle']['path'] - trainer = Trainer.load(local / bundle, embedding_models) + trainer = Trainer.load(local / bundle, embedding_models, + accelerator=accelerator) warn_if_other_bundle(model.metadata, trainer) return trainer diff --git a/packages/kumo-engine/kumo/engine/integrations/snowflake.py b/packages/kumo-engine/kumo/engine/integrations/snowflake.py index dfa463060..78b3e23b6 100644 --- a/packages/kumo-engine/kumo/engine/integrations/snowflake.py +++ b/packages/kumo-engine/kumo/engine/integrations/snowflake.py @@ -7,6 +7,7 @@ import logging import re import tempfile +import uuid from contextlib import ExitStack from pathlib import Path from typing import Any, Dict, List, Mapping, Optional, Set, Tuple @@ -36,16 +37,17 @@ class SnowflakeTracking: parameters are not changed. Args: - session: Snowpark session; needs CREATE EXPERIMENT on the schema, or - MODIFY on an existing experiment. - connector: a ``SnowflakeConnector`` whose identity opens the session - instead. + source: a ``SnowflakeConnector``, whose identity opens the session + and whose Snowflake Model Registry ``kumo.MLOps(registry=...)`` + registers models in; or a Snowpark session, which tracks only. + It needs CREATE EXPERIMENT on the schema, or MODIFY on an + existing experiment. experiment_name: experiment for the fit run, created if missing. database: `None` uses the session's current database. - schema: `None` uses the session's current schema, or ``PUBLIC`` when - `database` is given or the session has none. + schema_name: `None` uses the session's current schema, or ``PUBLIC`` + when `database` is given or the session has none. run_name: name of the fit run. `None` lets Snowflake pick one; a - taken name gets a suffix from the kumo run id. + taken name gets a random suffix. tags: extra tags for the fit run. native_model_lineage: keep the run open as a context for the whole fit, so models registered meanwhile are linked to it. @@ -54,26 +56,32 @@ class SnowflakeTracking: by name), once it has started. """ + platform = 'snowflake' + run_url = None + def __init__( self, - session: Any = None, + source: Any, *, - connector: Any = None, experiment_name: str, database: Optional[str] = None, - schema: Optional[str] = None, + schema_name: Optional[str] = None, run_name: Optional[str] = None, tags: Optional[Mapping[str, str]] = None, native_model_lineage: bool = False, ) -> None: + from kumo.engine.connector.base import Connector + self._module = import_provider('snowflake.ml.experiment', 'snowflake') - if (session is None) == (connector is None): - raise ValueError('SnowflakeTracking takes a session or a ' - 'connector.') - self.session = session if session is not None else connector.session() + if source is None: + raise ValueError('SnowflakeTracking takes a SnowflakeConnector ' + 'or a Snowpark session.') + self.connector = source if isinstance(source, Connector) else None + self.session = (source.session() + if self.connector is not None else source) self.experiment_name = experiment_name self.database = database - self.schema = schema + self.schema = schema_name self.run_name = run_name self.tags = dict(tags or {}) self.native_model_lineage = native_model_lineage @@ -92,6 +100,17 @@ def _reset(self) -> None: self._dropped_metrics: Set[str] = set() self._dropped_params = 0 + def load_model(self, location: Any, **kwargs: Any) -> Any: + r"""The `Trainer` of the Snowflake Model Registry version + ``location`` names, through this tracker's connector.""" + from kumo.engine.trainer.trainer import Trainer + + if self.connector is None: + raise ValueError('Loading a registered model needs ' + 'SnowflakeTracking(connector, ...).') + return Trainer.load_from_connector(self.connector, location=location, + **kwargs) + def check(self, registry: Any = None) -> None: r"""Raises before any work unless the session can see the experiment's schema (one metadata query, no warehouse).""" @@ -135,8 +154,7 @@ def start(self, record: RunRecord) -> None: database_name=database, schema_name=schema or 'PUBLIC') self._run = tracking.start_run( - self._new_run_name(tracking, experiment.name, - record.tags['kumo.run_id'])) + self._new_run_name(tracking, experiment.name)) self.run_id = str(self._run.name) if self.native_model_lineage: self._run_context.enter_context(self._run) @@ -191,8 +209,8 @@ def finish(self, record: Optional[RunRecord], self._run_context.close() self._reset() - def _new_run_name(self, tracking: Any, experiment: Any, - run_id: str) -> Optional[str]: + def _new_run_name(self, tracking: Any, + experiment: Any) -> Optional[str]: if self.run_name is None: return None # start_run(name) resumes, or raises on, any run matching name as a @@ -203,7 +221,7 @@ def _new_run_name(self, tracking: Any, experiment: Any, for row in tracking._sql_client.show_runs_in_experiment( experiment_name=experiment) ] - short = run_id.rsplit('-', 1)[-1] + short = uuid.uuid4().hex for name in (self.run_name, _suffixed(self.run_name, short[:8]), _suffixed(self.run_name, short)): if not _taken(name, existing): diff --git a/packages/kumo-engine/kumo/engine/trainer/_provenance.py b/packages/kumo-engine/kumo/engine/trainer/_provenance.py index 73e54e499..72c83129e 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_provenance.py +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -32,7 +32,7 @@ def write_provenance( warm_start: Optional[ModelManifest], started_at: datetime, fit_mode: Optional[str], - tracked_run_id: Optional[str] = None, + tracking: Optional[TrackingRef] = None, dependencies: Optional[List[Dependency]] = None, ) -> None: from kumo_core.artifacts import PROVENANCE_FILE, add_bundle_file @@ -54,8 +54,7 @@ def write_provenance( seconds=round((finished_at - started_at).total_seconds(), 3), ), fit_mode=fit_mode, - tracking=None if tracked_run_id is None else TrackingRef( - run_id=tracked_run_id), + tracking=tracking, dependencies=dependencies, **stamp(), )) diff --git a/packages/kumo-engine/kumo/engine/trainer/_run_spec.py b/packages/kumo-engine/kumo/engine/trainer/_run_spec.py index 7819797e2..1a916aad4 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_run_spec.py +++ b/packages/kumo-engine/kumo/engine/trainer/_run_spec.py @@ -17,7 +17,6 @@ class FitContext: r"""What a fit trains, known once its model plan is resolved.""" - run_id: str query: str model_plan: Mapping[str, Any] inputs: Mapping[str, Mapping[str, Optional[str]]] = field( @@ -28,7 +27,6 @@ def fit_spec(context: FitContext, *, tags: Optional[Mapping[str, Any]] = None) -> RunSpec: return RunSpec( kind='fit', - run_id=context.run_id, params=dict(_flatten(context.model_plan)), tags=dict(tags or {}), inputs=dict(context.inputs), diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 594303214..94ad871f1 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -66,7 +66,7 @@ with_text_encoders, ) from kumo.engine._events import Event, EventEmitter, query_sha256 -from kumo.engine._manifest import Dependency, ModelManifest, stub +from kumo.engine._manifest import Dependency, ModelManifest, TrackingRef, stub from kumo.engine._mlops import full_error, warn_failure from kumo.engine._predictor import TRAINER_PREDICTOR from kumo.engine._scratch import new_run_id, trainer_scratch_dir @@ -78,7 +78,12 @@ from kumo.engine.pquery.prediction_table import PredictionTable from kumo.engine.pquery.training_table import TrainingTable from kumo.engine.trainer._run_spec import FitContext, fit_spec, fitted_spec -from kumo.engine.trainer._scoring import ScoringDefaults, read_scoring_defaults +from kumo.engine.trainer._scoring import ( + ScoringDefaults, + normalized_scoring_defaults, + read_scoring_defaults, + write_scoring_defaults, +) from kumo.engine.trainer._spec import ( HOLDOUT_META_ROLE, HOLDOUT_ROLE, @@ -145,6 +150,13 @@ def snapshot_id(self) -> Optional[str]: return None return read_training_spec(self.bundle_dir).snapshot_id + @property + def tracking(self) -> Optional[TrackingRef]: + """The experiment tracker's run of the fit: its ``platform``, + ``run_id`` and ``url``; None when the fit was not tracked.""" + manifest = self.manifest + return None if manifest is None else manifest.tracking + def predict(self, *args: Any, **kwargs: Any) -> 'PredictionResult': """Score with this model, as :meth:`Trainer.predict`.""" return self._model('predict').predict(*args, **kwargs) @@ -158,14 +170,11 @@ def with_scoring_defaults(self, defaults: ScoringDefaults) -> 'TrainingResult': """A copy of this model whose ``predict`` applies ``defaults`` when a call omits them; a new bundle, whose manifest names this one in - ``parents.derived_from``.""" + ``parents.derived_from``. To register a model with its defaults, + pass ``fit(scoring_defaults=...)`` instead.""" from kumo_core.artifacts import bundle_task_type from kumo.engine.trainer._provenance import record_derivation - from kumo.engine.trainer._scoring import ( - normalized_scoring_defaults, - write_scoring_defaults, - ) if self._trainer is None: raise RuntimeError('This result has no model; use Trainer.load on ' @@ -440,13 +449,21 @@ def _require_publisher(export_config: ModelOutputConfig) -> None: def _mlflow_registry(settings: Any) -> Optional[ModelRegistryLocation]: - registry = None if settings is None else settings.registry - return registry if isinstance(registry, ModelRegistryLocation) else None + if settings is None or not _logs_models(settings.tracking): + return None + return settings.registry def _connector_registry(settings: Any) -> Optional[ModelOutputConfig]: - registry = None if settings is None else settings.registry - return registry if isinstance(registry, ModelOutputConfig) else None + if (settings is None or settings.registry is None + or _logs_models(settings.tracking)): + return None + return ModelOutputConfig(settings.tracking.connector, + location=settings.registry) + + +def _logs_models(tracking: Any) -> bool: + return callable(getattr(tracking, 'log_model', None)) def _publish(bundle_dir: Path, bundle_id: Optional[str], @@ -599,6 +616,7 @@ def fit( *, snapshot: GraphSnapshot, export_config: Optional[ModelOutputConfig] = None, + scoring_defaults: Optional[ScoringDefaults] = None, warm_start: Optional[Self] = None, mlops: Optional[Any] = None, callbacks: Sequence[Callable[[Event], None]] = (), @@ -622,6 +640,9 @@ def fit( generate_training_table(snapshot)`. snapshot: graph data and features to train against. export_config: if given, publishes the trained model bundle. + scoring_defaults: prediction arguments the bundle applies when + ``predict`` omits them, such as a classification threshold; + stored, registered and served with the model. warm_start: a fitted trainer or one restored with ``load`` or ``load_from_connector``. Reuses its selected weights, optimizer state, fitted encoders, architecture and neighbor sampling; @@ -653,7 +674,6 @@ def fit( session = TrackingSession( None if settings is None else settings.tracking, warn=warn) events = EventEmitter( - run_id, [*callbacks, session.on_event] if session.enabled else callbacks) fit_start = time.monotonic() try: @@ -667,6 +687,7 @@ def fit( try: result = self._fit(train_table, snapshot=snapshot, export_config=export_config, + scoring_defaults=scoring_defaults, warm_start=warm_start, run_id=run_id, events=events, tracking=session, deps=Dependencies(settings=settings)) @@ -702,6 +723,7 @@ def _fit( *, snapshot: GraphSnapshot, export_config: Optional[ModelOutputConfig], + scoring_defaults: Optional[ScoringDefaults], warm_start: Optional[Self], run_id: str, events: EventEmitter, @@ -764,6 +786,9 @@ def _fit( training_table_generation_plan=train_table.generation_plan, search=AutoMLSearchOptions(), ) + if scoring_defaults is not None: + scoring_defaults = normalized_scoring_defaults( + scoring_defaults, TaskType(workflow.query.task_type)) if warm_start_artifacts is not None: from kumo_core.artifacts import bundle_task_type @@ -842,7 +867,6 @@ def _fit( raise ValueError('The test split is empty: shorten the query ' 'window or add data.') tracking.start(lambda: fit_spec(FitContext( - run_id=run_id, query=train_table.query, model_plan=to_json_dict(materialized.model_plan), inputs=fit_inputs(train_table.manifest), @@ -905,6 +929,7 @@ def _fit( graph_config=snapshot.graph_config, training_table_generation=train_table.generation_plan, export_config=export_config, + scoring_defaults=scoring_defaults, model_plan=effective_plan, provenance=dict( train_table=train_table, @@ -912,7 +937,7 @@ def _fit( warm_start._training_result.bundle_dir, warm_start_bundle_id), started_at=started_at, fit_mode=fit_mode, - tracked_run_id=run_id if tracking.enabled else None), + tracking=tracking.pointer()), deps=deps, ) @@ -928,6 +953,7 @@ def _finish_fit( model_plan: Optional[ModelPlan] = None, provenance: Optional[Mapping[str, Any]] = None, deps: Optional[Dependencies] = None, + scoring_defaults: Optional[ScoringDefaults] = None, ) -> TrainingResult: """Commit durable fit outputs before making this trainer usable.""" from kumo_core.artifacts import bundle_id_from_manifest @@ -943,6 +969,11 @@ def _finish_fit( training_table_generation=training_table_generation, graph_config=graph_config, ) + if scoring_defaults is not None: + write_scoring_defaults(training_result.bundle_dir, + scoring_defaults) + training_result = replace(training_result, + scoring_defaults=scoring_defaults) training_result = _bundle_holdout(training_result) try: training_result = replace( @@ -1308,12 +1339,20 @@ def load_from_connector( from trusted publishers; checkpoint deserialization happens at predict. ``accelerator`` selects the device for this instance (default: ``auto``). A `kumo.MLOps` ``model_registry`` reads ``path`` under its artifact - root (``models/``). + root (``models/``) and ``location`` from the registry its + tracker registers in. """ from kumo.engine._mlops import MLOps, on_root - if isinstance(model_registry, MLOps) and path is None: - raise TypeError('Loading from a kumo.MLOps takes path=.') + if isinstance(model_registry, MLOps) and location is not None: + load = getattr(model_registry.tracking, 'load_model', None) + if path is not None or load is None: + raise TypeError( + 'Loading a registered model from a kumo.MLOps takes ' + 'location= and its tracking=MlflowTracking(...) or ' + 'SnowflakeTracking(connector, ...).') + return load(location, embedding_models=embedding_models, + accelerator=accelerator) if path is not None: model_registry, path = on_root(model_registry, path) address = ModelOutputConfig(model_registry, path=path, diff --git a/packages/kumo-engine/test/test_fit_events.py b/packages/kumo-engine/test/test_fit_events.py index c9efff26c..08a22729c 100644 --- a/packages/kumo-engine/test/test_fit_events.py +++ b/packages/kumo-engine/test/test_fit_events.py @@ -78,7 +78,6 @@ def broken(event): 'trial_started', 'epoch_ended', 'epoch_ended', 'trial_ended', 'fit_ended', ] - assert len({event.run_id for event in events}) == 1 assert set(events[0].params) == {'snapshot_id', 'query_sha256'} assert [(e.trial, e.epoch) for e in events if e.name == 'epoch_ended' ] == [(1, 1), (1, 2), (2, 1), (2, 2)] @@ -124,7 +123,7 @@ def test_event_metrics_keep_finite_numbers(): from kumo.engine._events import EventEmitter received = [] - EventEmitter('run', [received.append]).emit( + EventEmitter([received.append]).emit( 'epoch_ended', metrics={'val_loss': 'n/a', 'val_auc': float('nan'), 'val_mae': 0.5}) diff --git a/packages/kumo-engine/test/test_mlflow_model.py b/packages/kumo-engine/test/test_mlflow_model.py index efd932dec..2f3cff81d 100644 --- a/packages/kumo-engine/test/test_mlflow_model.py +++ b/packages/kumo-engine/test/test_mlflow_model.py @@ -65,7 +65,6 @@ def test_fit_logs_registers_and_serves_the_model(table_and_snapshot, assert {k: v for k, v in parent.data.tags.items() if k.startswith('kumo.')} == { 'kumo.kind': 'fit', - 'kumo.run_id': parent.data.tags['kumo.run_id'], 'kumo.bundle_id': result.id, } assert result.registered_model == PublishedModel( @@ -95,6 +94,9 @@ def test_fit_logs_registers_and_serves_the_model(table_and_snapshot, assert list(served) == ['ENTITY', 'TARGET_PRED'] assert len(served) == len(range(0, 400, 7)) assert load_model('models:/m@champion').result.id == result.id + assert kumo.Trainer.load_from_connector( + cfg, location=kumo.ModelRegistryLocation('m', alias='champion') + ).result.id == result.id original = mlflow.models.Model.load @@ -289,7 +291,7 @@ def finish(self, record, error): for fail in (False, True): tracker = Tracker(fail) session = TrackingSession(tracker) - session.start(lambda: fit_spec(FitContext('r', 'q', {}))) + session.start(lambda: fit_spec(FitContext('q', {}))) with warnings.catch_warnings(record=True) as caught: warnings.simplefilter('always') _, error = session.log_model(result.bundle_dir, None, None) @@ -329,7 +331,7 @@ def finish(self, record, error): tracker = Tracker() session = TrackingSession(tracker) - session.start(lambda: fit_spec(FitContext('r', 'q', {}))) + session.start(lambda: fit_spec(FitContext('q', {}))) if stall == 'dead thread': session._submit(lambda: (_ for _ in ()).throw(SystemExit())) session._worker.join(5) @@ -418,7 +420,7 @@ def finish(self, record, error): pass session = TrackingSession(Tracker()) - session.start(lambda: fit_spec(FitContext('r', 'q', {}))) + session.start(lambda: fit_spec(FitContext('q', {}))) deadline = time.monotonic() + 10 while session._tracker is not None and time.monotonic() < deadline: time.sleep(0.01) diff --git a/packages/kumo-engine/test/test_mlflow_target.py b/packages/kumo-engine/test/test_mlflow_target.py index c18a2df68..27be764b0 100644 --- a/packages/kumo-engine/test/test_mlflow_target.py +++ b/packages/kumo-engine/test/test_mlflow_target.py @@ -40,30 +40,34 @@ def _connector(identity): def test_no_target_reads_mlflows_own_uris_once(tracking_uri): - target = resolve(None, None) + target = resolve(None) assert target.tracking_uri == tracking_uri assert not target.explicit - assert resolve(None, 'databricks').registry_uri == 'databricks-uc' - assert resolve(None, 'databricks://prod').registry_uri == ( + assert target.platform == 'mlflow' and target.run_url('1', 'r') is None + assert resolve('databricks').registry_uri == 'databricks-uc' + assert resolve('databricks').platform == 'databricks' + assert resolve('databricks://prod').registry_uri == ( 'databricks-uc://prod') - assert resolve(None, 'http://mlflow:5000').registry_uri == ( - 'http://mlflow:5000') - with pytest.raises(ValueError, match='not both'): - resolve(SageMakerMLflow(ARN), 'http://x') + server = resolve('http://mlflow:5000') + assert server.registry_uri == 'http://mlflow:5000' + assert server.run_url('1', 'r') == ( + 'http://mlflow:5000/#/experiments/1/runs/r') + assert resolve(ARN).platform == 'sagemaker' + assert resolve(server) is server def test_ambient_identities_set_no_credentials(monkeypatch, tracking_uri): monkeypatch.setenv('DATABRICKS_TOKEN', 'mine') monkeypatch.setenv('MLFLOW_TRACKING_URI', tracking_uri) - target = resolve(_connector(DatabricksIdentity()), None) + target = resolve(_connector(DatabricksIdentity())) before = dict(os.environ) with target.scope(): assert dict(os.environ) == before assert mlflow.get_tracking_uri() == tracking_uri assert not target.explicit - for target in (resolve(SageMakerMLflow(ARN, S3Identity()), None), - resolve(SageMakerMLflow(ARN), None)): + for target in (resolve(SageMakerMLflow(ARN, S3Identity())), + resolve(SageMakerMLflow(ARN))): with target.scope(): assert mlflow.get_tracking_uri() == ARN assert target.artifact_repository('s3://b/run') is None @@ -85,7 +89,7 @@ def test_a_databricks_service_principal_lends_a_token_for_each_call( target = resolve( _connector( DatabricksIdentity(host='https://dbc.example', client_id='sp', - client_secret='secret')), None) + client_secret='secret'))) previous = mlflow.get_tracking_uri() with target.scope(): @@ -102,6 +106,8 @@ def test_a_databricks_service_principal_lends_a_token_for_each_call( assert 'DATABRICKS_TOKEN' not in os.environ assert (target.tracking_uri, target.registry_uri) == ('databricks', 'databricks-uc') + assert target.run_url('7', 'r') == ( + 'https://dbc.example/ml/experiments/7/runs/r') def test_an_aws_access_key_signs_sagemaker_calls_and_uploads( @@ -112,7 +118,7 @@ def test_an_aws_access_key_signs_sagemaker_calls_and_uploads( monkeypatch.setenv('SAGEMAKER_MLFLOW_ASSUME_ROLE_ARN', 'arn:aws:iam::x') target = resolve( SageMakerMLflow(ARN, S3Identity(access_key_id='AKIA', - secret_access_key='s3cret')), None) + secret_access_key='s3cret'))) with target.scope(): assert mlflow.get_tracking_uri() == ARN @@ -172,11 +178,14 @@ def test_an_unreachable_server_fails_when_mlops_is_built( from kumo.engine.integrations.mlflow import MlflowTracking monkeypatch.setenv('MLFLOW_HTTP_REQUEST_MAX_RETRIES', '0') - tracking = MlflowTracking(tracking_uri='http://127.0.0.1:9', + tracking = MlflowTracking(target='http://127.0.0.1:9', experiment_name='e') with pytest.raises(RuntimeError, match='MLflow at http://127.0.0.1:9 ' 'is not usable with these credentials'): kumo.MLOps(str(tmp_path), tracking=tracking) + with pytest.raises(ValueError, match='pass experiment_name'): + kumo.MLOps(str(tmp_path), + tracking=MlflowTracking(target='http://127.0.0.1:9')) def test_a_missing_unity_catalog_schema_fails_when_mlops_is_built( @@ -255,7 +264,7 @@ def test_registering_a_registered_bundle_reuses_its_version( tracking_uri, caplog): from kumo.engine.integrations.mlflow import _register - target = resolve(None, None) + target = resolve(None) client = target.client() run = client.create_run('0') info = SimpleNamespace(model_uri=f'runs:/{run.info.run_id}/model', diff --git a/packages/kumo-engine/test/test_mlflow_tracking.py b/packages/kumo-engine/test/test_mlflow_tracking.py index 9b459e264..920e92272 100644 --- a/packages/kumo-engine/test/test_mlflow_tracking.py +++ b/packages/kumo-engine/test/test_mlflow_tracking.py @@ -116,11 +116,12 @@ def test_fit_logs_a_run_per_fit_and_per_trial(table_and_snapshot, tmp_path, assert client.get_experiment(parent.info.experiment_id).name == 'ents' assert parent.info.status == 'FINISHED' assert {k for k in parent.data.tags if k.startswith('kumo.')} == { - 'kumo.kind', 'kumo.run_id', 'kumo.bundle_id', 'kumo.dropped_params'} + 'kumo.kind', 'kumo.bundle_id', 'kumo.dropped_params'} assert parent.data.tags['kumo.bundle_id'] == result.id stored = tmp_path / 'root' / 'models' / result.id / PROVENANCE_FILE assert json.loads(stored.read_text())['tracking'] == { - 'run_id': parent.data.tags['kumo.run_id']} + 'platform': 'mlflow', 'run_id': parent.info.run_id} + assert result.tracking.run_id == parent.info.run_id logged = client.download_artifacts(parent.info.run_id, PROVENANCE_FILE, str(tmp_path)) assert json.loads(Path(logged).read_text())['id'] == result.id @@ -231,7 +232,7 @@ def test_registry_model_names_are_checked_when_mlops_is_built( for name in good: check_model_name(name, uri) tracking = MlflowTracking(log_system_metrics=False, - tracking_uri=f'sqlite:///{tmp_path}/m.db', + target=f'sqlite:///{tmp_path}/m.db', experiment_name='e') tracking.target.registry_uri = uri for name in bad: diff --git a/packages/kumo-engine/test/test_mlops.py b/packages/kumo-engine/test/test_mlops.py index ea9974e7d..aeca7d9e9 100644 --- a/packages/kumo-engine/test/test_mlops.py +++ b/packages/kumo-engine/test/test_mlops.py @@ -36,26 +36,29 @@ def test_mlops_needs_an_artifact_root_it_can_write_files_to(tmp_path): kumo.MLOps(artifact_root=str(tmp_path), on_error='ignore') -def test_a_connector_registry_checks_the_name_when_mlops_is_built(tmp_path): +def test_the_tracker_registers_and_its_registry_checks_the_name(tmp_path): def check_model_location(connector, location): if '-' in location.model_name: raise ValueError('bad name') - connector = SimpleNamespace(sink_backend=SimpleNamespace( - check_model_location=check_model_location)) + tracker = SimpleNamespace(connector=SimpleNamespace( + sink_backend=SimpleNamespace( + check_model_location=check_model_location))) good = kumo.ModelRegistryLocation('USERS_CHURN', alias='candidate') - kumo.MLOps(artifact_root=str(tmp_path), - registry=kumo.ModelOutputConfig(connector, location=good)) + kumo.MLOps(str(tmp_path), tracking=tracker, registry=good) with pytest.raises(ValueError, match='bad name'): - kumo.MLOps(artifact_root=str(tmp_path), - registry=kumo.ModelOutputConfig( - connector, - location=kumo.ModelRegistryLocation('users-churn'))) + kumo.MLOps(str(tmp_path), tracking=tracker, + registry=kumo.ModelRegistryLocation('users-churn')) with pytest.raises(ValueError, match='model alias'): - kumo.MLOps(artifact_root=str(tmp_path), - registry=kumo.ModelOutputConfig( - connector, location=kumo.ModelRegistryLocation( - 'USERS_CHURN', alias='LATEST'))) + kumo.MLOps(str(tmp_path), tracking=tracker, + registry=kumo.ModelRegistryLocation('USERS_CHURN', + alias='LATEST')) + with pytest.raises(ValueError, match='registers through the tracker'): + kumo.MLOps(str(tmp_path), registry=good) + with pytest.raises(TypeError, match='ModelRegistryLocation'): + kumo.MLOps(str(tmp_path), tracking=tracker, + registry=kumo.ModelOutputConfig(tracker.connector, + location=good)) def test_a_root_is_a_uri_or_a_connector_and_path(tmp_path): diff --git a/packages/kumo-engine/test/test_model_agnostic_mlops.py b/packages/kumo-engine/test/test_model_agnostic_mlops.py index 648c1b66d..a73a30061 100644 --- a/packages/kumo-engine/test/test_model_agnostic_mlops.py +++ b/packages/kumo-engine/test/test_model_agnostic_mlops.py @@ -75,6 +75,8 @@ def on_trial(self, record, event): def log_model(self, bundle_dir, manifest, snapshot, registry): self.calls.append(('log_model', (bundle_dir, manifest.id, snapshot))) + return registry and PublishedModel(name=registry.model_name, + version='2', uri='fake://TOY/2') def finish(self, record, error): self.calls.append(('finish', record)) @@ -83,48 +85,62 @@ def finish(self, record, error): def test_any_run_is_tracked_from_its_spec(): tracker = _Tracker() session = TrackingSession(tracker) - spec = RunSpec(kind='custom', run_id='loop-1', params={'lr': 0.1}, + spec = RunSpec(kind='custom', params={'lr': 0.1}, inputs={'users': {'connector': 'S3'}}) session.start(lambda: spec) - session.on_event(Event('trial_started', 'loop-1', 0.0, trial=1, + session.on_event(Event('trial_started', 0.0, trial=1, params={'attempt': 1})) session.finish(lambda: replace(spec, metrics={'auc': 0.9}, parents={'bundle_id': 'b'}), None) session._worker.join(5) started, trial, finished = (record for _, record in tracker.calls) assert started == run_record(spec) - assert trial.tags['kumo.parent_run'] == 'loop-1' + assert trial.name == 'trial-1-attempt-1' + assert trial.tags == {'kumo.kind': 'trial'} assert finished.metrics == {'auc': 0.9} - assert finished.tags == {'kumo.kind': 'custom', 'kumo.run_id': 'loop-1', - 'kumo.bundle_id': 'b'} + assert finished.tags == {'kumo.kind': 'custom', 'kumo.bundle_id': 'b'} -def test_any_bundle_is_stored_logged_and_registered(tmp_path): +class _RegistryTracker(_Tracker): + log_model = None + + def __init__(self, connector): + super().__init__() + self.connector = connector + + +@pytest.mark.parametrize('logs', [False, True]) +def test_any_bundle_is_stored_and_registered_by_its_tracker(tmp_path, logs): bundle, manifest = _toy_bundle(tmp_path / 'bundle') published = [] sink = SimpleNamespace( publish_model=lambda connector, bundle_dir, **kw: published.append( (bundle_dir, kw)) or PublishedModel(name='TOY', version='1', uri='fake://TOY/1')) - registry = kumo.ModelOutputConfig(SimpleNamespace(sink_backend=sink), - location=ModelRegistryLocation('TOY')) - settings = kumo.MLOps(str(tmp_path / 'root'), registry=registry) - tracker = _Tracker() + tracker = (_Tracker() if logs else _RegistryTracker( + SimpleNamespace(sink_backend=sink))) + settings = kumo.MLOps(str(tmp_path / 'root'), tracking=tracker, + registry=ModelRegistryLocation('TOY')) session = TrackingSession(tracker) - session.start(lambda: RunSpec(kind='custom', run_id='loop-1')) + session.start(lambda: RunSpec(kind='custom')) registered, error = _store_and_register( bundle, manifest, settings, session=session, snapshot='snap', metrics={'auc': 0.9}) - assert error is None and registered.version == '1' + assert error is None stored = tmp_path / 'root' / 'models' / manifest.id assert (stored / 'weight.txt').read_text() == '2.0' - assert ('log_model', (bundle, manifest.id, 'snap')) in tracker.calls - assert published == [(str(bundle), { - 'path': None, 'location': ModelRegistryLocation('TOY'), - 'metadata': {'kumo_bundle_id': manifest.id}, 'metrics': {'auc': 0.9} - })] + if logs: + assert registered.version == '2' and published == [] + assert ('log_model', (bundle, manifest.id, 'snap')) in tracker.calls + else: + assert registered.version == '1' + assert published == [(str(bundle), { + 'path': None, 'location': ModelRegistryLocation('TOY'), + 'metadata': {'kumo_bundle_id': manifest.id}, + 'metrics': {'auc': 0.9} + })] session.finish(None, None) diff --git a/packages/kumo-engine/test/test_scoring_defaults.py b/packages/kumo-engine/test/test_scoring_defaults.py index 6ccfdf08e..fd8812113 100644 --- a/packages/kumo-engine/test/test_scoring_defaults.py +++ b/packages/kumo-engine/test/test_scoring_defaults.py @@ -137,3 +137,27 @@ def test_unlisted_or_unreadable_defaults_are_ignored(fitted, tmp_path): assert ignored.result.scoring_defaults == kumo.ScoringDefaults() with pytest.raises(ValueError, match='must be specified'): ignored.predict(table, snapshot=snapshot) + + +def test_a_fit_stores_its_defaults_in_the_bundle(fitted, tmp_path): + from kumoapi.model_plan import TrainingTableGenerationPlan + + from kumo.engine.trainer.trainer import TrainingResult + + trainer, table, snapshot = fitted + first = trainer.result + defaults = kumo.ScoringDefaults(binary_classification_threshold=0.3) + refit = kumo.Trainer(trainer.model_plan) + result = refit._finish_fit( + TrainingResult({}, dict(first.artifact_paths), + _copy(trainer, tmp_path / 'bundle')), + predictive_query=trainer._training_query, + snapshot_id=snapshot.id, export_config=None, + training_table_generation=TrainingTableGenerationPlan(), + graph_config=snapshot.graph_config, scoring_defaults=defaults) + assert result.scoring_defaults == defaults and result.id != first.id + loaded = kumo.Trainer.load(result.bundle_dir) + assert loaded.result.scoring_defaults == defaults + pd.testing.assert_frame_equal( + _frame(loaded, table, snapshot), + _frame(first, table, snapshot, binary_classification_threshold=0.3)) diff --git a/packages/kumo-engine/test/test_snowflake_tracking.py b/packages/kumo-engine/test/test_snowflake_tracking.py index aebe648c5..7cbbacec9 100644 --- a/packages/kumo-engine/test/test_snowflake_tracking.py +++ b/packages/kumo-engine/test/test_snowflake_tracking.py @@ -3,6 +3,7 @@ # SPDX-License-Identifier: Apache-2.0 import inspect +import re import threading import time from unittest import mock @@ -81,7 +82,7 @@ def _replay(tracking, result=RESULT, error=None, during=None): def test_native_model_lineage_holds_the_run_open_during_the_fit(fake): lineage = [] _replay( - _tracking(database='ml', schema='runs', run_name='fit', + _tracking(database='ml', schema_name='runs', run_name='fit', native_model_lineage=True), during=lambda: lineage.extend(fake._instance.lineage_stack)) @@ -154,8 +155,8 @@ def test_an_existing_run_is_never_resumed(fake, taken, caplog): tracking = _tracking(session, run_name=run_name) _replay(tracking) - new = f"{run_name}_{CONTEXT.run_id.rsplit('-', 1)[-1][:8]}".upper() - assert tracking.run_id == new + new = tracking.run_id + assert re.fullmatch(f'{run_name}_[0-9a-f]{{8}}'.upper(), new) assert set(user.runs) == {taken, new} assert user.runs[taken]['params'] == {} assert user.runs[new]['status'] == 'FINISHED' @@ -240,15 +241,43 @@ def sql(self, statement): def test_a_connector_opens_the_session_and_the_schema_is_checked_first( fake): import kumo.engine as kumo + from kumo.engine.connector.base import Connector session = _SchemaSession([('RUNS', )]) - tracking = _tracking(None, connector=mock.Mock(session=lambda: session)) + connector = mock.Mock(spec=Connector, session=lambda: session) + tracking = _tracking(connector) tracking.check() - assert tracking.session is session + assert tracking.session is session and tracking.connector is connector assert session.statements == ["SHOW SCHEMAS LIKE 'RUNS' IN DATABASE DB"] tracking.session = _SchemaSession([]) with pytest.raises(ValueError, match='needs schema DB.RUNS'): kumo.MLOps('/tmp/root', tracking=tracking) - with pytest.raises(ValueError, match='a session or a connector'): + with pytest.raises(ValueError, match='or a Snowpark session'): _tracking(None) + + +def test_mlops_loads_a_registered_model_through_the_tracker( + fake, monkeypatch, tmp_path): + import kumo.engine as kumo + from kumo.engine.connector.base import Connector + + loaded = [] + original = kumo.Trainer.load_from_connector.__func__ + monkeypatch.setattr( + kumo.Trainer, 'load_from_connector', + classmethod(lambda cls, source, **kwargs: loaded.append( + (source, kwargs)) if not isinstance(source, kumo.MLOps) else + original(cls, source, **kwargs))) + session = _SchemaSession([('RUNS', )]) + connector = mock.Mock(spec=Connector, session=lambda: session) + tracking = _tracking(connector) + cfg = kumo.MLOps(str(tmp_path), tracking=tracking) + location = kumo.ModelRegistryLocation('CHURN', alias='champion') + kumo.Trainer.load_from_connector(cfg, location=location) + assert loaded == [(connector, { + 'location': location, 'embedding_models': None, + 'accelerator': 'auto' + })] + with pytest.raises(ValueError, match='SnowflakeTracking\\(connector'): + _tracking().load_model(location) diff --git a/packages/kumo-engine/test/test_tracking_contract.py b/packages/kumo-engine/test/test_tracking_contract.py index bcf9cf3f0..0d7bedbf1 100644 --- a/packages/kumo-engine/test/test_tracking_contract.py +++ b/packages/kumo-engine/test/test_tracking_contract.py @@ -23,8 +23,8 @@ from kumo.engine.tracking import Event, RunRecord from kumo.engine.trainer._run_spec import FitContext, fit_spec, fitted_spec -RUN = 'trainer-fit-1' -USER_TAGS = {'team': 'ents', 'kumo.run_id': 'overridden'} +FIT = 'fit' +USER_TAGS = {'team': 'ents', 'kumo.kind': 'overridden'} QUERY = 'PREDICT ents.Y FOR EACH ents.ID' INPUTS = { 'ents': {'connector': 'S3', 'location': 's3://bucket/ents.parquet', @@ -36,7 +36,7 @@ 'version': '1790000000', 'version_kind': 'timestamp'}, } CONTEXT = FitContext( - run_id=RUN, query=QUERY, inputs=INPUTS, + query=QUERY, inputs=INPUTS, model_plan={ 'optimization': {'base_lr': [0.01, 0.001], 'max_epochs': 2}, 'training_job': {'tune_metric': 'mae'}, @@ -47,7 +47,7 @@ def _event(name, trial=None, epoch=None, metrics=None, **params): - return Event(name=name, run_id=RUN, time=0.0, trial=trial, epoch=epoch, + return Event(name=name, time=0.0, trial=trial, epoch=epoch, metrics=metrics or {}, params=params) @@ -125,7 +125,7 @@ def replay(tracker=None, events=EVENTS, result=RESULT, None if result is None else lambda: fitted_spec( session.spec, result, events[-1].metrics.get('seconds')), error) records = { - record.tags['kumo.run_id']: record + record.name: record for hook, record in recorder.calls if hook == 'on_trial' } spec = fit_spec(CONTEXT, tags=USER_TAGS) @@ -133,7 +133,7 @@ def replay(tracker=None, events=EVENTS, result=RESULT, spec = fitted_spec(spec, result, events[-1].metrics.get('seconds')) if 'error' in events[-1].params: spec = replace(spec, error=str(events[-1].params['error'])) - records[RUN] = run_record(spec) + records[FIT] = run_record(spec) return records @@ -196,7 +196,8 @@ def differ(what, have, want): and k not in (DROPPED_PARAMS, DROPPED_METRICS)), sorted(k for k in record.tags if k.startswith('kumo.'))) if adapter.nested_runs: - differ('parent', seen.parent, record.tags.get('kumo.parent_run')) + differ('parent', seen.parent, + FIT if record.tags['kumo.kind'] == 'trial' else None) if adapter.artifacts: differ('artifacts', dict(seen.artifacts), dict(record.artifacts)) return found @@ -215,7 +216,11 @@ def observe(): client = mlflow.MlflowClient() runs = client.search_runs( [e.experiment_id for e in client.search_experiments()]) - kumo_ids = {r.info.run_id: r.data.tags['kumo.run_id'] for r in runs} + kumo_ids = { + r.info.run_id: + FIT if r.data.tags['kumo.kind'] == 'fit' else r.info.run_name + for r in runs + } def history(run_id, key): return tuple((m.step, m.value) for m in sorted( @@ -280,7 +285,7 @@ def bucket(key): runs['steps'][name] = (*runs['steps'].get(name, ()), (step, value)) for prefix, runs in names.items(): - observed[runs['tags']['kumo.run_id']] = Observed( + observed[prefix or FIT] = Observed( name=prefix, artifacts=run['artifacts'] if prefix is None else {}, **runs) @@ -300,8 +305,8 @@ def adapter(request) -> Adapter: def test_records_of_the_recorded_fit(): records = replay() - fit = records[RUN] - assert fit.tags == {'team': 'ents', 'kumo.kind': 'fit', 'kumo.run_id': RUN, + fit = records[FIT] + assert fit.tags == {'team': 'ents', 'kumo.kind': 'fit', 'kumo.bundle_id': BUNDLE} assert fit.params == { 'optimization.base_lr': '[0.01, 0.001]', @@ -319,17 +324,15 @@ def test_records_of_the_recorded_fit(): assert json.loads(fit.artifacts['model_plan.json']) == CONTEXT.model_plan assert fit.inputs == INPUTS - oom = records[f'{RUN}-trial-2-attempt-1'] + oom = records['trial-2-attempt-1'] assert oom == RunRecord( name='trial-2-attempt-1', params={'optimization.base_lr': '0.001', 'attempt': '1'}, tags={ 'kumo.kind': 'trial', - 'kumo.run_id': f'{RUN}-trial-2-attempt-1', - 'kumo.parent_run': RUN, 'kumo.status': 'failed', 'kumo.error': 'OutOfMemoryError: CUDA out of memory', }) - retry = records[f'{RUN}-trial-2-attempt-2'] + retry = records['trial-2-attempt-2'] assert retry.params['attempt'] == '2' assert retry.metrics == {'best_val_mae': 0.3, 'seconds': 4.0} assert retry.steps == {'train_loss': ((1, 0.6), (2, 0.4)), @@ -342,7 +345,7 @@ def test_adapter_records_the_same_runs(adapter, tmp_path): bundle_dir=tmp_path) expected = replay(adapter.make(), result=result) - assert expected[RUN].artifacts[PROVENANCE_FILE] == '{"bundle": "b"}' + assert expected[FIT].artifacts[PROVENANCE_FILE] == '{"bundle": "b"}' assert gaps(expected, adapter.observe(), adapter) == [] @@ -367,7 +370,7 @@ def test_adapter_records_a_failed_fit(adapter): expected = replay(adapter.make(), [*EVENTS[:-1], failed], None, ValueError('stop')) - assert expected[RUN].tags['kumo.error'] == 'ValueError: stop' + assert expected[FIT].tags['kumo.error'] == 'ValueError: stop' assert gaps(expected, adapter.observe(), adapter) == [] diff --git a/packages/kumo-engine/test/test_trainer_registry.py b/packages/kumo-engine/test/test_trainer_registry.py index e20b5f387..fc1fd6185 100644 --- a/packages/kumo-engine/test/test_trainer_registry.py +++ b/packages/kumo-engine/test/test_trainer_registry.py @@ -699,6 +699,20 @@ def test_invalid_export_alias_is_rejected_before_publishing( assert sink.published == {} +class _RegistryTracker: + def __init__(self, sink): + self.connector = _FakeConnector(sink) + + def start(self, record): + pass + + def on_trial(self, record, event): + pass + + def finish(self, record, error): + pass + + @pytest.mark.parametrize('sink,registered', [ (_AliasingSink(), 'V1'), (SimpleNamespace(publish_model=lambda *a, **k: 1 / 0), None), @@ -716,17 +730,16 @@ def test_mlops_registry_registers_after_the_fit(tmp_path: Path, monkeypatch, monkeypatch.setattr( Trainer, '_fit', lambda self, *a, **k: TrainingResult( {}, {}, bundle, bundle_id='b1')) - registry = kumo.ModelOutputConfig( - _FakeConnector(sink), - location=ModelRegistryLocation('CHURN_GNN', alias='prod')) + tracker = _RegistryTracker(sink) + registry = ModelRegistryLocation('CHURN_GNN', alias='prod') trainer = Trainer(ModelPlan()) with warnings.catch_warnings(record=True) as caught: warnings.simplefilter('always') result = trainer.fit( SimpleNamespace(query='PREDICT COUNT(*)'), snapshot=SimpleNamespace(id='s'), - mlops=kumo.MLOps(str(tmp_path / 'root'), registry=registry, - store=False)) + mlops=kumo.MLOps(str(tmp_path / 'root'), tracking=tracker, + registry=registry, store=False)) failures = [w for w in caught if 'Could not register' in str(w.message)] assert getattr(result.registered_model, 'version', None) == registered assert result.published_model is None @@ -740,8 +753,9 @@ def test_mlops_registry_registers_after_the_fit(tmp_path: Path, monkeypatch, trainer.fit( SimpleNamespace(query='PREDICT COUNT(*)'), snapshot=SimpleNamespace(id='s'), - mlops=kumo.MLOps(str(tmp_path / 'root'), registry=registry, - store=False, on_error='raise')) + mlops=kumo.MLOps(str(tmp_path / 'root'), tracking=tracker, + registry=registry, store=False, + on_error='raise')) assert isinstance(raised.value.__cause__, ZeroDivisionError) @@ -750,10 +764,10 @@ def test_mlops_registry_alias_is_checked_before_the_fit(tmp_path: Path, import kumo.engine as kumo monkeypatch.setattr(Trainer, '_fit', lambda *a, **k: 1 / 0) - registry = kumo.ModelOutputConfig( - _FakeConnector(_AliasingSink()), - location=ModelRegistryLocation('CHURN_GNN', alias='LATEST')) with pytest.raises(ValueError, match='model alias'): Trainer(ModelPlan()).fit( SimpleNamespace(query='q'), snapshot=SimpleNamespace(id='s'), - mlops=kumo.MLOps(str(tmp_path), registry=registry, store=False)) + mlops=kumo.MLOps( + str(tmp_path), tracking=_RegistryTracker(_AliasingSink()), + registry=ModelRegistryLocation('CHURN_GNN', alias='LATEST'), + store=False)) From 170a39973f0bb0f0f2dcd21129d69e57f77ea5ff Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 28 Oct 2026 22:11:47 +0000 Subject: [PATCH 021/278] Store, register and authenticate safely on every platform Store failures (snapshot, training table, model, asset) follow on_error like tracking and registration, and a model whose bundle or assets were not stored is never registered. Assets go to object stores without bucket-level calls, so prefix-scoped S3 rights suffice. A second MLflow registration into a Unity Catalog model reads search results whose tags are a method. result.export(cfg) registers a model through the MLOps tracker and returns the version that already holds its bundle. An ambient Databricks target points MLflow at the workspace for each call, so a profile works without environment variables, and a service principal's client ignores the token lent to MLflow on another thread. docs/mlops.md is restructured around a quickstart per platform, with the details in a reference section. Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_databricks/session.py | 3 + packages/kumo-engine/docs/mlops.md | 497 +++++++----------- packages/kumo-engine/kumo/engine/_deps.py | 19 +- packages/kumo-engine/kumo/engine/_mlops.py | 10 +- .../kumo-engine/kumo/engine/_snapshot_io.py | 9 +- packages/kumo-engine/kumo/engine/graph.py | 6 +- .../engine/integrations/_mlflow_target.py | 10 +- .../kumo/engine/integrations/mlflow.py | 57 +- .../kumo/engine/pquery/training_table.py | 12 +- .../kumo/engine/trainer/trainer.py | 66 ++- .../kumo-engine/test/test_mlflow_model.py | 14 +- .../kumo-engine/test/test_mlflow_target.py | 48 +- .../test/test_model_agnostic_mlops.py | 81 +++ .../kumo-engine/test/test_snapshot_export.py | 14 + 14 files changed, 484 insertions(+), 362 deletions(-) diff --git a/packages/kumo-databricks/src/kumo_databricks/session.py b/packages/kumo-databricks/src/kumo_databricks/session.py index b274cd234..ceb299d0b 100644 --- a/packages/kumo-databricks/src/kumo_databricks/session.py +++ b/packages/kumo-databricks/src/kumo_databricks/session.py @@ -91,10 +91,13 @@ def _client_for(principal: Principal) -> WorkspaceClient: f"Identity {principal.name!r} holds " f"{type(credentials).__name__} credentials; Databricks " f"authenticates a service principal with an OAuth secret.") + # Explicit: a DATABRICKS_TOKEN another thread lends MLflow must not + # make this client's configuration ambiguous. return WorkspaceClient( host=config.host, client_id=config.client_id, client_secret=_secret(credentials.client_secret), + auth_type='oauth-m2m', ) diff --git a/packages/kumo-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index 061a9a253..1553fef8b 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -1,404 +1,289 @@ # MLOps with Kumo Engine -## Turning MLOps on +Turn MLOps on with one `kumo.MLOps` and the normal `kumo.engine` calls store their snapshots, training tables and models under an artifact root, record fits in your experiment tracker, and register models in its registry. + +Install with `pip install "kumo-engine[mlflow]"` (add `s3`, `gcs`, `databricks`, `sagemaker` or `snowflake` as needed; Snowflake needs only `kumo-engine[snowflake]`). GPU, CUDA and Databricks Runtime notes are under [Install](#install). + +## Quickstart -`kumo.MLOps` says where artifacts go and what records them. Use it as a block, which notebooks and scheduled scripts alike can do, or pass it as `mlops=` on a single call: +Each block assumes a `graph` and a `pquery` (see the README). It trains, stores and registers a model, then scores with it. + +### Laptop or VM ```python import kumo.engine as kumo +from kumo.engine.integrations.mlflow import MlflowTracking # MLflow's own tracking URI, else ./mlruns -cfg = kumo.MLOps(artifact_root="s3://bucket/kumo") # or (connector, path), e.g. a volume or stage +cfg = kumo.MLOps("artifacts", tracking=MlflowTracking("churn"), + registry=kumo.ModelRegistryLocation("users-churn", alias="candidate")) with kumo.mlops(cfg): snapshot = graph.snapshot() + result = kumo.Trainer(pquery.suggest_model_plan(snapshot)).fit(pquery.generate_training_table(snapshot), + snapshot=snapshot, scoring_defaults=kumo.ScoringDefaults(binary_classification_threshold=0.5)) + scores = result.predict(snapshot).frame + print(result.registered_model.version, result.tracking.run_id, len(scores)) ``` -- **`artifact_root`** is required. It is a URI (`s3://`, `gs://`, `file://`, or a local path, including `/Volumes/...` on a Databricks cluster) or a `(connector, path)` pair for a store a URI cannot reach. -- **Volumes and stages:** `(DatabricksConnector(...), "/Volumes////kumo")` and `(SnowflakeConnector(...), "@DB.SCHEMA.STAGE/kumo")` work as roots once the installed kumo-databricks or kumo-snowflake sink holds files at those paths; with an older plugin, building `kumo.MLOps` raises. -- **`export`** names the artifact kinds stored under the root when they are produced: `snapshot`, `training_table` and `model` (default: all). -- **`tracking`** and **`registry`** record fits and register models. -- **`on_error`** says what a tracking or registration failure in `fit` does. `"warn"` (default): the call warns, naming the error, and returns as usual. `"raise"`, for scheduled jobs: the call finishes and stores its result first, then raises a `RuntimeError` whose `__cause__` is the failure; after `fit`, `trainer.result` holds the fitted model. -- **Precedence:** a call's `mlops=` replaces the enclosing block, and an inner block replaces an outer one. A block covers the calling thread or asyncio task. -- **Off by default.** Without a block or `mlops=`, every call behaves as it always has. Plain exports (`export(config)`, `export_config=`, `output_config=`) never write to the artifact root. A path export that finds a complete copy of the same artifact at its destination reuses it instead of uploading again (logged at info): a snapshot with the same bytes, a training table or prediction directory with `_SUCCESS` and the same `_KUMO_EXPORT` identity, or a model bundle with the same files and `_provenance.json`. - -## Lineage - -Every artifact carries a `_provenance.json` manifest, metadata only. The leading underscore keeps Spark and Hive Parquet readers from treating it as data. - -A graph snapshot's manifest, next to its index, records each source table's connector, location and version, plus the engine build that wrote it. `graph.snapshot(record_sources=False)` keeps only connector and table names, plus a SHA-256 of each full location so a source moved elsewhere is still noticed, and keeps Custom SQL only as a hash. - -### Training tables - -A generated or caller-supplied training table carries `_provenance.json` in its staging directory: -- the query, the effective generation plan and the task type; -- the split seed; -- the snapshot it was generated from (that snapshot's manifest, embedded); -- for `TrainingTable.from_source_table`, the source table's connector, location and version. +### Databricks -`TrainingTable.id` is deterministic: the same snapshot, query, plan, split seed and engine build give the same id. A caller-supplied source without a version gets a random id. An export to files (`TrainingTable.export`) writes the manifest next to the Parquet parts. `TrainingTable.load_from_connector(connector, path=..., pq=..., snapshot=...)` materializes an export again. +In a notebook or job on Databricks Runtime 17.3 LTS ML or later: -With MLOps on, each table is stored once under `/training_tables//`. - -### Models +```python +import kumo.engine as kumo +from kumo.engine.integrations.mlflow import MlflowTracking -Every fitted bundle carries `_provenance.json` (`ModelManifest`): -- its `id`, the bundle id (a content hash that leaves the manifest file out); -- when it was fitted and with which fit mode; -- its parents: the training table's manifest, embedded with that table's snapshot, and the warm-start model's manifest, if any. +cfg = kumo.MLOps("/Volumes/main/ml/kumo", tracking=MlflowTracking("/Users//churn"), + registry=kumo.ModelRegistryLocation("main.ml.users_churn", alias="candidate")) +with kumo.mlops(cfg): + snapshot = graph.snapshot() + result = kumo.Trainer(pquery.suggest_model_plan(snapshot)).fit(pquery.generate_training_table(snapshot), + snapshot=snapshot, scoring_defaults=kumo.ScoringDefaults(binary_classification_threshold=0.5)) + print(result.registered_model.version, result.tracking.url) +``` -An embedded model keeps only the ids of its own model parents, so a long chain of warm-start retrains does not grow each manifest. `fit` returns a new `TrainingResult` per call (`.id`, `.manifest`); `Trainer.result` is the latest fit's result or the loaded bundle's (`trainer.result.id`, `trainer.result.manifest`). +### SageMaker -With MLOps on, each bundle is stored under `/models//`. A complete copy that is already there is kept; a partial one is replaced. +With the managed MLflow tracking server's ARN (setup: [SageMaker managed MLflow](#sagemaker-managed-mlflow)): -The bundle also carries the published model's holdout predictions (`holdout.parquet`, with its truncation flag), so `result.holdout_df()` works on a model loaded with `Trainer.load` or `Trainer.load_from_connector` (a path or a registry), not only on the result of `fit`. Bundles fitted before this have no copy, and `holdout_df()` raises `FileNotFoundError` for them. +```python +import kumo.engine as kumo +from kumo.engine.integrations.mlflow import MlflowTracking, SageMakerMLflow -### Artifacts, bundles and dependencies +cfg = kumo.MLOps("s3:///kumo", tracking=MlflowTracking("users-churn", target=SageMakerMLflow("")), + registry=kumo.ModelRegistryLocation("users-churn", alias="candidate")) +with kumo.mlops(cfg): + snapshot = graph.snapshot() + result = kumo.Trainer(pquery.suggest_model_plan(snapshot)).fit(pquery.generate_training_table(snapshot), + snapshot=snapshot, scoring_defaults=kumo.ScoringDefaults(binary_classification_threshold=0.5)) + print(result.registered_model.version, result.tracking.url) +``` -- An **artifact** is anything with its own id and `_provenance.json`: a snapshot, a training table, a model, a prediction run, or an asset. -- A model's **bundle** holds only the model's own state: weights, encoders, the training spec and the holdout. -- Its manifest's **`parents`** are lineage: the training table and the warm-start model it came from. -- Its **`dependencies`** are what `predict` needs besides the bundle and a snapshot. The snapshot is an input to each call, never a dependency. Each entry has a `kind`: - - `predictor`: the entry point that scores the bundle; - - `package`: the `kumo-engine` version and its extras, with `[text]` only when the model encodes a text column; - - `asset`: a content-addressed file set, by `asset_kind`, `name`, `version` and `sha256`. -- Bundles fitted before dependencies were recorded list none. They load as before, with the built-in predictor. -- **Assets** are not copied into bundles. With MLOps on, each asset is stored once under `/assets///` with its own `_provenance.json`, and every model that uses it shares that copy. The dependency records where it was stored. -- **Resolving assets:** `Trainer.load`, the MLflow pyfunc and scoring jobs look for each asset at the location the dependency records, then under the artifact root of the enclosing `kumo.mlops` block, then in the embedding-model store. They check its sha256 and cache non-local copies under `~/.cache/kumo/assets//`. If an asset is stored but cannot be found, loading fails with `MissingDependencyError`, which names the asset kind, name and sha256. A bundle without stored assets resolves its models from the store, as before. -- **Text encoders** are the first assets (`asset_kind="text_encoder"`): the word vectors of each text column's embedding model. Word-vector models (`GloVeVectors`) are stored. Other `TokenEmbeddingModel`s are recorded by fingerprint and must be registered wherever the model is loaded. +### Snowflake -To train with your own word vectors and serve the model without registering them again, register them under a name, fit with MLOps on, and load the stored model anywhere that can reach the artifact root: +Tracking and the Snowflake Model Registry through one connector (`kumo-engine[snowflake,s3]`): ```python -store = kumo.EmbeddingModelStore() -store.add("colors", kumo.GloVeVectors("/data/vectors")) # a .zip release or a directory of converted files -plan.column_processing.encoder_overrides = {"users.bio": TextEmbedding("colors")} -with kumo.mlops(kumo.MLOps(artifact_root="s3://bucket/kumo")): - result = kumo.Trainer(plan, embedding_models=store).fit(table, snapshot=snapshot) -# Elsewhere, with no store: the vectors come from /assets/text_encoder//. -trainer = kumo.Trainer.load_from_connector(kumo.S3Connector(root_dir="s3://bucket/kumo"), - path=f"models/{result.id}") +import kumo.engine as kumo +from kumo.engine.integrations.snowflake import SnowflakeTracking + +sf = kumo.SnowflakeConnector(warehouse="WH", database="ML", schema_name="MODELS") +cfg = kumo.MLOps("s3:///kumo", tracking=SnowflakeTracking(sf, experiment_name="CHURN"), + registry=kumo.ModelRegistryLocation("CHURN_GNN", alias="candidate")) +with kumo.mlops(cfg): + snapshot = graph.snapshot() + result = kumo.Trainer(pquery.suggest_model_plan(snapshot)).fit(pquery.generate_training_table(snapshot), + snapshot=snapshot, scoring_defaults=kumo.ScoringDefaults(binary_classification_threshold=0.5)) + print(result.registered_model.version, result.tracking.url) ``` -### Scoring defaults +To serve only a better model, move its alias yourself: [Promote a model](#promote-a-model). -A model can carry the prediction arguments `predict` uses when a call omits them: +## How it fits together -```python -from kumo.engine import ModelOutputConfig, ScoringDefaults +- The **artifact root** holds every artifact a block produces, each in its own directory with a `_provenance.json` manifest that embeds its parents' manifests. +- The **tracker** records each fit as a run of its own; the **registry** is the tracker's: MLflow (Unity Catalog, SageMaker MLflow, any MLflow server) or the Snowflake Model Registry. -tuned = result.with_scoring_defaults( - ScoringDefaults(binary_classification_threshold=0.37)) -tuned.predict(snapshot) # thresholds at 0.37 -tuned.scoring_defaults # ScoringDefaults(binary_classification_threshold=0.37, ...) -``` +| id | names | where to find it | +|---|---|---| +| `snapshot.id` | a graph snapshot | `/snapshots///`; `result.snapshot_id` | +| `table.id` | a training table | `/training_tables//` | +| `result.id` | a model (its bundle id) | `/models//`; the `kumo.bundle_id` run tag; `kumo_bundle_id` on the registered version | +| `result.registered_model.version` | the registered version | the registry | +| `result.tracking` | the tracker's own run of the fit: `platform`, `run_id`, `url` | the tracker's UI | -`with_scoring_defaults` returns a `TrainingResult` for a copy of the bundle with its own id. The copy's manifest embeds the original's in `parents.derived_from`. The threshold applies to binary and multilabel classification, `num_classes_to_return` to multiclass, multilabel and link prediction. Explicit `predict` arguments and a query's `TOP k` win; `ScoringDefaults()` removes the defaults. `Trainer.export(config)` publishes a fitted or loaded bundle like `fit(export_config=...)`. +Find runs with the platform's UI, or by tag: `kumo.kind` (`fit`, `trial`) and `kumo.bundle_id`. -### Registering models +## Turning MLOps on -With `registry=ModelOutputConfig(connector, location=...)`, every fit in the block registers its model in the connector's registry (Snowflake Model Registry, or Unity Catalog through a `DatabricksConnector`) after it is stored and tracked: +`kumo.MLOps` says where artifacts go and what records them. Use it as a block, or pass it as `mlops=` on a single call (`fit`, `graph.snapshot`, `generate_training_table`): ```python -from kumo.engine import ModelOutputConfig, ModelRegistryLocation - -cfg = kumo.MLOps("s3://bucket/kumo", registry=ModelOutputConfig( - registry, location=ModelRegistryLocation("CHURN_GNN", alias="challenger"))) +cfg = kumo.MLOps("s3://bucket/kumo", tracking=..., registry=..., on_error="raise") with kumo.mlops(cfg): - result = trainer.fit(training_table, snapshot=snapshot) -result.registered_model # PublishedModel(name='CHURN_GNN', version='V3', uri=...) - -# Scoring jobs load whichever version an alias points at. -trainer = kumo.Trainer.load_from_connector(registry, location=ModelRegistryLocation( - "CHURN_GNN", alias="challenger")) + snapshot = graph.snapshot() +trainer.fit(table, snapshot=snapshot, mlops=cfg) ``` -- Registration is independent of `export`. If it fails, `fit` warns with the error and returns the result with `registered_model=None` (with `on_error="raise"`, it raises instead), and a tracked fit run ends failed with `kumo.error`; an invalid alias or a connector that cannot publish models raises before training, and an invalid alias or model name raises when `kumo.MLOps` is built. Snowflake model names are Snowflake identifiers (a letter or underscore, then letters, digits, underscores or `$`, e.g. `USERS_CHURN`; or a double-quoted name); Unity Catalog model names are at most 255 characters with no space, `/`, `.` or control character, optionally qualified with the connector's own catalog and schema. -- `result.export(ModelOutputConfig(connector, location=...))` on a `TrainingResult`, or `trainer.export(...)` for its latest result, publishes the model with or without a block and also accepts an `alias`; `Trainer.load_from_connector` loads it back. `fit(export_config=...)` is shorthand for fitting then exporting, and its version is `result.published_model`. -- An alias is a letter or underscore followed by letters, digits or underscores; `DEFAULT`, `FIRST`, `LAST`, `LATEST` and `v` are reserved. The alias is set on the new version after publishing, moving it from any version that held it; if that fails, the publish still stands, with a warning. Snowflake allows one alias per version; if moving an alias fails, it is put back on the version that held it. Snowflake's `DEFAULT`, `FIRST` and `LAST` can be loaded but not set. File, S3 and GCS connectors have no aliases. -- Each registered version records only pointers: the `kumo_engine` marker and `kumo_bundle_id`, as version metrics in the Snowflake Model Registry and as version tags in Unity Catalog through a `DatabricksConnector`. (A version registered through MLflow, below, gets a `kumo_bundle_id` version tag only.) The bundle's own `_provenance.json` holds its lineage. `list_models` and `fetch_model` return `kumo_bundle_id` in `PublishedModel.metadata`; `load_from_connector` warns if the fetched bundle differs from it. Through a `DatabricksConnector`, publishing a bundle that a ready version of the model already holds (same `kumo_bundle_id`) returns that version, logged at info, instead of adding another; an alias then moves to it. +- **`artifact_root`** (required): a URI (`s3://`, `gs://`, `file://`, or a local path, including `/Volumes/...` on a Databricks cluster) or a `(connector, path)` pair for a store a URI cannot reach: a volume through a `DatabricksConnector` (`"/Volumes////kumo"`) or a named stage through a `SnowflakeConnector` (`"@DB.SCHEMA.STAGE/kumo"`) once the installed sink holds files there. +- **`tracking`**: `MlflowTracking` or `SnowflakeTracking` (below), or your own [`Tracking`](#custom-trackers). +- **`registry`**: a `ModelRegistryLocation(name, alias=...)`; every fit registers its model there, in the tracker's registry. +- **`store`**: the artifact kinds stored when produced: `snapshot`, `training_table` and `model` (default: all). +- **`on_error`**: `"warn"` (default) warns about a failed store, tracking call or registration and returns; `"raise"`, for scheduled jobs, raises once the call has finished, with the failure as the `__cause__`. After `fit`, `trainer.result` still holds the model. A model whose bundle or assets could not be stored under the root is never registered. +- **Precedence:** a call's `mlops=` replaces the enclosing block, and an inner block replaces an outer one. A block covers the calling thread or asyncio task. +- **Off by default.** Without a block or `mlops=`, every call behaves as it always has; plain exports (`export(config)`, `export_config=`) never write to the artifact root. +- `kumo.MLOps(...)` checks the tracker and registry when it is built (one experiment lookup, one registry listing, the model name's rules), so wrong credentials or a bad name fail before any work. -### Promote a Kumo model +### Reading artifacts back -To register every fit as a candidate and promote only a better one, register under a candidate alias (or none) and move the serving alias yourself with the registry's own API. Compare the new model's test metrics (`result.metrics`, also on its tracked fit run) with the current champion's (`load_from_connector(...).result.metrics`): +Pass the `kumo.MLOps` where a connector goes, with paths relative to the artifact root: ```python -champion = kumo.Trainer.load_from_connector( - registry, location=ModelRegistryLocation("CHURN_GNN", alias="champion")) -better = result.metrics["test_auroc"] > champion.result.metrics["test_auroc"] -version = result.registered_model.version +snapshot = kumo.GraphSnapshot.load_from_connector(cfg, id=result.snapshot_id) +trainer = kumo.Trainer.load_from_connector(cfg, path=f"models/{result.id}") +champion = kumo.Trainer.load_from_connector(cfg, location=kumo.ModelRegistryLocation("users-churn", alias="champion")) +table = kumo.TrainingTable.load_from_connector(cfg, path=f"training_tables/{table_id}", pq=pquery, snapshot=snapshot) ``` -Unity Catalog (registered through MLflow or a `DatabricksConnector`) and SageMaker managed MLflow: +`location=` loads the version an alias (or `version=`) names from the tracker's registry, one way on every platform. A loaded model is a `Trainer`; its `result` is the `TrainingResult` (`.id`, `.manifest`, `.metrics`, `.scoring_defaults`, `.holdout_df()`). -```python -from mlflow import MlflowClient +## Train and register -if better: - MlflowClient().set_registered_model_alias("main.ml.churn_gnn", "champion", version) -``` +`fit` inside the block stores the bundle under `models//`, records the fit as one run (a child run per hyperparameter-search trial attempt on MLflow, `trial--attempt-/` names inside the run on Snowflake), and registers the model as a new version of the registry location, moving its alias to it. -`set_registered_model_alias` moves the alias from the version that held it. For models registered through MLflow, the champion's metrics are also on its run: `client.get_run(client.get_model_version_by_alias(name, "champion").run_id).data.metrics`. +- `result.registered_model` is the version `MLOps(registry=)` registered; `result.published_model` the one `fit(export_config=...)` published. +- `result.tracking` is the tracker's run: `TrackingRef(platform="databricks", run_id="...", url="https://.../ml/experiments/.../runs/...")`. The model's `_provenance.json` records the same. +- `fit(scoring_defaults=kumo.ScoringDefaults(...))` stores prediction defaults in the bundle before it is stored and registered: `binary_classification_threshold` (binary and multilabel classification) and `num_classes_to_return` (multiclass, multilabel, link prediction). `predict` and the MLflow pyfunc apply them when a call omits them; explicit arguments and a query's `TOP k` win. `result.with_scoring_defaults(...)` retunes a fitted model as a new bundle whose manifest names the original in `parents.derived_from`. +- `result.export(cfg)` (or `trainer.export(cfg)` for a loaded model) registers a model at `cfg.registry` through its tracker. On every platform, registering a bundle the registry already holds under that name returns its version instead of adding one. +- If registration fails, `fit` warns with the full error and returns with `registered_model=None` (with `on_error="raise"`, it raises), and the fit run ends failed with `kumo.error`. Registering a bundle the registry already holds under the name returns that version (logged at info) and moves the alias to it. -Snowflake Model Registry, where a version holds at most one alias, so the alias is unset first: +**Model names** follow the registry's rules, checked when `kumo.MLOps` is built: -```sql -ALTER MODEL DB.SCHEMA.CHURN_GNN VERSION champion UNSET ALIAS; -ALTER MODEL DB.SCHEMA.CHURN_GNN VERSION V3 SET ALIAS = champion; -``` +| registry | rule | example | +|---|---|---| +| Unity Catalog | `catalog.schema.model`, each part ≤ 255 characters without space, `/` or control characters | `main.ml.users_churn` | +| SageMaker MLflow | letters, digits and hyphens, ≤ 57, starting and ending with a letter or digit | `users-churn` | +| Snowflake Model Registry | a Snowflake identifier | `USERS_CHURN` | -Scoring jobs that load `alias="champion"` (or `models:/@champion`) pick up the new version on their next run. +An alias is a letter or underscore followed by letters, digits or underscores; `DEFAULT`, `FIRST`, `LAST`, `LATEST` and `v` are reserved. -### Moving snapshots between machines +### Promote a model -To build an index on one machine and train or score on another, export it to a file destination and load it there: +Register every fit under a candidate alias and move the serving alias yourself once the new model is better, comparing `result.metrics` with the champion's: ```python -from kumo.engine import GraphSnapshot, S3Connector, SnapshotOutputConfig +champion = kumo.Trainer.load_from_connector(cfg, location=kumo.ModelRegistryLocation(name, alias="champion")) +if result.metrics["test_auroc"] > champion.result.metrics["test_auroc"]: + mlflow.MlflowClient().set_registered_model_alias(name, "champion", result.registered_model.version) +``` + +In the Snowflake Model Registry a version holds at most one alias, so unset it first: -location = graph.snapshot().export(SnapshotOutputConfig(S3Connector(), "s3://bucket/snapshots")) -# On the training machine: -snapshot = GraphSnapshot.load_from_connector(S3Connector(), path=location) +```sql +ALTER MODEL ML.MODELS.CHURN_GNN VERSION champion UNSET ALIAS; +ALTER MODEL ML.MODELS.CHURN_GNN VERSION V3 SET ALIAS = champion; ``` -- **Layout:** `export` writes `///`, the two halves of `snapshot.id`, and returns that location. It holds `kumo_snapshot.json` and one directory per export generation (`index/` and `_provenance.json`). -- **Transfer:** files move in parallel with a CRC32 per range. Each export uploads into a new generation and commits by writing `kumo_snapshot.json` last, naming that generation. Until then the previous export stays loadable, and concurrent exports never mix files: the last commit wins whole. Superseded generations, and abandoned ones over a day old, are removed best-effort. A load either completes or leaves nothing; if a newer export commits while it downloads, it switches to that one. -- **Reuse:** loading again reuses the local copy only if it holds this export's bytes (the same range checksums) and passes DiskGraph's integrity check; otherwise it is downloaded again. Exporting again skips the upload when the destination already holds a complete export (its `kumo_snapshot.json` and `_DONE`) with the local index's bytes, and logs that it reused it. -- **Staleness:** `snapshot.source_changes(graph)` reports, per table of `graph`, whether its source is `unchanged`, `changed` or `unknown` since the snapshot was built. `load_from_connector(..., graph=graph)` warns when the snapshot was taken of another graph, naming the schema changes (tables, columns, types, keys, edges), or when any of its tables changed. A graph id also covers each table's connector and location, so a graph that differs only there also warns, saying the schema is the same. -- **Manifest:** the snapshot's `_provenance.json` travels with the export. -- **Trust:** load only from paths whose writers you trust. Checksums detect corruption, not tampering. -- **Destinations:** any destination whose sink holds files at the path works: S3, GCS, local files, or a volume or stage once those sinks are installed. A warehouse connector raises. +Jobs that load `alias="champion"` pick up the new version on their next run. -With MLOps on (`graph.snapshot(mlops=...)` or a `kumo.mlops` block), each snapshot is exported under `/snapshots/`, unless a complete export of it is already there. +### Trackers -### Scoring without the graph +`MlflowTracking(experiment_name=None, run_name=None, log_system_metrics=True, tags=None, log_model=False, *, target=None)`: +- With no `target`, it records where MLflow's own configuration points (`MLFLOW_TRACKING_URI`, `mlflow.set_tracking_uri`; automatic on Databricks). `target` names another server with its credentials: a `DatabricksConnector`, a `SageMakerMLflow(arn)`, or a tracking URI. With a `target`, `experiment_name` is required to record runs. +- `log_model=True` logs the model on the fit run without registering it; `MLOps(registry=)` implies it. -A scoring job does not need the training job's `Graph` code. It loads the snapshot the training job exported and the model, then scores: +`SnowflakeTracking(source, *, experiment_name, database=None, schema_name=None, run_name=None, tags=None, native_model_lineage=False)`: +- `source` is a `SnowflakeConnector`, whose identity opens the session and whose Snowflake Model Registry `MLOps(registry=)` registers in, or a Snowpark session, which tracks only. +- `database` and `schema_name` default to the session's (`PUBLIC` when only `database` is given). The fit is in Snowsight under AI & ML » Experiments. +- The registered version carries the bundle; load it with `Trainer.load_from_connector`. -```python -snapshot = GraphSnapshot.load_from_connector(store, path=location) -trainer = Trainer.load_from_connector(store, path="models/") -result = trainer.predict(snapshot, binary_classification_threshold=0.5) -``` +## Score -- `predict(snapshot)` reuses the trained query and table-generation plan, and scores its population at `prediction_time`. -- `entities=` scores given rows instead. -- It warns when the snapshot was taken of a different graph than the one the model was trained on, naming the schema changes, for example `Schema changes: events edges.` A graph that differs only in table connectors or locations (a dev graph and a prod graph of the same schema) warns that the schema is the same. The model records its training graph from this release on; for an older model the warning cannot say what differs. -- `predict(prediction_table, snapshot=...)` scores a prediction table generated beforehand, as before. -- `fit` results score too: `result.predict(snapshot, ...)`. +`predict(snapshot)` scores the trained query's population at its prediction time; `entities=` scores given rows; `predict(prediction_table, snapshot=...)` scores a table generated beforehand. A scoring job needs no `Graph` code: `GraphSnapshot.load_from_connector(cfg, id=...)` and `Trainer.load_from_connector(cfg, ...)`, then `predict`. It warns when the snapshot was taken of a different graph than the model's, naming the schema changes. -## Progress events +### MLflow pyfunc -`fit(callbacks=[...])` calls each callback with every progress -`kumo.engine.tracking.Event`: `fit_started`, then per search trial -`trial_started`, `epoch_ended` and `trial_ended`, then `fit_ended`. Ended events -carry `params["status"]` and, on failure, a short `params["error"]`. A callback -that raises is logged and not called again; it never fails the fit. For JSON -lines on stdout: +A model registered through `MlflowTracking` is an MLflow pyfunc: it serves with `mlflow.pyfunc.load_model(uri)`, `spark_udf` and Model Serving. ```python -import json - -trainer.fit(training_table, snapshot=snapshot, - callbacks=[lambda e: print(json.dumps(e.to_dict()))]) +model = mlflow.pyfunc.load_model("models:/main.ml.users_churn@champion") +scores = model.predict(pd.DataFrame({"ENTITY": [1, 2, 3]}), params={ + "snapshot": "/Volumes/main/ml/kumo/snapshots//", + "prediction_time": "2024-06-01"}) # temporal queries ``` -## Experiment tracking +- Input rows have an `ENTITY` column and optionally a `TIMESTAMP`; scores come back in input order, one per scored row. Binary classification needs `params={"binary_classification_threshold": ...}` unless the model carries scoring defaults. +- Each replica loads a snapshot on its first `predict` with that URI and keeps it. `log_model(trainer, snapshot=..., serving_snapshot=...)` logs a model whose replicas load that snapshot at startup and use it when no `snapshot` param is given; `serving_snapshot` takes a URI, or a `GraphSnapshot` or snapshot id stored under the block's artifact root. +- Serving needs an image with kumo-engine installed: `mlflow models serve -m --env-manager local`. Its pip requirement is `kumo-engine[s3,gcs]` of the same version (plus `[text]` for text columns). -`kumo.MLOps(tracking=...)` records each fit in an experiment tracker: one run -per fit, with each hyperparameter-search trial attempt as a child run (MLflow) -or inside the fit run (Snowflake). Tracking is on only -inside a `kumo.mlops(cfg)` block or for a `fit` given `mlops=cfg`. +## Credentials and running off-platform + +`MlflowTracking(target=...)` and `SnowflakeTracking(connector)` take the same identities as the connectors, so a fit on a laptop or any VM can record to Databricks, SageMaker or Snowflake. An identity with no credentials (or no identity) uses the platform's ambient chain (a notebook, job or instance role, `aws configure`, `databricks auth login`); one with credentials uses exactly those. ```python -from kumo.engine.integrations.mlflow import MlflowTracking # kumo-engine[mlflow] +from kumo.engine import DatabricksConnector, DatabricksIdentity, S3Identity, SnowflakeConnector, SnowflakeIdentity -cfg = kumo.MLOps(artifact_root="s3://bucket/kumo", - tracking=MlflowTracking(experiment_name="users-spend")) -with kumo.mlops(cfg): - result = trainer.fit(training_table, snapshot=snapshot) +workspace = DatabricksConnector(catalog="main", schema_name="ml", warehouse_id="", + identity=DatabricksIdentity(host="https://", client_id="", client_secret="")) +databricks = MlflowTracking("/Users//users_churn", target=workspace) +sagemaker = MlflowTracking("users-churn", target=SageMakerMLflow("", identity=S3Identity(role_arn=""))) +snowflake = SnowflakeTracking(SnowflakeConnector(warehouse="", database="ML", schema_name="RUNS", + identity=SnowflakeIdentity(account="", username="", private_key="")), + experiment_name="USERS_CHURN") ``` -- A tracking or registration failure does not fail the fit: it is logged, and reported once as a `UserWarning` naming the full error, so both notebooks and job logs show it. With `on_error="raise"` the call raises once instead, after storing its result, with the error as the cause. Only global rank 0 tracks. -- Tracker calls run on a background thread, so a slow or unreachable tracking server does not hold up training. At the end of the fit, `fit` waits at most 30 s for tracking to finish; the rest continues in the background. -- During the fit, MLflow HTTP calls use 2 retries and a 30 s timeout unless `MLFLOW_HTTP_REQUEST_MAX_RETRIES` / `MLFLOW_HTTP_REQUEST_TIMEOUT` are set. -- `MlflowTracking()` records where MLflow's own configuration points when it is built (`MLFLOW_TRACKING_URI`, or `mlflow.set_tracking_uri`): on Databricks this is automatic, and with nothing configured MLflow writes a local store in the working directory. `MlflowTracking(target=...)` names the server and the credentials instead; see [Running from a laptop](#running-from-a-laptop). -- `kumo.MLOps(...)` checks the tracker when it is built: one experiment lookup, plus one registry listing and (Unity Catalog) the model's catalog and schema when it registers. An unreachable server, refused credentials or a missing catalog or schema raise there, before any work. -- The `View run ... at` links MLflow prints do not open on SageMaker; after a tracked fit kumo-engine warns how to open the UI: from SageMaker Studio, or at the URL `aws sagemaker create-presigned-mlflow-tracking-server-url --tracking-server-name ` prints. -- MLflow rejects `@` in metric names, so link-prediction metrics such as `test_map@10` are logged as `test_map_at_10`. -- Per-epoch metrics use step = epoch. After the fit, `tracking.run_id` is the fit run's MLflow run id. -- The fit run is named `fit-` (for example `fit-20261027T061838Z`) unless `run_name` is given; trial runs are `trial--attempt-`. -- The fit's source tables are logged as MLflow run inputs, read from the training table's embedded snapshot manifest (plus a custom training table's source). Databricks tables are logged as Delta tables at the snapshot's Delta version, so Unity Catalog links each table to a model registered from the run; file and object-store sources by URI. Other warehouse tables and Custom SQL sources are not logged as inputs; they stay in the model's `_provenance.json`. +`load_model(uri, target=...)` and `log_model(..., target=...)` take the same targets. A `DatabricksConnector` without an identity uses a notebook, a job or a `~/.databrickscfg` profile with no environment variables to set. MLflow has no per-client credentials, so with explicit credentials each MLflow call kumo-engine makes sets `DATABRICKS_HOST` and a short-lived `DATABRICKS_TOKEN` (never the secret), or the SageMaker tracking URI and AWS session, and restores them afterwards: don't call MLflow or Databricks yourself on other threads during a fit. `examples/mlops/smoke.py databricks|sagemaker|snowflake` is a manual live check per platform. -Snowflake records the fit in Snowflake Experiment Tracking (Snowsight: AI & ML -» Experiments), using a Snowpark session: +### SageMaker managed MLflow -```python -from kumo.engine.integrations.snowflake import SnowflakeTracking # kumo-engine[snowflake] +1. Create a tracking server once (15-25 minutes): `aws sagemaker create-mlflow-tracking-server --tracking-server-name --artifact-store-uri s3:///mlflow --role-arn `; the server role needs read and write on that prefix. +2. Give the job's identity `sagemaker-mlflow:*` on the server ARN and S3 read and write on the artifact store and the kumo artifact root. Rights scoped to the root's prefix (`s3:GetObject`, `s3:PutObject` and `s3:DeleteObject` on `//*`, `s3:ListBucket` limited to that prefix) suffice: kumo-engine makes no bucket-level calls. +3. Install `kumo-engine[mlflow,sagemaker]` and pass `MlflowTracking(..., target=SageMakerMLflow(arn))`; the registry is the same server. +4. MLflow's run links do not open on SageMaker: open the UI from SageMaker Studio, or at the URL `aws sagemaker create-presigned-mlflow-tracking-server-url --tracking-server-name ` prints. A tracked fit warns with this command, and `result.tracking.url` is `None`. -cfg = kumo.MLOps(artifact_root="s3://bucket/kumo", - tracking=SnowflakeTracking(session, experiment_name="users_spend")) -with kumo.mlops(cfg): - result = trainer.fit(training_table, snapshot=snapshot) -``` +## Reference -- Snowflake has no child runs or tags, so trial attempts are logged in the fit run as `trial-1-attempt-1/` params and metrics, and tags are params of the same name. -- `native_model_lineage=True` links a model the fit publishes through Snowflake lineage. -- Names beyond Snowflake's 200 metrics and 1000 params per run are counted in `kumo.dropped_metrics` / `kumo.dropped_params`. -- If a Snowflake run is already active in the process, the fit is not tracked there; an existing run is never resumed (a taken `run_name` gets a suffix). -- Snowflake Experiment Tracking has no dataset inputs, so source tables are not recorded there. -- The Snowpark session's parameters are left as they are. After the fit, `tracking.run_id` is the fit run's name. +### Install -Runs carry pointers, not lineage: +- A bare `kumo-engine[mlflow]` with `--find-links ` can resolve to an unrelated `kumo-engine` on PyPI; install the engine wheel by path: `pip install --find-links "/[mlflow]"`. +- **GPU:** on a plain GPU VM or SageMaker image, install a torch build that matches the NVIDIA driver first, e.g. `pip install --index-url https://download.pytorch.org/whl/cu126 "torch==2.12.*" torchvision`. PyPI's default torch (cu130) needs driver ≥ 580; cu126 runs on ≥ 525.60. With `accelerator="auto"`, a torch that cannot use the driver runs on CPU: check `torch.cuda.is_available()`. +- **Databricks:** Runtime 17.3 LTS ML or later keeps its torch, torchvision, MLflow and pyarrow (kumo-engine accepts pyarrow 16 through 24; a custom CA bundle for S3 needs 21 or later). On 16.4 LTS ML pip replaces torch and breaks torchvision. The parent folder of an experiment path must exist. Registering to Unity Catalog from outside Databricks also needs `boto3`. +- **MLflow:** the `mlflow` extra installs `mlflow-skinny` 3.0.1 or later; a 3.x client works against SageMaker's MLflow 3.0 servers. +- **SageMaker serving:** MLflow's stock serving containers cannot install kumo-engine; on SageMaker endpoints serve with `-m /opt/ml/model -p 8080`. GPU endpoints need a driver for CUDA 12.6 (set `InferenceAmiVersion`, e.g. `al2023-ami-sagemaker-inference-gpu-4-1`); batch transform takes JSON (`dataframe_split` plus `params`). -| tag | on | value | -|---|---|---| -| `kumo.kind` | every run | `fit` or `trial` | -| `kumo.run_id` | every run | the kumo 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 (including a failed model logging or registration) | outcome | +### Exporting and loading artifacts -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`. +Every artifact has one way out and one way back. A `kumo.MLOps` stands in for the connector to read the artifact root. -Which id is which: - -| id | what it names | where to read it | +| Artifact | Export | Load | |---|---|---| -| kumo run id | one `fit` call | the `kumo.run_id` tag; `tracking.run_id` in the model's `_provenance.json` | -| bundle id | the fitted model's content | `result.id`; the `kumo.bundle_id` tag; `kumo_bundle_id` in the registry; `models//` under the artifact root | -| tracker run id | the run in the tracker | the tracker's `run_id` attribute (`tracking.run_id` on the `MlflowTracking` object: the MLflow run id; Snowflake: the run name) | -| registry version | the registered model version | `result.registered_model.version` | +| Snapshot | `snapshot.export(SnapshotOutputConfig(connector, path))`, returns the location | `GraphSnapshot.load_from_connector(connector, path=location)`, `(cfg, id=...)` | +| Training table | `table.export(TrainingTableOutputConfig(connector, name))` | `TrainingTable.load_from_connector(connector, path=name, pq=pq, snapshot=snapshot)` | +| Model | `result.export(ModelOutputConfig(connector, path=... or location=...))`, or `fit(export_config=...)` | `Trainer.load_from_connector(connector, path=... or location=...)` | -To find a bundle's MLflow fit run, search for its tag: ``mlflow.search_runs(filter_string="tags.`kumo.bundle_id` = '' AND tags.`kumo.kind` = 'fit'", search_all_experiments=True)``. +Loading needs a file destination; a table export is for consumers outside kumo-engine. Exporting again to a destination that holds a complete copy of the same artifact reuses it (logged at info). `ModelOutputConfig(connector, location=...)` publishes to a connector's registry outside a block. -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 -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`, `check(registry)` (called when `kumo.MLOps` is built) and `log_model(result, snapshot, registry)` (see MLflow models below). 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. +### Lineage and manifests -### MLflow models +Every artifact carries `_provenance.json`, metadata only (the underscore keeps Spark and Hive readers off it). Parents are embedded, minus their statistics and mutable fields; a parent without a manifest (a bundle from an older release) is an id stub. -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 bundle the registry already holds under `name` (same `kumo_bundle_id` version tag) is not registered again: its version is returned, the alias moves to it, and an INFO log says it was reused. A bare `ModelRegistryLocation` without such a tracker, a fixed `version`, an invalid alias, or a name the registry (`mlflow.get_registry_uri()`) rejects raises when `kumo.MLOps` is built: +- **Snapshot:** each source table's connector, location and version, and the engine build. `graph.snapshot(record_sources=False)` keeps only connector and table names, a SHA-256 of each location, and Custom SQL as a hash. +- **Training table:** the query, generation plan, task type and split seed, its snapshot, and for `TrainingTable.from_source_table` the source table. `table.id` is deterministic over snapshot, query, plan, seed and engine build; a source without a version gets a random id. +- **Model** (`ModelManifest`, `result.manifest`): the bundle id (a content hash excluding the manifest), fit times and fit mode, the training table and warm-start model (an embedded model keeps only ids of its own model parents), `tracking` (the tracker's run) and `dependencies`. The bundle carries the holdout predictions (`holdout.parquet`), so `holdout_df()` works on a loaded model. -- SageMaker managed MLflow (a tracking server or MLflow app ARN): letters, digits and hyphens, starting and ending with a letter or digit, at most 57 letters and digits (`^[a-zA-Z0-9](-*[a-zA-Z0-9]){0,56}$`), so no `_` or `.`: `users-churn`. -- Unity Catalog (`databricks-uc`): `catalog.schema.model`, each part at most 255 characters with no space, `/` or control character: `main.ml.users_spend`. On a cluster a shorter name takes the current catalog and schema. -- Other MLflow registries: no extra rule. +### Artifacts, bundles and dependencies -The rules conflict, so no one name fits every registry: SageMaker forbids the `.` Unity Catalog needs between its parts and the `_` Snowflake needs instead of `-`. The examples use `users-churn` for SageMaker and other MLflow registries, `..users_churn` for Unity Catalog and `USERS_CHURN` for Snowflake. +- A model's **bundle** holds only its own state: weights, encoders, the training spec and the holdout. Its manifest's **`parents`** are lineage; its **`dependencies`** are what `predict` needs besides the bundle and a snapshot: the `predictor` entry point, the `kumo-engine` `package` (with `[text]` when a text column is encoded) and content-addressed `asset`s. +- **Assets** are stored once under `/assets///` and shared by every model that uses them. `Trainer.load` and the pyfunc look for each at the location recorded, then under the enclosing block's root, then in the embedding-model store, check its sha256 and cache non-local copies under `~/.cache/kumo/assets/`. A stored asset that cannot be found raises `MissingDependencyError`. +- **Text encoders** are the first assets: word-vector models (`GloVeVectors`) are stored; other `TokenEmbeddingModel`s are recorded by fingerprint and must be registered wherever the model loads: ```python -import mlflow -import pandas as pd -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")) +store = kumo.EmbeddingModelStore() +store.add("colors", kumo.GloVeVectors("/data/vectors")) +plan.column_processing.encoder_overrides = {"users.bio": TextEmbedding("colors")} 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(pd.DataFrame({"ENTITY": [1, 2, 3]}), params={ - "snapshot": "s3://bucket/kumo/snapshots//", - "prediction_time": "2024-06-01", # temporal queries - "binary_classification_threshold": 0.5, # binary classification only -}) + result = kumo.Trainer(plan, embedding_models=store).fit(table, snapshot=snapshot) +trainer = kumo.Trainer.load_from_connector(cfg, path=f"models/{result.id}") # anywhere, no store ``` -- `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 once (`Could not log the model: ...` or `Could not register the model: ...`, with the full error) and still returns its result with `registered_model=None` (with `on_error="raise"`, it raises instead), the run ends `FAILED` with the error in `kumo.error`, 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, the registered version's tags and `result.registered_model.metadata` hold `kumo_bundle_id`. The model's `provenance` artifact is the bundle's `_provenance.json`. Its pip requirement is `kumo-engine[s3,gcs]` of the same version, plus `[text]` when the model encodes a text column. The pyfunc scores with the predictor the bundle's manifest names. It resolves the model's assets, such as non-default word vectors, as described in [Artifacts, bundles and dependencies](#artifacts-bundles-and-dependencies). -- 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=` (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. -- Input rows have an `ENTITY` column and, optionally, a `TIMESTAMP` column. Scores come back in the input's row order, one per input row (repeated rows included); rows whose entity is not in the snapshot are left out. A temporal query scores entities at a time: pass `params={"prediction_time": ...}` or a `TIMESTAMP` column; without either, `predict` fails with `temporal query: pass params={'prediction_time': ...} or a TIMESTAMP column`. Binary classification needs `params={"binary_classification_threshold": ...}` unless the model's `scoring_defaults` set one. -- Serving needs an image with kumo-engine already installed; MLflow's stock serving containers cannot install it. Serve with `mlflow models serve -m --env-manager local` (on SageMaker endpoints, `-m /opt/ml/model -p 8080`). SageMaker batch transform takes JSON input (`dataframe_split` plus `params`), not CSV, which carries no params. -- GPU real-time endpoints on SageMaker need an inference image whose NVIDIA driver supports CUDA 12.6. The default GPU hosting image ships driver 470, so set a newer `InferenceAmiVersion` on the production variant, for example `al2023-ami-sagemaker-inference-gpu-4-1`. Not yet tested live. -- The `mlflow` extra installs `mlflow-skinny` (3.0.1 or later), the client without the MLflow server's dependencies. 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, use Databricks Runtime 17.3 LTS ML or later: its torch already satisfies kumo-core, so installing kumo-engine keeps the runtime's torch, torchvision and MLflow. kumo-engine accepts pyarrow 16 through 24, so a runtime's pyarrow in that range (19 on 17.3 LTS ML) is kept too; a custom CA bundle for S3 needs pyarrow 21 or later. On 16.4 LTS ML pip replaces torch 2.6 and breaks the preinstalled torchvision. -- 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. +### Moving snapshots between machines -#### Which registry +- **Layout:** `export` writes `///`, holding `kumo_snapshot.json` and one directory per export generation (`index/` and `_provenance.json`). +- **Transfer:** parallel, with a CRC32 per range. Each export commits by writing `kumo_snapshot.json` last, so the previous export stays loadable and concurrent exports never mix: the last commit wins. Superseded generations are removed best-effort. A load completes or leaves nothing, and follows a newer export that commits meanwhile. +- **Reuse:** a load reuses the local copy only if it holds the export's bytes and passes DiskGraph's integrity check; an export skips the upload when the destination holds a complete export with the same bytes. +- **Staleness:** `snapshot.source_changes(graph)` reports each table as `unchanged`, `changed` or `unknown`; `load_from_connector(..., graph=graph)` warns about another graph's schema or changed tables. +- **Trust:** load only from paths whose writers you trust; checksums detect corruption, not tampering. A warehouse connector cannot hold a snapshot. -Pick one per registered model, since the two store different formats. On Databricks both work; use the MLflow model unless only kumo-engine jobs load it, since it also serves through `mlflow.pyfunc` (Model Serving, `spark_udf`). 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`). +### Custom trackers -| | 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 | the registry's own API (see [Promote a Kumo model](#promote-a-kumo-model)) | MLflow's `set_registered_model_alias` | -| Databricks | `DatabricksConnector` (Unity Catalog) | Unity Catalog through MLflow | -| SageMaker | not available | SageMaker managed MLflow (`MlflowTracking(target=SageMakerMLflow(arn), ...)`) | -| Snowflake | `SnowflakeConnector` (Snowflake Model Registry) | not available | -| Plain VM | not available (store under the artifact root instead) | any MLflow tracking server | - -#### SageMaker managed MLflow - -1. Create a tracking server once and reuse it; creation takes about 15-25 minutes: - `aws sagemaker create-mlflow-tracking-server --tracking-server-name --artifact-store-uri s3:///mlflow --role-arn ` - (the server role needs read and write on that S3 prefix). -2. Give the identity that runs the job (notebook, VM role or pipeline execution role) `sagemaker-mlflow:*` on the tracking server ARN, and S3 read and write on the artifact store prefix and on the kumo artifact root. -3. Install `pip install "kumo-engine[mlflow,sagemaker]"` (the MLflow client, S3 and the `sagemaker-mlflow` plugin) and pass `MlflowTracking(target=SageMakerMLflow("arn:aws:sagemaker:::mlflow-tracking-server/"), experiment_name=...)`; the registry is the same server. Without an identity, the job's AWS identity (a role or `aws configure`) signs; `MLFLOW_TRACKING_URI` need not be set. The plugin signs and routes with the process-wide tracking URI, so each MLflow call kumo-engine makes sets it to the ARN and restores it afterwards: other threads that call MLflow during a fit see the SageMaker URI. -4. Model names use letters, digits and hyphens only (`users-churn`); `kumo.MLOps` checks this when it is built. -5. The run links MLflow prints do not open on SageMaker. Open the UI from SageMaker Studio or at the URL `aws sagemaker create-presigned-mlflow-tracking-server-url --tracking-server-name ` prints; a tracked fit warns with this command. - -On S3 the artifact root, source tables and prediction outputs are plain `s3://` locations read and written with the job's AWS identity (instance profile, `AWS_PROFILE` or the job role): +Other systems plug in through `kumo.engine.tracking.Tracking`, 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 attempt after each of its events, and `finish(record, error)` with the final record (`None` when tracking stopped early). Optional members: `prepare()`, `tags`, `check(registry)` (when `kumo.MLOps` is built), `log_model(bundle_dir, manifest, snapshot, registry)` (logs and registers, returning the version) or `connector` (a connector whose model registry `registry` registers in), `load_model(location)` and `register(result, location)` (for loading and `result.export(cfg)`), `platform`/`run_id`/`run_url` (recorded as `result.tracking`). A tracker that raises is not called again for that fit and its run ends failed. `RunRecord.inputs` maps each source table to its `connector`, `location`, `version` and `version_kind`; a tracker's `tags` are added except `kumo.*` and param names. -```python -from kumo.engine import PredictionOutputConfig, S3Connector - -data = S3Connector(root_dir="s3:///data") # data["orders.parquet"] reads s3:///data/orders.parquet; data["orders"] a directory of files -cfg = kumo.MLOps("s3:///kumo", - tracking=MlflowTracking(target=SageMakerMLflow(""), experiment_name="users-churn"), - registry=ModelRegistryLocation("users-churn", alias="candidate")) -output = PredictionOutputConfig(S3Connector(), output_paths={"predictions": "s3:///scores/"}) -store = S3Connector() # monitoring: load with absolute s3:// paths -``` - -### Running from a laptop +### Progress events -`MlflowTracking(target=...)` and `SnowflakeTracking(connector=...)` take the -same identities as the connectors, so a fit on a laptop or any VM can record -to Databricks, SageMaker or Snowflake. An identity with no credentials uses -the platform's ambient chain (a notebook, job or instance role, `aws -configure`, `databricks auth login`); one with credentials uses exactly -those. +`fit(callbacks=[...])` calls each callback with every `kumo.engine.tracking.Event`: `fit_started`, per trial `trial_started`, `epoch_ended` and `trial_ended`, then `fit_ended`. Ended events carry `params["status"]` and, on failure, `params["error"]`; a callback that raises is logged and not called again. For JSON lines: `callbacks=[lambda e: print(json.dumps(e.to_dict()))]`. -```python -import kumo.engine as kumo -from kumo.engine import ModelRegistryLocation -from kumo.engine.connector import DatabricksConnector, DatabricksIdentity, S3Identity, SnowflakeConnector, SnowflakeIdentity -from kumo.engine.integrations.mlflow import MlflowTracking, SageMakerMLflow -from kumo.engine.integrations.snowflake import SnowflakeTracking - -# Databricks MLflow and Unity Catalog, as a service principal (kumo-engine[mlflow,databricks]). -workspace = DatabricksConnector(catalog="main", schema_name="ml", warehouse_id="", - identity=DatabricksIdentity(host="https://", client_id="", client_secret="")) -cfg = kumo.MLOps("s3:///kumo", - tracking=MlflowTracking(target=workspace, experiment_name="/Users//users_churn"), - registry=ModelRegistryLocation("main.ml.users_churn", alias="candidate")) - -# SageMaker managed MLflow, with an access key or a role to assume (kumo-engine[mlflow,sagemaker]). -cfg = kumo.MLOps("s3:///kumo", - tracking=MlflowTracking(target=SageMakerMLflow("", identity=S3Identity(role_arn="")), - experiment_name="users-churn"), - registry=ModelRegistryLocation("users-churn", alias="candidate")) - -# Snowflake Experiment Tracking, as a key-pair user (kumo-engine[snowflake]). -snowflake = SnowflakeConnector(warehouse="", database="ML", schema_name="RUNS", - identity=SnowflakeIdentity(account="", username="", private_key="")) -cfg = kumo.MLOps("s3:///kumo", tracking=SnowflakeTracking(connector=snowflake, experiment_name="USERS_CHURN")) -``` +### Limits -- Each target is checked when `kumo.MLOps` is built, so wrong credentials or a missing catalog, schema or experiment folder parent fail before the fit starts. -- MLflow has no per-client credentials: its Databricks store reads `DATABRICKS_*` variables or a profile, and the SageMaker plugin signs with the process-wide tracking URI. So with explicit credentials, each MLflow call kumo-engine makes sets `DATABRICKS_HOST`, a short-lived `DATABRICKS_TOKEN` (never the client secret) and the process-wide tracking URI `databricks` (Unity Catalog registration downloads the logged model through it), or the process-wide tracking URI and the AWS session, and restores them when the call returns. With explicit credentials, don't call MLflow or Databricks yourself on other threads during a fit. Ambient identities change no credentials; a SageMaker target still sets the process-wide tracking URI for each call, so with a SageMaker target, other threads calling MLflow during a fit see the SageMaker URI. -- SageMaker model artifacts are uploaded and downloaded with the identity's credentials. `load_model(uri, tracking=MlflowTracking(target=..., experiment_name=...))` and `log_model(..., tracking=...)` take the same target. -- `examples/mlops/smoke.py databricks|sagemaker|snowflake` is a manual live check per platform: a small fit recorded and registered with an explicit identity, then loaded back, in a few minutes on a CPU. It is not run in CI. +- Tracker calls run on a background thread; `fit` waits at most 30 s for tracking to finish, and up to an hour for model logging and registration (600 s MLflow HTTP timeout during the upload). During a fit MLflow HTTP calls use 2 retries and a 30 s timeout unless `MLFLOW_HTTP_REQUEST_MAX_RETRIES` / `MLFLOW_HTTP_REQUEST_TIMEOUT` are set. Only global rank 0 tracks. +- MLflow rejects `@` in metric names: `test_map@10` is logged as `test_map_at_10`. Params MLflow rejects are counted in `kumo.dropped_params`. +- Snowflake keeps 200 metrics and 1000 params per run; the rest are counted in `kumo.dropped_metrics` / `kumo.dropped_params`. Snowflake has no tags (they are params of the same name), no child runs and no dataset inputs. A taken `run_name` gets a random suffix; an active run elsewhere in the process is never resumed. +- The fit's source tables are MLflow run inputs: Databricks tables as Delta tables at the snapshot's version (so Unity Catalog links them to the registered model), file sources by URI. diff --git a/packages/kumo-engine/kumo/engine/_deps.py b/packages/kumo-engine/kumo/engine/_deps.py index 189614856..f8d6b5713 100644 --- a/packages/kumo-engine/kumo/engine/_deps.py +++ b/packages/kumo-engine/kumo/engine/_deps.py @@ -10,7 +10,6 @@ import os import shutil import tempfile -import warnings from pathlib import Path from typing import ( Any, @@ -126,6 +125,7 @@ def __init__(self, records: Sequence[Dependency] = (), settings: Any = None, locations: Sequence[Location] = ()) -> None: self.records: List[Dependency] = list(records) + self.errors: List[BaseException] = [] self._settings = settings self._locations = list(locations) @@ -144,7 +144,10 @@ def require(self, asset: Asset) -> Dependency: location = None settings = self._settings if settings is not None and settings.wants('model'): - location = _store(asset, sha256, settings) + try: + location = _store(asset, sha256, settings) + except Exception as error: + self.errors.append(error) return self.add( Dependency(kind='asset', name=asset.name, version=asset.version, asset_kind=asset.asset_kind, sha256=sha256, @@ -214,8 +217,9 @@ def asset_cache() -> Path: def _store(asset: Asset, sha256: str, settings: Any) -> Optional[str]: r"""Writes ``asset`` once to ``/assets///``, - its ``_provenance.json`` last; returns where, or None (with a - warning) if it failed.""" + its ``_provenance.json`` last; returns where.""" + from kumo_core.object_store_sink import _prepared + from kumo.engine.connector import file_store relative = f'assets/{asset.asset_kind}/{sha256}' @@ -226,7 +230,7 @@ def _store(asset: Asset, sha256: str, settings: Any) -> Optional[str]: marker = f'{root}/{PROVENANCE_FILE}' if filesystem.get_file_info(marker).is_file: return settings.root.uri(relative) - filesystem.create_dir(root, recursive=True) + _prepared(filesystem, root) for name, path in asset.files().items(): with open(path, 'rb') as reader, \ filesystem.open_output_stream(f'{root}/{name}') as writer: @@ -235,10 +239,9 @@ def _store(asset: Asset, sha256: str, settings: Any) -> Optional[str]: writer.write(manifest_bytes(_manifest(asset, sha256))) return settings.root.uri(relative) except Exception as error: - warnings.warn( + raise RuntimeError( f"Could not store {asset.asset_kind} {asset.name!r} under the " - f"artifact root: {error}", UserWarning, stacklevel=4) - return None + f"artifact root: {error}") from error def _manifest(asset: Asset, sha256: str) -> AssetManifest: diff --git a/packages/kumo-engine/kumo/engine/_mlops.py b/packages/kumo-engine/kumo/engine/_mlops.py index bf3ef384a..0855838c4 100644 --- a/packages/kumo-engine/kumo/engine/_mlops.py +++ b/packages/kumo-engine/kumo/engine/_mlops.py @@ -237,13 +237,13 @@ def resolve(settings: Optional[MLOps] = None) -> Optional[MLOps]: def raise_on_error(settings: Optional[MLOps], error: Optional[BaseException], call: str, stacklevel: int = 2) -> None: - r"""Reports ``error``, a tracking or registration failure of a finished - ``call``: raises when ``settings.on_error`` is ``'raise'``, else warns - unless `warn_failure` already warned about it.""" + r"""Reports ``error``, a storing, tracking or registration failure of a + finished ``call``: raises when ``settings.on_error`` is ``'raise'``, + else warns unless `warn_failure` already warned about it.""" if error is None or settings is None: return - message = (f'{call} finished, but MLOps tracking or registration ' - f'failed: {full_error(error)}') + message = (f'{call} finished, but MLOps storing, tracking or ' + f'registration failed: {full_error(error)}') if settings.on_error == 'raise': raise RuntimeError(message) from error warn_failure(message, error, stacklevel + 1) diff --git a/packages/kumo-engine/kumo/engine/_snapshot_io.py b/packages/kumo-engine/kumo/engine/_snapshot_io.py index 5b8060244..cce308d7f 100644 --- a/packages/kumo-engine/kumo/engine/_snapshot_io.py +++ b/packages/kumo-engine/kumo/engine/_snapshot_io.py @@ -194,9 +194,7 @@ def export_location(snapshot: Any, path: str) -> str: def store(snapshot: Any, settings: Any) -> None: r"""Export ``snapshot`` under ``/snapshots/`` unless a complete - export of it is there; failures warn.""" - import warnings - + export of it is there.""" try: connector = settings.root.connector path = settings.root.location('snapshots') @@ -210,9 +208,8 @@ def store(snapshot: Any, settings: Any) -> None: and marker.get('snapshot_id') == snapshot.id): export(snapshot, connector, path) except Exception as error: - warnings.warn( - f"Could not store snapshot {snapshot.id} under the artifact " - f"root: {error}", UserWarning, stacklevel=3) + raise RuntimeError(f"Could not store snapshot {snapshot.id} under the " + f"artifact root: {error}") from error def export(snapshot: Any, connector: Any, path: str) -> str: diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index ea5c66bed..982ee731b 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -720,9 +720,13 @@ def snapshot(self, *, record_sources: bool = True, snapshot = self._snapshot_no_lock(graph_key, persist_dir, record_sources) if settings is not None and settings.wants('snapshot'): + from kumo.engine._mlops import raise_on_error from kumo.engine._snapshot_io import store - store(snapshot, settings) + try: + store(snapshot, settings) + except Exception as error: + raise_on_error(settings, error, 'Graph.snapshot') return snapshot def _snapshot_no_lock(self, graph_key: str, persist_dir: Path, diff --git a/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py b/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py index 8372eec35..62ecac2aa 100644 --- a/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py +++ b/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py @@ -9,8 +9,9 @@ the SageMaker plugin signs with the process-wide tracking URI and a boto3 session (``sagemaker_mlflow.auth_provider.AuthProvider.get_auth``). With explicit credentials, `Target.scope` sets exactly that state around each -call, under one lock, and restores it. Ambient credentials set nothing, -except that a SageMaker target still sets the process-wide tracking URI. +call, under one lock, and restores it. Ambient credentials set no +credentials; a Databricks or SageMaker target still sets the process-wide +tracking URI. """ import os @@ -165,7 +166,10 @@ def _databricks(connector: Any) -> Target: principal = _principal(connector.identity) if (principal is None or principal.config.authentication is DatabricksAuthenticationKind.AMBIENT): - return Target('databricks', 'databricks-uc') + # Unity Catalog registration downloads the logged model through the + # process-wide tracking URI, so it must point at the workspace. + return Target('databricks', 'databricks-uc', + swap=lambda: global_tracking_uri('databricks')) client = import_provider('kumo_databricks.session', 'databricks')._client_for(principal) return Target('databricks', 'databricks-uc', diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index 49d8a4b4e..502e1d3f1 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -106,9 +106,10 @@ class MlflowTracking: from for its duration (Databricks: ``DATABRICKS_HOST`` and a short-lived ``DATABRICKS_TOKEN``; SageMaker: the process-wide tracking URI), so do not call MLflow or Databricks yourself on other threads during a fit. - Ambient credentials (a notebook, job or instance role) change nothing, - except that a SageMaker target still sets the process-wide tracking URI - for each call, since the plugin signs and routes with it. + Ambient credentials (a notebook, job, profile or instance role) set no + credentials, but a Databricks or SageMaker target still sets the + process-wide tracking URI for each call: Unity Catalog downloads and the + SageMaker plugin route with it. ``run_id`` is the MLflow id of the latest fit run, once it has started, and ``run_url`` its link in the MLflow UI, where the server has one. @@ -217,6 +218,50 @@ def load_model(self, location: ModelRegistryLocation, return load_model(f'models:/{location.model_name}{reference}', target=self.target, **kwargs) + def register(self, result: Any, + location: ModelRegistryLocation) -> PublishedModel: + r"""Registers a fitted or loaded model (a `TrainingResult`) at + ``location``: the version that already holds its bundle, else a new + version of the model logged in a new run of this experiment.""" + from kumo.engine._manifest import stub + from kumo.engine.graph import GraphSnapshot + from kumo.engine.trainer._spec import read_training_spec + from kumo.engine.trainer.trainer import registry_pointers + + self._require_experiment() + client = self.target.client() + pointers = registry_pointers(result.id) + with _limited_http(frozenset(name for name in _HTTP_LIMITS + if name not in os.environ)): + with self.target.scope(): + held = _registered_version(client, location.model_name, + pointers) + info = None + if held is None: + spec = read_training_spec(result.bundle_dir) + snapshot = GraphSnapshot(spec.snapshot_id, Path(), + spec.graph_config or {}) + with self.target.scope(): + experiment_id = self._default_experiment or '0' + if self.experiment_name is not None: + experiment = client.get_experiment_by_name( + self.experiment_name) + experiment_id = (experiment.experiment_id + if experiment else + client.create_experiment( + self.experiment_name)) + run_id = client.create_run( + experiment_id, + tags={'kumo.bundle_id': str(result.id)}).info.run_id + info = _log_model(self.target, client, result.bundle_dir, + result.manifest or stub('model', result.id), + snapshot, name='model', run_id=run_id, + experiment_id=experiment_id, + serving_snapshot=None) + with self.target.scope(): + client.set_terminated(run_id, 'FINISHED') + return _register(self.target, client, info, location, pointers) + def start(self, record: RunRecord) -> None: self._require_experiment() self._reset() @@ -804,8 +849,10 @@ def _registered_version(client: Any, name: str, return None raise for version in page: - tags = version.tags or client.get_model_version( - name, version.version).tags + # Unity Catalog's search results carry tags as a method. + tags = version.tags + if not isinstance(tags, Mapping) or not tags: + tags = client.get_model_version(name, version.version).tags if (tags or {}).get(BUNDLE_ID_KEY) == bundle: return str(version.version) token = page.token diff --git a/packages/kumo-engine/kumo/engine/pquery/training_table.py b/packages/kumo-engine/kumo/engine/pquery/training_table.py index db3888293..82d616e11 100644 --- a/packages/kumo-engine/kumo/engine/pquery/training_table.py +++ b/packages/kumo-engine/kumo/engine/pquery/training_table.py @@ -148,9 +148,9 @@ def _store(self, settings: Any) -> None: self.export( TrainingTableOutputConfig(settings.root.connector, path)) except Exception as error: - warnings.warn( + raise RuntimeError( f"Could not store training table {self.id} under the " - f"artifact root: {error}", UserWarning, stacklevel=3) + f"artifact root: {error}") from error @classmethod def from_source_table( @@ -312,7 +312,13 @@ def _finish(table: TrainingTable, snapshot: GraphSnapshot, request: Any, _record_manifest(table, snapshot, request, seed, source_table=source_table) settings = resolve(mlops) if settings is not None and settings.wants('training_table'): - table._store(settings) + from kumo.engine._mlops import raise_on_error + + try: + table._store(settings) + except Exception as error: + raise_on_error(settings, error, 'Training table generation', + stacklevel=3) return table diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 94ad871f1..d3a818880 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -161,8 +161,14 @@ def predict(self, *args: Any, **kwargs: Any) -> 'PredictionResult': """Score with this model, as :meth:`Trainer.predict`.""" return self._model('predict').predict(*args, **kwargs) - def export(self, config: ModelOutputConfig) -> PublishedModel: - """Publish this model's bundle, as :meth:`Trainer.export`.""" + def export(self, config: Any) -> PublishedModel: + """Publish this model's bundle to a ``ModelOutputConfig``, or + register it at a `kumo.MLOps` ``registry`` through its tracker; a + registry that already holds the bundle returns that version.""" + from kumo.engine._mlops import MLOps + + if isinstance(config, MLOps): + return _register_result(self, config) _require_publisher(config) return _publish(self.bundle_dir, self.bundle_id, self.metrics, config) @@ -352,15 +358,25 @@ def _store_and_register( snapshot: Any = None, metrics: Mapping[str, float] = {}, warn: bool = True, + stored: Optional[BaseException] = None, ) -> Tuple[Optional[PublishedModel], Optional[BaseException]]: r"""Store a fitted bundle under the artifact root, log it in the session's run and register it, as MLOps ``settings`` say; returns the - registered version and the first failure.""" - if settings is not None and settings.wants('model'): - _store_model(bundle_dir, manifest, settings) + registered version and the first failure. A model whose bundle or + assets (``stored``, the asset failure) could not be stored is not + registered.""" + if stored is None and settings is not None and settings.wants('model'): + try: + _store_model(bundle_dir, manifest, settings) + except Exception as error: + stored = error + if stored is not None and warn: + warn_failure(f'{stored}; the model is not registered', stored, 3) registered, error = (None, None) if session is None else session.log_model( - bundle_dir, manifest, snapshot, _mlflow_registry(settings)) - registry = _connector_registry(settings) + bundle_dir, manifest, snapshot, + None if stored else _mlflow_registry(settings)) + error = stored or error + registry = None if stored else _connector_registry(settings) if registry is None: return registered, error errors: List[BaseException] = [] @@ -379,7 +395,7 @@ def _store_and_register( def _store_model(bundle_dir: Path, manifest: Optional[ModelManifest], settings: Any) -> None: r"""Copy the bundle to ``/models//`` unless a complete - copy is there; failures warn.""" + copy is there.""" from kumo_core.object_store_sink import holds_bundle, upload_bundle from pyarrow.fs import FileSelector @@ -398,9 +414,8 @@ def _store_model(bundle_dir: Path, manifest: Optional[ModelManifest], filesystem.delete_dir_contents(root, missing_dir_ok=True) upload_bundle(filesystem, root, str(bundle_dir)) except Exception as error: - warnings.warn( - f"Could not store model {manifest.id} under the artifact " - f"root: {error}", UserWarning, stacklevel=4) + raise RuntimeError(f"Could not store model {manifest.id} under the " + f"artifact root: {error}") from error def _published_holdout(trainer: Any) -> tuple[Optional[Path], Optional[bool]]: @@ -448,6 +463,19 @@ def _require_publisher(export_config: ModelOutputConfig) -> None: check_model_location(export_config) +def _register_result(result: TrainingResult, settings: Any) -> PublishedModel: + if settings.registry is None: + raise ValueError('export(cfg) registers at cfg.registry, which is ' + 'not set.') + register = getattr(settings.tracking, 'register', None) + if callable(register): + return register(result, settings.registry) + registry = cast(ModelOutputConfig, _connector_registry(settings)) + _require_publisher(registry) + return _publish(result.bundle_dir, result.bundle_id, result.metrics, + registry) + + def _mlflow_registry(settings: Any) -> Optional[ModelRegistryLocation]: if settings is None or not _logs_models(settings.tracking): return None @@ -602,12 +630,10 @@ def _fitted(self, method: str) -> TrainingResult: "this same Trainer instance.") return self._training_result - def export(self, config: ModelOutputConfig) -> PublishedModel: - """Publish the fitted or loaded model bundle.""" - training_result = self._fitted('export') - _require_publisher(config) - return _publish(training_result.bundle_dir, training_result.bundle_id, - training_result.metrics, config) + def export(self, config: Any) -> PublishedModel: + """Publish the fitted or loaded model bundle, as + :meth:`TrainingResult.export`.""" + return self._fitted('export').export(config) @_without_torch_scatter_hint def fit( @@ -675,6 +701,7 @@ def fit( None if settings is None else settings.tracking, warn=warn) events = EventEmitter( [*callbacks, session.on_event] if session.enabled else callbacks) + deps = Dependencies(settings=settings) fit_start = time.monotonic() try: started = { @@ -690,7 +717,7 @@ def fit( scoring_defaults=scoring_defaults, warm_start=warm_start, run_id=run_id, events=events, tracking=session, - deps=Dependencies(settings=settings)) + deps=deps) except BaseException as error: events.emit('fit_ended', params={ 'status': 'failed', @@ -706,7 +733,8 @@ def fit( result.bundle_dir, result.manifest or ( None if result.bundle_id is None else stub( 'model', result.bundle_id)), settings, session=session, - snapshot=snapshot, metrics=result.metrics, warn=warn) + snapshot=snapshot, metrics=result.metrics, warn=warn, + stored=(deps.errors or [None])[0]) if registered is not None: result = replace(result, registered_model=registered) self._training_result = result diff --git a/packages/kumo-engine/test/test_mlflow_model.py b/packages/kumo-engine/test/test_mlflow_model.py index 2f3cff81d..a62f90694 100644 --- a/packages/kumo-engine/test/test_mlflow_model.py +++ b/packages/kumo-engine/test/test_mlflow_model.py @@ -94,9 +94,17 @@ def test_fit_logs_registers_and_serves_the_model(table_and_snapshot, assert list(served) == ['ENTITY', 'TARGET_PRED'] assert len(served) == len(range(0, 400, 7)) assert load_model('models:/m@champion').result.id == result.id - assert kumo.Trainer.load_from_connector( - cfg, location=kumo.ModelRegistryLocation('m', alias='champion') - ).result.id == result.id + champion = kumo.Trainer.load_from_connector( + cfg, location=kumo.ModelRegistryLocation('m', alias='champion')) + assert champion.result.id == result.id + assert champion.export(cfg).version == str(version.version) + client = mlflow.MlflowClient() + assert len(client.search_model_versions("name = 'm'")) == 1 + other = kumo.MLOps(str(tmp_path / 'root'), tracking=cfg.tracking, + registry=kumo.ModelRegistryLocation('m2')) + first = champion.export(other) + assert first.version == '1' and champion.export(other) == first + assert load_model('models:/m2/1').result.id == result.id original = mlflow.models.Model.load diff --git a/packages/kumo-engine/test/test_mlflow_target.py b/packages/kumo-engine/test/test_mlflow_target.py index 27be764b0..80aeb006d 100644 --- a/packages/kumo-engine/test/test_mlflow_target.py +++ b/packages/kumo-engine/test/test_mlflow_target.py @@ -62,9 +62,13 @@ def test_ambient_identities_set_no_credentials(monkeypatch, tracking_uri): target = resolve(_connector(DatabricksIdentity())) before = dict(os.environ) with target.scope(): - assert dict(os.environ) == before - assert mlflow.get_tracking_uri() == tracking_uri - assert not target.explicit + # A profile or notebook identity: Unity Catalog registration + # downloads through the process-wide URI, so it points at Databricks. + assert mlflow.get_tracking_uri() == 'databricks' + assert {k: v for k, v in os.environ.items() + if k.startswith('DATABRICKS_')} == {'DATABRICKS_TOKEN': 'mine'} + assert mlflow.get_tracking_uri() == tracking_uri + assert dict(os.environ) == before for target in (resolve(SageMakerMLflow(ARN, S3Identity())), resolve(SageMakerMLflow(ARN))): @@ -297,3 +301,41 @@ def test_the_smoke_script_compiles(tmp_path): doraise=True) return pytest.skip('examples/mlops is not in this checkout') + + +def test_a_service_principal_client_ignores_a_token_lent_to_mlflow( + monkeypatch): + from kumo_databricks import session + + from kumo.engine.integrations._mlflow_target import _principal + + made = [] + monkeypatch.setattr(session, 'WorkspaceClient', + lambda **kwargs: made.append(kwargs)) + monkeypatch.setenv('DATABRICKS_TOKEN', 'short-lived') + session._client_for(_principal(DatabricksIdentity( + host='https://dbc.example', client_id='sp', client_secret='secret'))) + assert made[0]['auth_type'] == 'oauth-m2m' + + +def test_unity_catalog_search_results_with_tags_as_a_method(): + from kumo.engine.integrations.mlflow import _registered_version + + class Search: + def __init__(self, version): + self.version = version + + def tags(self): + raise AssertionError('not a mapping') + + class Page(list): + token = None + + held = {'1': {'kumo_bundle_id': 'other'}, '2': {'kumo_bundle_id': 'b'}} + client = SimpleNamespace( + search_model_versions=lambda *a, **k: Page([Search('1'), + Search('2')]), + get_model_version=lambda name, version: SimpleNamespace( + tags=held[version])) + assert _registered_version(client, 'm', {'kumo_bundle_id': 'b'}) == '2' + assert _registered_version(client, 'm', {'kumo_bundle_id': 'x'}) is None diff --git a/packages/kumo-engine/test/test_model_agnostic_mlops.py b/packages/kumo-engine/test/test_model_agnostic_mlops.py index a73a30061..4bfa258ad 100644 --- a/packages/kumo-engine/test/test_model_agnostic_mlops.py +++ b/packages/kumo-engine/test/test_model_agnostic_mlops.py @@ -4,11 +4,13 @@ """The MLOps seams take any bundle, run and predictor, not only the Trainer's.""" +import io import json import os import subprocess import sys import textwrap +import warnings from dataclasses import replace from pathlib import Path from types import SimpleNamespace @@ -283,3 +285,82 @@ def test_a_model_serves_from_its_bundle_snapshot_and_assets_alone( capture_output=True, text=True, timeout=600) assert "MissingDependencyError: The model needs text_encoder 'colors'" \ in done.stderr + + +class _PrefixOnlyS3: + r"""An S3 filesystem whose credentials cover only the root's prefix: + HeadBucket and parent marker objects are refused.""" + type_name = 's3' + + def __init__(self): + self.objects = {} + + def get_file_info(self, path): + return SimpleNamespace(is_file=path in self.objects) + + def create_dir(self, path, recursive=True): + raise OSError('AWS Error ACCESS_DENIED during HeadBucket operation') + + def open_output_stream(self, path): + objects = self.objects + + class Writer(io.BytesIO): + def close(self): + objects[path] = self.getvalue() + super().close() + + return Writer() + + +def test_an_asset_is_stored_on_s3_with_prefix_scoped_rights( + tmp_path, monkeypatch): + from kumo.engine import connector + + s3 = _PrefixOnlyS3() + monkeypatch.setattr(connector, 'file_store', + lambda connector, path: (s3, f'bucket/kumo/{path}')) + (tmp_path / 'vocab').mkdir() + (tmp_path / 'vocab' / 'words.txt').write_text('a b c') + deps = Dependencies(settings=kumo.MLOps(str(tmp_path / 'root'))) + record = deps.require(FileAsset('toy', 'vocab', tmp_path / 'vocab')) + assert deps.errors == [] and record.location + assert sorted(s3.objects) == [ + f'bucket/kumo/assets/toy/{record.sha256}/_provenance.json', + f'bucket/kumo/assets/toy/{record.sha256}/words.txt'] + + +@pytest.mark.parametrize('on_error', ['warn', 'raise']) +def test_a_model_that_was_not_stored_is_not_registered( + tmp_path, monkeypatch, on_error): + from kumo.engine._mlops import raise_on_error + from kumo.engine.trainer import trainer + + bundle, manifest = _toy_bundle(tmp_path / 'bundle') + published = [] + sink = SimpleNamespace( + publish_model=lambda *a, **k: published.append(a) or 1 / 0) + settings = kumo.MLOps(str(tmp_path / 'root'), on_error=on_error, + tracking=_RegistryTracker( + SimpleNamespace(sink_backend=sink)), + registry=ModelRegistryLocation('TOY')) + + def refuse(*args): + raise RuntimeError('Could not store model m under the artifact ' + 'root: validate: oauth and pat') + + monkeypatch.setattr(trainer, '_store_model', refuse) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter('always') + registered, error = _store_and_register( + bundle, manifest, settings, warn=on_error == 'warn') + assert registered is None and published == [] + assert 'oauth and pat' in str(error) + assert any('not registered' in str(w.message) + for w in caught) is (on_error == 'warn') + if on_error == 'raise': + with pytest.raises(RuntimeError, match='oauth and pat'): + raise_on_error(settings, error, 'Trainer.fit') + asset_failure = OSError('ACCESS_DENIED') + registered, error = _store_and_register(bundle, manifest, settings, + warn=False, stored=asset_failure) + assert registered is None and error is asset_failure diff --git a/packages/kumo-engine/test/test_snapshot_export.py b/packages/kumo-engine/test/test_snapshot_export.py index 0e67c3286..e5e3cc07c 100644 --- a/packages/kumo-engine/test/test_snapshot_export.py +++ b/packages/kumo-engine/test/test_snapshot_export.py @@ -794,3 +794,17 @@ def test_loading_a_missing_export_names_its_full_path(tmp_path, exports): with pytest.raises(FileNotFoundError, match=f'No complete snapshot export at {missing}: '): _load(exports, 'snapshots/graph/version') + + +@pytest.mark.parametrize('on_error', ['warn', 'raise']) +def test_a_failed_store_follows_on_error(tmp_path, scratch, graph, + monkeypatch, on_error): + monkeypatch.setattr(snapshot_io, 'export', + Mock(side_effect=OSError('ACCESS_DENIED'))) + cfg = kumo.MLOps(str(tmp_path / 'root'), on_error=on_error) + if on_error == 'raise': + with pytest.raises(RuntimeError, match='Could not store snapshot'): + graph.snapshot(mlops=cfg) + else: + with pytest.warns(UserWarning, match='ACCESS_DENIED'): + graph.snapshot(mlops=cfg) From 58269b24b0e6b854a7edf38d833c88d597500c09 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 00:36:27 +0000 Subject: [PATCH 022/278] Give your own MLflow calls the tracker's identity and fix the promotion recipes MlflowTracking.client() yields an MlflowClient for the tracker's server with its credentials and tracking URI set for the block, so promotion and run lookups reach the server as the identity kumo-engine uses. The promotion recipes handle the first promotion, and the Snowflake one moves candidate off the new version. kumo.TextEmbedding is exported, and a text column left unencoded says when no word-vector model is installed. Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_core/encoder_selection.py | 15 +++++++- .../kumo-core/tests/test_encoder_selection.py | 14 ++++++- packages/kumo-engine/docs/mlops.md | 37 +++++++++++++------ packages/kumo-engine/kumo/engine/__init__.py | 2 + .../kumo/engine/integrations/mlflow.py | 11 ++++++ .../kumo-engine/test/test_embedding_models.py | 7 +++- .../kumo-engine/test/test_mlflow_target.py | 22 +++++++++++ 7 files changed, 92 insertions(+), 16 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/encoder_selection.py b/packages/kumo-core/src/kumo_core/encoder_selection.py index d0956b1cb..e0933c42e 100644 --- a/packages/kumo-core/src/kumo_core/encoder_selection.py +++ b/packages/kumo-core/src/kumo_core/encoder_selection.py @@ -440,12 +440,25 @@ def select_text_encoder( logger.warning( "Text column %r will not be encoded by default: its values are too " "short or have too few known words. To encode it anyway, set its " - "encoder to TextEmbedding().", + "encoder to TextEmbedding().%s", column_name, + _no_default_model(), ) return Null() +def _no_default_model() -> str: + from kumo_core.embedding_models import ( + DEFAULT_MODEL, + INSTALL_DEFAULT_MODEL, + _installed, + ) + if DEFAULT_MODEL in _installed(): + return "" + return (f" No word-vector model is installed for it either: " + f"{INSTALL_DEFAULT_MODEL}, or pass embedding_models= to Trainer.") + + def select_sequence_encoder( dtype: Dtype, stats: _EncoderStats, diff --git a/packages/kumo-core/tests/test_encoder_selection.py b/packages/kumo-core/tests/test_encoder_selection.py index 5edf9476b..a05c5b186 100644 --- a/packages/kumo-core/tests/test_encoder_selection.py +++ b/packages/kumo-core/tests/test_encoder_selection.py @@ -298,7 +298,14 @@ def test_target_selection_still_requires_numerical_stats(): select_target_encoder(Stype.numerical, {}) -def test_text_left_unencoded_is_reported_by_every_inference(caplog): +@pytest.mark.parametrize("installed", [True, False]) +def test_text_left_unencoded_is_reported_by_every_inference( + caplog, monkeypatch, installed): + from kumo_core import embedding_models + + monkeypatch.setattr( + embedding_models, "_installed", + lambda: {embedding_models.DEFAULT_MODEL: None} if installed else {}) stats = { stat_type: _EncoderStat(value=value, approx=True) for stat_type, value in ( @@ -314,7 +321,10 @@ def test_text_left_unencoded_is_reported_by_every_inference(caplog): assert [record.getMessage() for record in caplog.records] == [ "Text column 'short' will not be encoded by default: its values are " "too short or have too few known words. To encode it anyway, set its " - "encoder to TextEmbedding()." + "encoder to TextEmbedding()." + + ("" if installed else " No word-vector model is installed for it " + "either: pip install \"kumo-engine[text]\", or pass embedding_models= " + "to Trainer.") ] * 2 diff --git a/packages/kumo-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index 1553fef8b..a430f6caf 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -58,6 +58,16 @@ with kumo.mlops(cfg): print(result.registered_model.version, result.tracking.url) ``` +Tables on S3 are read through an `S3Connector`: + +```python +s3 = kumo.S3Connector(root_dir="s3:///data") +users = kumo.Table.from_source_table(s3["users.parquet"], primary_key="USER_ID").infer_metadata() +orders = kumo.Table.from_source_table(s3["orders.parquet"], primary_key="ORDER_ID", time_column="ORDER_TIME").infer_metadata() +graph = kumo.Graph(tables={"users": users, "orders": orders.with_column("USER_ID", stype="ID")}, + edges=[kumo.Edge(src_table="orders", fkey="USER_ID", dst_table="users")]) +``` + ### Snowflake Tracking and the Snowflake Model Registry through one connector (`kumo-engine[snowflake,s3]`): @@ -151,15 +161,18 @@ An alias is a letter or underscore followed by letters, digits or underscores; ` Register every fit under a candidate alias and move the serving alias yourself once the new model is better, comparing `result.metrics` with the champion's: ```python -champion = kumo.Trainer.load_from_connector(cfg, location=kumo.ModelRegistryLocation(name, alias="champion")) -if result.metrics["test_auroc"] > champion.result.metrics["test_auroc"]: - mlflow.MlflowClient().set_registered_model_alias(name, "champion", result.registered_model.version) +with cfg.tracking.client() as client: # the tracker's MLflow server, as its identity + first = "champion" not in client.get_registered_model(name).aliases + if first or result.metrics["test_auroc"] > kumo.Trainer.load_from_connector( + cfg, location=kumo.ModelRegistryLocation(name, alias="champion")).result.metrics["test_auroc"]: + client.set_registered_model_alias(name, "champion", result.registered_model.version) ``` -In the Snowflake Model Registry a version holds at most one alias, so unset it first: +In the Snowflake Model Registry a version holds at most one alias: take `champion` off the old version (there is none on the first promotion) and `candidate` off the new one, then set it: ```sql -ALTER MODEL ML.MODELS.CHURN_GNN VERSION champion UNSET ALIAS; +ALTER MODEL ML.MODELS.CHURN_GNN VERSION champion UNSET ALIAS; -- skip on the first promotion +ALTER MODEL ML.MODELS.CHURN_GNN VERSION V3 UNSET ALIAS; -- V3 holds candidate ALTER MODEL ML.MODELS.CHURN_GNN VERSION V3 SET ALIAS = champion; ``` @@ -206,16 +219,16 @@ workspace = DatabricksConnector(catalog="main", schema_name="ml", warehouse_id=" identity=DatabricksIdentity(host="https://", client_id="", client_secret="")) databricks = MlflowTracking("/Users//users_churn", target=workspace) sagemaker = MlflowTracking("users-churn", target=SageMakerMLflow("", identity=S3Identity(role_arn=""))) -snowflake = SnowflakeTracking(SnowflakeConnector(warehouse="", database="ML", schema_name="RUNS", +snowflake = SnowflakeTracking(SnowflakeConnector(warehouse="", database="ML", schema_name="RUNS", role="", identity=SnowflakeIdentity(account="", username="", private_key="")), experiment_name="USERS_CHURN") ``` -`load_model(uri, target=...)` and `log_model(..., target=...)` take the same targets. A `DatabricksConnector` without an identity uses a notebook, a job or a `~/.databrickscfg` profile with no environment variables to set. MLflow has no per-client credentials, so with explicit credentials each MLflow call kumo-engine makes sets `DATABRICKS_HOST` and a short-lived `DATABRICKS_TOKEN` (never the secret), or the SageMaker tracking URI and AWS session, and restores them afterwards: don't call MLflow or Databricks yourself on other threads during a fit. `examples/mlops/smoke.py databricks|sagemaker|snowflake` is a manual live check per platform. +`load_model(uri, target=...)` and `log_model(..., target=...)` take the same targets. A `DatabricksConnector` without an identity uses a notebook, a job or a `~/.databrickscfg` profile with no environment variables to set. MLflow has no per-client credentials, so with explicit credentials each MLflow call kumo-engine makes sets `DATABRICKS_HOST` and a short-lived `DATABRICKS_TOKEN` (never the secret), or the SageMaker tracking URI and AWS session, and restores them afterwards: don't call MLflow or Databricks yourself on other threads during a fit. For your own MLflow calls as that identity, use `with tracking.client() as client:`, an `MlflowClient` for the tracker's server with the same state set for the block (`mlflow.search_runs` inside it reaches that server too). A `SnowflakeConnector` without `role=` uses the user's default role. `examples/mlops/smoke.py databricks|sagemaker|snowflake` is a manual live check per platform. ### SageMaker managed MLflow -1. Create a tracking server once (15-25 minutes): `aws sagemaker create-mlflow-tracking-server --tracking-server-name --artifact-store-uri s3:///mlflow --role-arn `; the server role needs read and write on that prefix. +1. Create a tracking server once (15-25 minutes): `aws sagemaker create-mlflow-tracking-server --tracking-server-name --artifact-store-uri s3:///mlflow --role-arn `; the server role needs read and write on that prefix. A stopped server starts with `aws sagemaker start-mlflow-tracking-server --tracking-server-name ` (about 20 minutes); `kumo.MLOps` fails until `aws sagemaker describe-mlflow-tracking-server --tracking-server-name ` reports `Started`. 2. Give the job's identity `sagemaker-mlflow:*` on the server ARN and S3 read and write on the artifact store and the kumo artifact root. Rights scoped to the root's prefix (`s3:GetObject`, `s3:PutObject` and `s3:DeleteObject` on `//*`, `s3:ListBucket` limited to that prefix) suffice: kumo-engine makes no bucket-level calls. 3. Install `kumo-engine[mlflow,sagemaker]` and pass `MlflowTracking(..., target=SageMakerMLflow(arn))`; the registry is the same server. 4. MLflow's run links do not open on SageMaker: open the UI from SageMaker Studio, or at the URL `aws sagemaker create-presigned-mlflow-tracking-server-url --tracking-server-name ` prints. A tracked fit warns with this command, and `result.tracking.url` is `None`. @@ -225,7 +238,7 @@ snowflake = SnowflakeTracking(SnowflakeConnector(warehouse="", database="ML" ### Install - A bare `kumo-engine[mlflow]` with `--find-links ` can resolve to an unrelated `kumo-engine` on PyPI; install the engine wheel by path: `pip install --find-links "/[mlflow]"`. -- **GPU:** on a plain GPU VM or SageMaker image, install a torch build that matches the NVIDIA driver first, e.g. `pip install --index-url https://download.pytorch.org/whl/cu126 "torch==2.12.*" torchvision`. PyPI's default torch (cu130) needs driver ≥ 580; cu126 runs on ≥ 525.60. With `accelerator="auto"`, a torch that cannot use the driver runs on CPU: check `torch.cuda.is_available()`. +- **GPU:** on a plain GPU VM or SageMaker image, install a torch build that matches the NVIDIA driver first, e.g. `pip install --index-url https://download.pytorch.org/whl/cu126 "torch==2.12.*" torchvision`. PyPI's default torch is built for a recent CUDA (`torch.version.cuda`; CUDA 13 needs driver ≥ 580); cu126 runs on ≥ 525.60. With `accelerator="auto"`, a torch that cannot use the driver runs on CPU: check `torch.cuda.is_available()`. - **Databricks:** Runtime 17.3 LTS ML or later keeps its torch, torchvision, MLflow and pyarrow (kumo-engine accepts pyarrow 16 through 24; a custom CA bundle for S3 needs 21 or later). On 16.4 LTS ML pip replaces torch and breaks torchvision. The parent folder of an experiment path must exist. Registering to Unity Catalog from outside Databricks also needs `boto3`. - **MLflow:** the `mlflow` extra installs `mlflow-skinny` 3.0.1 or later; a 3.x client works against SageMaker's MLflow 3.0 servers. - **SageMaker serving:** MLflow's stock serving containers cannot install kumo-engine; on SageMaker endpoints serve with `-m /opt/ml/model -p 8080`. GPU endpoints need a driver for CUDA 12.6 (set `InferenceAmiVersion`, e.g. `al2023-ami-sagemaker-inference-gpu-4-1`); batch transform takes JSON (`dataframe_split` plus `params`). @@ -258,13 +271,15 @@ Every artifact carries `_provenance.json`, metadata only (the underscore keeps S ```python store = kumo.EmbeddingModelStore() -store.add("colors", kumo.GloVeVectors("/data/vectors")) -plan.column_processing.encoder_overrides = {"users.bio": TextEmbedding("colors")} +store.add("glove100", kumo.GloVeVectors("glove.6B.zip", dim=100)) # the official GloVe release .zip, unchanged +plan.column_processing.encoder_overrides = {"users.bio": kumo.TextEmbedding("glove100")} with kumo.mlops(cfg): result = kumo.Trainer(plan, embedding_models=store).fit(table, snapshot=snapshot) trainer = kumo.Trainer.load_from_connector(cfg, path=f"models/{result.id}") # anywhere, no store ``` +A text column `infer_metadata()` reads as categorical (short values from a small vocabulary) takes a text encoder once its stype is text: `users.with_column("bio", stype="text")`. + ### Moving snapshots between machines - **Layout:** `export` writes `///`, holding `kumo_snapshot.json` and one directory per export generation (`index/` and `_provenance.json`). diff --git a/packages/kumo-engine/kumo/engine/__init__.py b/packages/kumo-engine/kumo/engine/__init__.py index b9e467763..642cdece9 100644 --- a/packages/kumo-engine/kumo/engine/__init__.py +++ b/packages/kumo-engine/kumo/engine/__init__.py @@ -21,6 +21,7 @@ TokenEmbeddingModel, ) from kumo_core.utils.nofile import ensure_nofile +from kumoapi.encoder import TextEmbedding from kumoapi.graph import Edge from kumoapi.plugin import ModelRegistryLocation from kumoapi.table import Column @@ -69,6 +70,7 @@ "configure", "EmbeddingModelStore", "GloVeVectors", + "TextEmbedding", "TokenEmbeddingModel", "ModelRegistryLocation", "Connector", diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index 502e1d3f1..96c616790 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -43,6 +43,7 @@ from kumo.engine.integrations._mlflow_target import ( SAGEMAKER_ARN, SageMakerMLflow, + global_tracking_uri, ) from kumo.engine.integrations._mlflow_target import resolve as resolve_target @@ -161,6 +162,16 @@ def run_url(self) -> Optional[str]: return None return self.target.run_url(self._experiment_id, self.run_id) + @contextmanager + def client(self) -> Iterator[Any]: + r"""An ``MlflowClient`` for this tracker's server, with its + credentials and tracking URI set for the block, so your own MLflow + calls (the client's, or ``mlflow.search_runs``) reach the server + kumo-engine records to, as the same identity.""" + with self.target.scope(), global_tracking_uri( + self.target.tracking_uri): + yield self.target.client() + def check_model_name(self, name: str) -> None: r"""Raises :class:`ValueError` unless this tracker's registry accepts ``name``.""" diff --git a/packages/kumo-engine/test/test_embedding_models.py b/packages/kumo-engine/test/test_embedding_models.py index ee2eb68da..a51475c8a 100644 --- a/packages/kumo-engine/test/test_embedding_models.py +++ b/packages/kumo-engine/test/test_embedding_models.py @@ -14,7 +14,6 @@ from kumo_core import embedding_models from kumo_core.embedding_models import EmbeddingModelNotFoundError, Release from kumo_core.glove import TokenEmbeddingTable -from kumoapi.encoder import TextEmbedding from kumoapi.typing import Stype from kumo.engine import ( @@ -24,6 +23,7 @@ Graph, PredictiveQuery, Table, + TextEmbedding, Trainer, ) @@ -148,7 +148,10 @@ def test_a_missing_model_fails_fit_before_any_work(fitted, monkeypatch): def test_each_fit_reports_what_inference_leaves_unencoded(fitted): unencoded = ("Text column 'users.code' will not be encoded by default: its " "values are too short or have too few known words. To " - "encode it anyway, set its encoder to TextEmbedding().") + "encode it anyway, set its encoder to TextEmbedding(). No " + "word-vector model is installed for it either: pip install " + "\"kumo-engine[text]\", or pass embedding_models= to " + "Trainer.") assert fitted[-1] == { 'tables': [], 'failed fit': [unencoded], diff --git a/packages/kumo-engine/test/test_mlflow_target.py b/packages/kumo-engine/test/test_mlflow_target.py index 80aeb006d..4b9b360a8 100644 --- a/packages/kumo-engine/test/test_mlflow_target.py +++ b/packages/kumo-engine/test/test_mlflow_target.py @@ -139,6 +139,28 @@ def test_an_aws_access_key_signs_sagemaker_calls_and_uploads( assert isinstance(session, boto3.Session) +def test_your_own_mlflow_calls_reach_the_trackers_server_as_its_identity( + monkeypatch, tracking_uri, tmp_path): + plugin = pytest.importorskip('sagemaker_mlflow.session') + from kumo.engine.integrations.mlflow import MlflowTracking + + monkeypatch.setattr(MlflowTracking, '_require_experiment', lambda _: 0) + tracking = MlflowTracking('e', target=SageMakerMLflow( + ARN, S3Identity(access_key_id='AKIA', secret_access_key='s'))) + with tracking.client() as client: + assert client._tracking_client.tracking_uri == ARN + assert mlflow.get_tracking_uri() == ARN + assert plugin._get_current_session().get_credentials( + ).access_key == 'AKIA' + assert mlflow.get_tracking_uri() == tracking_uri + + other = f'sqlite:///{tmp_path / "other.db"}' + with MlflowTracking('e', target=other).client() as client: + assert mlflow.get_tracking_uri() == other + assert client._tracking_client.tracking_uri == other + assert mlflow.get_tracking_uri() == tracking_uri + + def test_explicit_calls_take_turns(monkeypatch): target = target_module.Target('databricks', 'databricks-uc', swap=lambda: target_module.environ( From 70d4adb6b27b72edb2baa1b38bba429ac976bac6 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 00:58:59 +0000 Subject: [PATCH 023/278] Require MLflow 3.5 for the mlflow extra and test against that floor Loading a model by alias passes download_artifacts(registry_uri=), which MLflow added in 3.5, so on Databricks Runtime 17.3 (mlflow-skinny 3.0.1) it failed with a TypeError. The extra now requires mlflow-skinny 3.5, which %pip installs over the runtime's, and a CI job runs the MLflow suites at that floor. WriteMode is exported from kumo.engine, and the Databricks quickstart shows its install line. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/testing.yml | 35 ++++++++++++++++++++ packages/kumo-engine/docs/mlops.md | 6 ++-- packages/kumo-engine/kumo/engine/__init__.py | 2 ++ packages/kumo-engine/pyproject.toml | 5 +-- uv.lock | 2 +- 5 files changed, 44 insertions(+), 6 deletions(-) diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml index 179decec6..1c59eb78b 100644 --- a/.github/workflows/testing.yml +++ b/.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] diff --git a/packages/kumo-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index a430f6caf..5b9dccab0 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -26,7 +26,7 @@ with kumo.mlops(cfg): ### Databricks -In a notebook or job on Databricks Runtime 17.3 LTS ML or later: +In a notebook or job on Databricks Runtime 17.3 LTS ML or later, after `%pip install "/[databricks,mlflow,text]"` and `dbutils.library.restartPython()`: ```python import kumo.engine as kumo @@ -239,8 +239,8 @@ snowflake = SnowflakeTracking(SnowflakeConnector(warehouse="", database="ML" - A bare `kumo-engine[mlflow]` with `--find-links ` can resolve to an unrelated `kumo-engine` on PyPI; install the engine wheel by path: `pip install --find-links "/[mlflow]"`. - **GPU:** on a plain GPU VM or SageMaker image, install a torch build that matches the NVIDIA driver first, e.g. `pip install --index-url https://download.pytorch.org/whl/cu126 "torch==2.12.*" torchvision`. PyPI's default torch is built for a recent CUDA (`torch.version.cuda`; CUDA 13 needs driver ≥ 580); cu126 runs on ≥ 525.60. With `accelerator="auto"`, a torch that cannot use the driver runs on CPU: check `torch.cuda.is_available()`. -- **Databricks:** Runtime 17.3 LTS ML or later keeps its torch, torchvision, MLflow and pyarrow (kumo-engine accepts pyarrow 16 through 24; a custom CA bundle for S3 needs 21 or later). On 16.4 LTS ML pip replaces torch and breaks torchvision. The parent folder of an experiment path must exist. Registering to Unity Catalog from outside Databricks also needs `boto3`. -- **MLflow:** the `mlflow` extra installs `mlflow-skinny` 3.0.1 or later; a 3.x client works against SageMaker's MLflow 3.0 servers. +- **Databricks:** Runtime 17.3 LTS ML or later keeps its torch, torchvision and pyarrow (kumo-engine accepts pyarrow 16 through 24; a custom CA bundle for S3 needs 21 or later). On 16.4 LTS ML pip replaces torch and breaks torchvision. The parent folder of an experiment path must exist. Registering to Unity Catalog from outside Databricks also needs `boto3`. +- **MLflow:** the `mlflow` extra installs `mlflow-skinny` 3.5 or later (Databricks Runtime 17.3 ships 3.0.1, which `%pip` upgrades); a 3.x client works against SageMaker's MLflow 3.0 servers. - **SageMaker serving:** MLflow's stock serving containers cannot install kumo-engine; on SageMaker endpoints serve with `-m /opt/ml/model -p 8080`. GPU endpoints need a driver for CUDA 12.6 (set `InferenceAmiVersion`, e.g. `al2023-ami-sagemaker-inference-gpu-4-1`); batch transform takes JSON (`dataframe_split` plus `params`). ### Exporting and loading artifacts diff --git a/packages/kumo-engine/kumo/engine/__init__.py b/packages/kumo-engine/kumo/engine/__init__.py index 642cdece9..ca5a6b00f 100644 --- a/packages/kumo-engine/kumo/engine/__init__.py +++ b/packages/kumo-engine/kumo/engine/__init__.py @@ -23,6 +23,7 @@ from kumo_core.utils.nofile import ensure_nofile from kumoapi.encoder import TextEmbedding from kumoapi.graph import Edge +from kumoapi.jobs import WriteMode from kumoapi.plugin import ModelRegistryLocation from kumoapi.table import Column @@ -71,6 +72,7 @@ "EmbeddingModelStore", "GloVeVectors", "TextEmbedding", + "WriteMode", "TokenEmbeddingModel", "ModelRegistryLocation", "Connector", diff --git a/packages/kumo-engine/pyproject.toml b/packages/kumo-engine/pyproject.toml index b051564ae..c55d4f8f5 100644 --- a/packages/kumo-engine/pyproject.toml +++ b/packages/kumo-engine/pyproject.toml @@ -71,9 +71,10 @@ text=[ gcs=[ "kumo-gcs>=0.1.0.dev1", ] -# nvidia-ml-py adds GPU system metrics. +# mlflow-skinny 3.5 added download_artifacts(registry_uri=), which loading +# by alias uses; nvidia-ml-py adds GPU system metrics. mlflow=[ - "mlflow-skinny>=3.0.1,<4", + "mlflow-skinny>=3.5,<4", "nvidia-ml-py", ] # SageMaker managed MLflow: 0.5.0 added the per-thread boto3 session diff --git a/uv.lock b/uv.lock index 79523d2dd..321bd266a 100644 --- a/uv.lock +++ b/uv.lock @@ -3541,7 +3541,7 @@ requires-dist = [ { name = "kumo-glove-6b-50d", marker = "extra == 'text'", editable = "tools/glove-6b-50d" }, { name = "kumo-s3", marker = "extra == 's3'", editable = "packages/kumo-s3" }, { name = "kumo-snowflake", extras = ["sdk"], marker = "extra == 'snowflake'", editable = "packages/kumo-snowflake" }, - { name = "mlflow-skinny", marker = "extra == 'mlflow'", specifier = ">=3.0.1,<4" }, + { name = "mlflow-skinny", marker = "extra == 'mlflow'", specifier = ">=3.5,<4" }, { name = "numpy", marker = "extra == 'test'", specifier = ">=1.26,<3" }, { name = "nvidia-ml-py", marker = "extra == 'mlflow'" }, { name = "pandas" }, From 14bca8e26b61e0ca1790051c5ce9cc9cab2d7b44 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 00:59:46 +0000 Subject: [PATCH 024/278] Give a stored snapshot's location: MLOps.snapshot_uri The pyfunc's snapshot param takes the snapshot's location under the artifact root, which users rebuilt by hand from the snapshot id. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/docs/mlops.md | 2 +- packages/kumo-engine/kumo/engine/_mlops.py | 5 +++++ packages/kumo-engine/test/test_mlops.py | 3 +++ 3 files changed, 9 insertions(+), 1 deletion(-) diff --git a/packages/kumo-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index 5b9dccab0..de9348704 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -200,7 +200,7 @@ A model registered through `MlflowTracking` is an MLflow pyfunc: it serves with ```python model = mlflow.pyfunc.load_model("models:/main.ml.users_churn@champion") scores = model.predict(pd.DataFrame({"ENTITY": [1, 2, 3]}), params={ - "snapshot": "/Volumes/main/ml/kumo/snapshots//", + "snapshot": cfg.snapshot_uri(snapshot), # "/Volumes/main/ml/kumo/snapshots//" "prediction_time": "2024-06-01"}) # temporal queries ``` diff --git a/packages/kumo-engine/kumo/engine/_mlops.py b/packages/kumo-engine/kumo/engine/_mlops.py index 0855838c4..fb6e501fe 100644 --- a/packages/kumo-engine/kumo/engine/_mlops.py +++ b/packages/kumo-engine/kumo/engine/_mlops.py @@ -148,6 +148,11 @@ def _check_registry(self) -> None: '...).') check_model_location(ModelOutputConfig(connector, location=registry)) + def snapshot_uri(self, snapshot: Any) -> str: + r"""Where a `GraphSnapshot`, or a snapshot id, is stored under this + artifact root, e.g. for the pyfunc's ``snapshot`` param.""" + return self.root.uri(snapshot_path(snapshot)) + def wants(self, kind: str) -> bool: r"""Whether artifacts of ``kind`` are stored under the root.""" return kind in self.kinds diff --git a/packages/kumo-engine/test/test_mlops.py b/packages/kumo-engine/test/test_mlops.py index aeca7d9e9..53f929086 100644 --- a/packages/kumo-engine/test/test_mlops.py +++ b/packages/kumo-engine/test/test_mlops.py @@ -75,6 +75,9 @@ def test_a_root_is_a_uri_or_a_connector_and_path(tmp_path): assert pair.root.location('models/m') == 'volume/kumo/models/m' assert pair.wants('model') and not pair.wants('snapshot') assert not kumo.MLOps(artifact_root=str(tmp_path), store=False).kinds + assert uri.snapshot_uri('ab12@cd34') == f'{tmp_path}/snapshots/ab12/cd34' + assert kumo.MLOps('s3://bucket/kumo').snapshot_uri( + SimpleNamespace(id='ab12@cd34')) == 's3://bucket/kumo/snapshots/ab12/cd34' def test_a_connector_root_outside_its_root_dir_fails_when_built(): From f8d28d087686784a9d4240a9b8b7708f41de470a Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 00:59:58 +0000 Subject: [PATCH 025/278] Say where Python's temporary files go next to the scratch directory Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/docs/mlops.md | 1 + 1 file changed, 1 insertion(+) diff --git a/packages/kumo-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index de9348704..7db3f3f77 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -240,6 +240,7 @@ snowflake = SnowflakeTracking(SnowflakeConnector(warehouse="", database="ML" - A bare `kumo-engine[mlflow]` with `--find-links ` can resolve to an unrelated `kumo-engine` on PyPI; install the engine wheel by path: `pip install --find-links "/[mlflow]"`. - **GPU:** on a plain GPU VM or SageMaker image, install a torch build that matches the NVIDIA driver first, e.g. `pip install --index-url https://download.pytorch.org/whl/cu126 "torch==2.12.*" torchvision`. PyPI's default torch is built for a recent CUDA (`torch.version.cuda`; CUDA 13 needs driver ≥ 580); cu126 runs on ≥ 525.60. With `accelerator="auto"`, a torch that cannot use the driver runs on CPU: check `torch.cuda.is_available()`. - **Databricks:** Runtime 17.3 LTS ML or later keeps its torch, torchvision and pyarrow (kumo-engine accepts pyarrow 16 through 24; a custom CA bundle for S3 needs 21 or later). On 16.4 LTS ML pip replaces torch and breaks torchvision. The parent folder of an experiment path must exist. Registering to Unity Catalog from outside Databricks also needs `boto3`. +- **Disk:** `kumo.configure(scratch_dir=...)` holds snapshots, tables and bundles; Python's own temporary files (multiprocessing) follow `TMPDIR`, so point both at a large local disk. - **MLflow:** the `mlflow` extra installs `mlflow-skinny` 3.5 or later (Databricks Runtime 17.3 ships 3.0.1, which `%pip` upgrades); a 3.x client works against SageMaker's MLflow 3.0 servers. - **SageMaker serving:** MLflow's stock serving containers cannot install kumo-engine; on SageMaker endpoints serve with `-m /opt/ml/model -p 8080`. GPU endpoints need a driver for CUDA 12.6 (set `InferenceAmiVersion`, e.g. `al2023-ami-sagemaker-inference-gpu-4-1`); batch transform takes JSON (`dataframe_split` plus `params`). From a2c4398e17f1fc1ba893e30471c8c7b0a25fd68d Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 22:00:07 +0000 Subject: [PATCH 026/278] Check a root's write rights without leaving an empty S3 object behind pyarrow's S3 delete recreates an emptied parent as a zero-byte / object, so the write check left one at a new artifact root. The S3 sink deletes the probe through boto3 instead. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_mlops.py | 6 +++++- packages/kumo-s3/src/kumo_s3/sink.py | 16 ++++++++++++++++ packages/kumo-s3/tests/test_sink_delete.py | 22 ++++++++++++++++++++++ 3 files changed, 43 insertions(+), 1 deletion(-) create mode 100644 packages/kumo-s3/tests/test_sink_delete.py diff --git a/packages/kumo-engine/kumo/engine/_mlops.py b/packages/kumo-engine/kumo/engine/_mlops.py index 120b71792..e38d35329 100644 --- a/packages/kumo-engine/kumo/engine/_mlops.py +++ b/packages/kumo-engine/kumo/engine/_mlops.py @@ -134,7 +134,11 @@ def _require_writable(root: _Root) -> None: except FileNotFoundError: # a store with directories, before first use filesystem.create_dir(posixpath.dirname(probe), recursive=True) filesystem.open_output_stream(probe).close() - filesystem.delete_file(probe) + delete = getattr(root.connector.sink_backend, 'delete_file', None) + if delete is None: + filesystem.delete_file(probe) + else: + delete(root.connector, probe) except Exception as error: raise PermissionError( f'kumo.MLOps cannot write under artifact_root ' diff --git a/packages/kumo-s3/src/kumo_s3/sink.py b/packages/kumo-s3/src/kumo_s3/sink.py index 5e2a5a2d8..cc650cb4d 100644 --- a/packages/kumo-s3/src/kumo_s3/sink.py +++ b/packages/kumo-s3/src/kumo_s3/sink.py @@ -2,8 +2,11 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from typing import Any + from kumo_core.object_store_sink import ObjectStoreSinkBackend from kumoapi.data_source import DataSourceType +from pyarrow.fs import S3FileSystem from .backend import S3Sessions @@ -16,3 +19,16 @@ class S3SinkBackend(S3Sessions, ObjectStoreSinkBackend): def __repr__(self) -> str: return f'{type(self).__name__}(identity={self._identity!r})' + + def delete_file(self, connector: Any, path: str) -> None: + r"""Deletes the object at ``path`` (``/``). Unlike + pyarrow's delete, which recreates an emptied parent as a zero-byte + ``/`` object, this leaves no object behind.""" + bucket, _, key = path.partition('/') + session, _, connection = self._resolved_session(connector, bucket) + if isinstance(session, S3FileSystem): + session.delete_file(path) + return + session.client( + 's3', endpoint_url=connection.endpoint_url if connection else None + ).delete_object(Bucket=bucket, Key=key) diff --git a/packages/kumo-s3/tests/test_sink_delete.py b/packages/kumo-s3/tests/test_sink_delete.py new file mode 100644 index 000000000..b3f97e114 --- /dev/null +++ b/packages/kumo-s3/tests/test_sink_delete.py @@ -0,0 +1,22 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +import boto3 +from kumo_s3.sink import S3SinkBackend +from moto import mock_s3 + + +@mock_s3 +def test_deleting_the_last_object_under_a_prefix_leaves_no_marker( + monkeypatch): + client = boto3.client('s3', region_name='us-east-1') + client.create_bucket(Bucket='bucket') + client.put_object(Bucket='bucket', Key='root/.probe', Body=b'probe') + sink = S3SinkBackend() + monkeypatch.setattr( + sink, '_resolved_session', lambda connector, bucket: + (boto3.Session(region_name='us-east-1'), True, None)) + + sink.delete_file(None, 'bucket/root/.probe') + + assert client.list_objects_v2(Bucket='bucket')['KeyCount'] == 0 From 13a4e9b2b66a72910ac2127d64c8bcb1d8aad162 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 03:53:39 +0000 Subject: [PATCH 027/278] Stay silent about inferred encoders for columns the model plan overrides Fit logged "Text column 'users.BIO' will not be encoded by default" for a column whose encoder the plan set to TextEmbedding, because encoder inference runs for every column before the overrides apply. Fit now drops inference warnings for overridden columns and keeps them for the rest. Co-Authored-By: Claude Opus 5.5 --- .../kumo-engine/kumo/engine/_embeddings.py | 15 ++++++++--- .../kumo/engine/trainer/trainer.py | 26 ++++++++++--------- .../kumo-engine/test/test_embedding_models.py | 6 +++-- 3 files changed, 29 insertions(+), 18 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_embeddings.py b/packages/kumo-engine/kumo/engine/_embeddings.py index 1a51a42c8..ffeb42ff7 100644 --- a/packages/kumo-engine/kumo/engine/_embeddings.py +++ b/packages/kumo-engine/kumo/engine/_embeddings.py @@ -5,18 +5,25 @@ import logging from contextlib import contextmanager from pathlib import Path -from typing import Any, Dict, Iterator, List, Optional +from typing import Any, Collection, Dict, Iterator, List, Optional from kumo_core import embedding_models from kumo_core.embedding_models import EmbeddingModelStore @contextmanager -def no_encoder_inference_warnings() -> Iterator[None]: +def no_encoder_inference_warnings( + overridden: Optional[Collection[str]] = None) -> Iterator[None]: # Only fit reports what inference chose: training tables do not use it, - # and a trained model encodes columns as it was trained to. + # and a trained model encodes columns as it was trained to. Given the + # columns whose encoders the model plan overrides, fit stays silent on + # those alone. def drop(record: logging.LogRecord) -> bool: - return record.levelno < logging.WARNING + if overridden is None: + return record.levelno < logging.WARNING + return (record.levelno < logging.WARNING + or not isinstance(record.args, tuple) or not record.args + or record.args[0] not in overridden) inference = logging.getLogger('kumo_core.encoder_selection') inference.addFilter(drop) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index d3a818880..7d32b71b4 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -863,18 +863,20 @@ def _fit( train_table.generation_plan.lag_timesteps) year_over_year = inferred_or_missing_type_to_none( train_table.generation_plan.year_over_year) - task_dataset = build_generated_training_task_dataset( - query=workflow.query, - graph=workflow.graph, - graph_index_path=workflow.prepared_graph.index_path, - run_id=run_id, - artifact_dir=str(run_dir), - tt_id=ingest_tt_id, - required_stages=train_table.required_stages, - lag_timesteps=(None - if lag_timesteps is None else int(lag_timesteps)), - year_over_year=bool(year_over_year), - ) + with no_encoder_inference_warnings( + effective_plan.column_processing.encoder_overrides or {}): + task_dataset = build_generated_training_task_dataset( + query=workflow.query, + graph=workflow.graph, + graph_index_path=workflow.prepared_graph.index_path, + run_id=run_id, + artifact_dir=str(run_dir), + tt_id=ingest_tt_id, + required_stages=train_table.required_stages, + lag_timesteps=(None if lag_timesteps is None else + int(lag_timesteps)), + year_over_year=bool(year_over_year), + ) require_embedding_models( self._embedding_models, task_dataset.feature_schema, effective_plan.column_processing.encoder_overrides) diff --git a/packages/kumo-engine/test/test_embedding_models.py b/packages/kumo-engine/test/test_embedding_models.py index a51475c8a..681fef6c9 100644 --- a/packages/kumo-engine/test/test_embedding_models.py +++ b/packages/kumo-engine/test/test_embedding_models.py @@ -14,6 +14,7 @@ from kumo_core import embedding_models from kumo_core.embedding_models import EmbeddingModelNotFoundError, Release from kumo_core.glove import TokenEmbeddingTable +from kumoapi.encoder import Null from kumoapi.typing import Stype from kumo.engine import ( @@ -70,7 +71,8 @@ def fitted(tmp_path_factory): warnings['failed fit'] = warned.take() plan.column_processing.encoder_overrides = { - 'users.bio': TextEmbedding('words') + 'users.bio': TextEmbedding('words'), + 'users.code': Null(), } store = EmbeddingModelStore() store.add('words', GloVeVectors(archive, cache_dir=root / 'emb')) @@ -155,7 +157,7 @@ def test_each_fit_reports_what_inference_leaves_unencoded(fitted): assert fitted[-1] == { 'tables': [], 'failed fit': [unencoded], - 'fit': [unencoded], + 'fit': [], } From 6fb6b82f6b7ab10755d45f14c962fb92b59f2384 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 21:35:37 +0000 Subject: [PATCH 028/278] List the artifacts a manifest names, with their locations A manifest records where it is stored (location). artifacts() walks it, its parents, their source tables, dependencies and outputs, and lists each artifact's kind, id, location, tag and whether it is a dataset, declared once per kind. A named model is listed but not opened. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_manifest.py | 71 +++++++++++++++++++ packages/kumo-engine/test/test_artifacts.py | 62 ++++++++++++++++ 2 files changed, 133 insertions(+) create mode 100644 packages/kumo-engine/test/test_artifacts.py diff --git a/packages/kumo-engine/kumo/engine/_manifest.py b/packages/kumo-engine/kumo/engine/_manifest.py index a6e8a871d..64d73bddc 100644 --- a/packages/kumo-engine/kumo/engine/_manifest.py +++ b/packages/kumo-engine/kumo/engine/_manifest.py @@ -12,7 +12,9 @@ from pathlib import Path from typing import ( Any, + ClassVar, Dict, + Iterator, List, Mapping, Optional, @@ -39,6 +41,23 @@ def omit_none() -> Any: return field(default=None, metadata={_OMIT_NONE: True}) +@dataclass(frozen=True) +class Artifact: + r"""An artifact a manifest names, stored at ``location``, which tracked + runs tag as ``tag``. A dataset is data a run reads, or writes when its + ``context`` is ``output``.""" + + kind: str + id: str + location: Optional[str] + is_dataset: bool + tag: Optional[str] = None + connector: Optional[str] = None + version: Optional[str] = None + version_kind: Optional[str] = None + context: Optional[str] = None + + @dataclass(frozen=True, kw_only=True) class Record: """A JSON object; keys this version does not know are kept in ``extra``.""" @@ -46,6 +65,11 @@ class Record: extra: Mapping[str, Any] = field(default_factory=dict, repr=False, compare=False) + def as_artifact(self, key: str) -> Optional[Artifact]: + r"""The artifact this record names, if any, as ``key`` in its + parent.""" + return None + def to_dict(self) -> Dict[str, Any]: data = dict(self.extra) for item in fields(self): @@ -125,6 +149,11 @@ class SourceTableEntry(Record): dialect: Optional[str] = omit_none() namespace: Optional[List[Any]] = omit_none() + def as_artifact(self, key: str) -> Artifact: + return Artifact('source_table', key, self.location, True, + connector=self.connector, version=self.version, + version_kind=self.version_kind) + #: The artifact kinds, as manifests and the artifact root name them. KINDS = ('snapshot', 'training_table', 'model') @@ -141,9 +170,48 @@ class Manifest(Record): created_at: Optional[str] = None environment: Optional[Environment] = None reconstructed: Optional[bool] = omit_none() + location: Optional[str] = omit_none() #: Fields left out where this manifest is embedded in a child's. EMBED_OMITS = () + #: Whether this kind is data a run reads or writes. + IS_DATASET: ClassVar[bool] = False + #: The tag of its location in tracked runs. + URI_TAG: ClassVar[Optional[str]] = None + + def as_artifact(self, key: str) -> Artifact: + return Artifact(self.kind, self.id, self.location, self.IS_DATASET, + self.URI_TAG or f'kumo.{self.kind}_uri', + version=self.id, version_kind='exact') + + +def artifacts(manifest: Manifest) -> List[Artifact]: + r"""``manifest``'s artifact, then each one it names (parents, their + source tables, dependencies, outputs), depth first, each once. A named + manifest that is not a dataset, such as a warm-start model, is listed + but not opened: its own manifest names what it was made from.""" + found: Dict[Any, Artifact] = {} + for artifact in _walk(manifest, manifest.kind, manifest): + found.setdefault((artifact.kind, artifact.id), artifact) + return list(found.values()) + + +def _walk(value: Any, key: str, root: Manifest) -> Iterator[Artifact]: + if isinstance(value, Record): + artifact = value.as_artifact(key) + if artifact is not None: + yield artifact + if isinstance(value, Manifest) and not value.IS_DATASET and \ + value is not root: + return + for item in fields(value): + yield from _walk(getattr(value, item.name), item.name, root) + elif isinstance(value, AbcMapping): + for name, item in value.items(): + yield from _walk(item, str(name), root) + elif isinstance(value, (list, tuple)): + for item in value: + yield from _walk(item, key, root) def stamp() -> Dict[str, Any]: @@ -206,6 +274,7 @@ class SnapshotManifest(Manifest): location and version, and the engine build that wrote it.""" kind: str = 'snapshot' + IS_DATASET: ClassVar[bool] = True parents: Optional[SnapshotParents] = None record_sources: Optional[bool] = omit_none() @@ -222,6 +291,7 @@ class TrainingTableManifest(Manifest): and, for caller-supplied rows, which source table.""" kind: str = 'training_table' + IS_DATASET: ClassVar[bool] = True parents: Optional[TrainingTableParents] = None query: Optional[str] = None generation_plan: Optional[Dict[str, Any]] = None @@ -249,6 +319,7 @@ class ModelManifest(Manifest): mode. Its ``id`` is the bundle id.""" kind: str = 'model' + URI_TAG: ClassVar[Optional[str]] = 'kumo.bundle_uri' parents: Optional[ModelParents] = None fit: Optional[FitTiming] = None fit_mode: Optional[str] = None diff --git a/packages/kumo-engine/test/test_artifacts.py b/packages/kumo-engine/test/test_artifacts.py new file mode 100644 index 000000000..0d51cc262 --- /dev/null +++ b/packages/kumo-engine/test/test_artifacts.py @@ -0,0 +1,62 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from dataclasses import dataclass +from typing import List, Optional + +from kumo.engine._manifest import ( + Artifact, + Manifest, + ModelManifest, + ModelParents, + Record, + SnapshotManifest, + SnapshotParents, + SourceTableEntry, + TrainingTableManifest, + artifacts, +) + + +@dataclass(frozen=True, kw_only=True) +class ReportParents(Record): + snapshot: SnapshotManifest + model: ModelManifest + + +@dataclass(frozen=True, kw_only=True) +class ReportManifest(Manifest): + kind: str = 'report' + parents: ReportParents + pages: Optional[List[SnapshotManifest]] = None + + IS_DATASET = True + + +def test_a_new_kind_names_its_artifacts_without_other_changes(): + orders = SourceTableEntry(connector='S3', location='s3://data/orders', + version='7', version_kind='exact') + snapshot = SnapshotManifest( + id='g@d', location='s3://root/snapshots/g/d', + parents=SnapshotParents(source_tables={'orders': orders})) + model = ModelManifest( + id='m', location='s3://root/models/m', parents=ModelParents( + training_table=TrainingTableManifest(id='t'))) + report = ReportManifest( + id='r', location='s3://root/reports/r', + parents=ReportParents(snapshot=snapshot, model=model), + pages=[snapshot, SnapshotManifest(id='p', reconstructed=True)]) + + assert artifacts(report) == [ + Artifact('report', 'r', 's3://root/reports/r', True, + 'kumo.report_uri', version='r', version_kind='exact'), + Artifact('snapshot', 'g@d', 's3://root/snapshots/g/d', True, + 'kumo.snapshot_uri', version='g@d', version_kind='exact'), + Artifact('source_table', 'orders', 's3://data/orders', True, + connector='S3', version='7', version_kind='exact'), + Artifact('model', 'm', 's3://root/models/m', False, + 'kumo.bundle_uri', version='m', version_kind='exact'), + Artifact('snapshot', 'p', None, True, 'kumo.snapshot_uri', + version='p', version_kind='exact'), + ] From b561a4914b31d8ef29aca448a40098c80c637cae Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 07:37:27 +0000 Subject: [PATCH 029/278] Score binary models without a threshold: probabilities, labels only with one A model fitted without scoring defaults (the notebook flow) failed every default scoring path: predict, the pyfunc and the score job raised "binary_classification_threshold must be specified". Binary predict now always writes True_PROB and sets TARGET_PRED only from a threshold the call passes or the bundle stores; without one TARGET_PRED is null, so no threshold is invented and the output schema stays the same. Verified: test_scoring_defaults, test_mlflow_model and test_semantic_carry_forward, kumo-core test_prediction. Co-Authored-By: Claude Opus 5.5 --- .../kumo-core/src/kumo_core/prediction.py | 1 + .../train/diskgraph_nested_trainer.py | 5 ++++ packages/kumo-engine/docs/mlops.md | 4 ++-- .../kumo/engine/integrations/_mlflow_model.py | 16 +++---------- .../kumo/engine/trainer/trainer.py | 13 ++++++---- .../kumo-engine/test/test_filestore_sink.py | 3 +-- .../kumo-engine/test/test_mlflow_model.py | 24 +++++++------------ .../kumo-engine/test/test_scoring_defaults.py | 4 ++-- .../test/test_semantic_carry_forward.py | 1 - 9 files changed, 30 insertions(+), 41 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/prediction.py b/packages/kumo-core/src/kumo_core/prediction.py index effb6b69e..14f18de0d 100644 --- a/packages/kumo-core/src/kumo_core/prediction.py +++ b/packages/kumo-core/src/kumo_core/prediction.py @@ -583,6 +583,7 @@ def _score_prediction_shard( runtime_num_workers=config.worker_count, runtime_worker_id=worker_id, task_manifest=manifest, + binary_labels=config.binary_threshold is not None, ) return _publish_scored_shard( collector, output_dir=config.output_dir, run_id=config.run_id, diff --git a/packages/kumo-core/src/kumo_core/train/diskgraph_nested_trainer.py b/packages/kumo-core/src/kumo_core/train/diskgraph_nested_trainer.py index ce07cbd03..4786678e8 100644 --- a/packages/kumo-core/src/kumo_core/train/diskgraph_nested_trainer.py +++ b/packages/kumo-core/src/kumo_core/train/diskgraph_nested_trainer.py @@ -1366,6 +1366,7 @@ def predict( artifact_stats_collector: PredictionArtifactStatsCollector | None = None, task_manifest: DiskgraphTaskManifest | None = None, + binary_labels: bool = True, ) -> Any: """TT-NG batch prediction over the wheel-backed predict stage. @@ -1542,6 +1543,10 @@ def _build_output(prediction: dict[str, Any]) -> PredOutput: out_timestamps=out_timestamps, array_targets=allow_array_targets(), ) + if (not binary_labels and self.task.task_type + == TaskType.BINARY_CLASSIFICATION): + df['TARGET_PRED'] = pd.Series(pd.NA, index=df.index, + dtype='boolean') embedding = None if return_embeddings: diff --git a/packages/kumo-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index 7db3f3f77..1a7ab98c8 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -142,7 +142,7 @@ table = kumo.TrainingTable.load_from_connector(cfg, path=f"training_tables/{tabl - `result.registered_model` is the version `MLOps(registry=)` registered; `result.published_model` the one `fit(export_config=...)` published. - `result.tracking` is the tracker's run: `TrackingRef(platform="databricks", run_id="...", url="https://.../ml/experiments/.../runs/...")`. The model's `_provenance.json` records the same. -- `fit(scoring_defaults=kumo.ScoringDefaults(...))` stores prediction defaults in the bundle before it is stored and registered: `binary_classification_threshold` (binary and multilabel classification) and `num_classes_to_return` (multiclass, multilabel, link prediction). `predict` and the MLflow pyfunc apply them when a call omits them; explicit arguments and a query's `TOP k` win. `result.with_scoring_defaults(...)` retunes a fitted model as a new bundle whose manifest names the original in `parents.derived_from`. +- `fit(scoring_defaults=kumo.ScoringDefaults(...))` stores prediction defaults in the bundle before it is stored and registered: `binary_classification_threshold` (binary and multilabel classification) and `num_classes_to_return` (multiclass, multilabel, link prediction). `predict` and the MLflow pyfunc apply them when a call omits them; explicit arguments and a query's `TOP k` win. A binary model always returns `True_PROB`; with no threshold passed or stored, its `TARGET_PRED` is null. `result.with_scoring_defaults(...)` retunes a fitted model as a new bundle whose manifest names the original in `parents.derived_from`. - `result.export(cfg)` (or `trainer.export(cfg)` for a loaded model) registers a model at `cfg.registry` through its tracker. On every platform, registering a bundle the registry already holds under that name returns its version instead of adding one. - If registration fails, `fit` warns with the full error and returns with `registered_model=None` (with `on_error="raise"`, it raises), and the fit run ends failed with `kumo.error`. Registering a bundle the registry already holds under the name returns that version (logged at info) and moves the alias to it. @@ -204,7 +204,7 @@ scores = model.predict(pd.DataFrame({"ENTITY": [1, 2, 3]}), params={ "prediction_time": "2024-06-01"}) # temporal queries ``` -- Input rows have an `ENTITY` column and optionally a `TIMESTAMP`; scores come back in input order, one per scored row. Binary classification needs `params={"binary_classification_threshold": ...}` unless the model carries scoring defaults. +- Input rows have an `ENTITY` column and optionally a `TIMESTAMP`; scores come back in input order, one per scored row. Binary classification returns `True_PROB`; `TARGET_PRED` needs `params={"binary_classification_threshold": ...}` or a stored threshold, and is null without one. - Each replica loads a snapshot on its first `predict` with that URI and keeps it. `log_model(trainer, snapshot=..., serving_snapshot=...)` logs a model whose replicas load that snapshot at startup and use it when no `snapshot` param is given; `serving_snapshot` takes a URI, or a `GraphSnapshot` or snapshot id stored under the block's artifact root. - Serving needs an image with kumo-engine installed: `mlflow models serve -m --env-manager local`. Its pip requirement is `kumo-engine[s3,gcs]` of the same version (plus `[text]` for text columns). diff --git a/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py b/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py index a1be11e94..75d063c23 100644 --- a/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py +++ b/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py @@ -70,7 +70,7 @@ def predict(trainer: Any, snapshot: GraphSnapshot, rows: pd.DataFrame, 'binary_classification_threshold'), num_classes_to_return=params.get('num_classes_to_return'), ) - _require_inputs(trainer, snapshot, rows, params, time) + _require_inputs(trainer, snapshot, rows, time) trainer.scratch_dir.mkdir(parents=True, exist_ok=True) with tempfile.TemporaryDirectory(dir=trainer.scratch_dir) as tmp, \ scoped_scratch_dir(Path(tmp) / 'scratch'): @@ -138,28 +138,18 @@ def _seconds(values: pd.Series) -> pd.Series: def _require_inputs(trainer: Any, snapshot: GraphSnapshot, rows: pd.DataFrame, - params: Mapping[str, Any], time: Optional[pd.Timestamp]) -> None: from kumo_core.pquery import query_is_temporal - from kumo_core.types import TaskType from kumo.engine.pquery import PredictiveQuery - untimed = len(rows) > 0 and time is None and 'TIMESTAMP' not in rows - unthresholded = ( - params.get('binary_classification_threshold') is None and - trainer.result.scoring_defaults.binary_classification_threshold is None) - if not (untimed or unthresholded): + if len(rows) == 0 or time is not None or 'TIMESTAMP' in rows: return parsed = PredictiveQuery._from_snapshot(snapshot, trainer._training_query)._parsed - if untimed and query_is_temporal(parsed, snapshot._native_graph): + if query_is_temporal(parsed, snapshot._native_graph): raise ValueError("temporal query: pass params={'prediction_time': " "...} or a TIMESTAMP column") - if unthresholded and (TaskType(parsed.task_type) - is TaskType.BINARY_CLASSIFICATION): - raise ValueError("binary classification: pass " - "params={'binary_classification_threshold': ...}") def snapshot_location(uri: str) -> Tuple[Connector, str]: diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 7d32b71b4..6e33eee1e 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -1089,9 +1089,10 @@ def predict( Link predictions require an explicit count of at most 500. Embeddings alone do not require a count. - Binary classification requires an explicit threshold. Multilabel - classification defaults to 0.5. Both accept thresholds from 0 to 1; - thresholds change predicted labels without changing class scores. + Binary classification always returns ``True_PROB``; ``TARGET_PRED`` + is set only with a threshold, else null. Multilabel classification + defaults to 0.5. Both accept thresholds from 0 to 1; thresholds + change predicted labels without changing class scores. Omitted arguments fall back to the bundle's :attr:`scoring_defaults`. """ @@ -1183,8 +1184,10 @@ def _predict( 'Prediction query target/entity contract differs from the trained model' ) task_type = TaskType(contract[1]) - threshold = resolve_classification_threshold( - task_type, binary_classification_threshold, binary_required=True) + threshold = (None if binary_classification_threshold is None + and task_type is TaskType.BINARY_CLASSIFICATION else + resolve_classification_threshold( + task_type, binary_classification_threshold)) destinations = (output_config.destinations() if output_config is not None else {}) artifacts = frozenset( diff --git a/packages/kumo-engine/test/test_filestore_sink.py b/packages/kumo-engine/test/test_filestore_sink.py index 770b1ac2d..2cafd1cfb 100644 --- a/packages/kumo-engine/test/test_filestore_sink.py +++ b/packages/kumo-engine/test/test_filestore_sink.py @@ -158,8 +158,7 @@ def test_predictions_export_through_a_filesystem_only_sink(fitted, connector): trainer, table, snapshot = fitted output = PredictionOutputConfig( connector, output_paths={'predictions': '@stage/preds/'}) - result = trainer.predict(table, snapshot=snapshot, output_config=output, - binary_classification_threshold=0.5) + result = trainer.predict(table, snapshot=snapshot, output_config=output) store = connector.sink_backend.store assert any(name.startswith('@stage/preds/') for name in store.objects) result.export( diff --git a/packages/kumo-engine/test/test_mlflow_model.py b/packages/kumo-engine/test/test_mlflow_model.py index a62f90694..cb93defd8 100644 --- a/packages/kumo-engine/test/test_mlflow_model.py +++ b/packages/kumo-engine/test/test_mlflow_model.py @@ -181,32 +181,24 @@ def predict_beside_another_call(self, *args, **kwargs): assert set(tables.glob('*')) == before | {tables / 'other-call'} -def test_pyfunc_names_missing_temporal_and_threshold_inputs( - fitted, monkeypatch): - from dataclasses import replace - +def test_pyfunc_names_a_missing_prediction_time_and_needs_no_threshold( + fitted): from kumo.engine.integrations._mlflow_model import predict trainer, table, snapshot = fitted rows = pd.DataFrame({'ENTITY': range(5)}) - time = str(pd.Timestamp(table.timestamps[0])) with pytest.raises(ValueError) as error: predict(trainer, snapshot, rows, {'binary_classification_threshold': .5}) assert str(error.value) == ("temporal query: pass params={" "'prediction_time': ...} or a TIMESTAMP " "column") - with pytest.raises(ValueError, match=r"^binary classification: pass " - r"params=\{'binary_classification_threshold': " - r"\.\.\.\}$"): - predict(trainer, snapshot, rows, {'prediction_time': time}) - - monkeypatch.setattr( - trainer, '_training_result', - replace(trainer.result, scoring_defaults=kumo.ScoringDefaults( - binary_classification_threshold=.5))) - assert len(predict(trainer, snapshot, rows, {'prediction_time': time - })) == 5 + + time = str(pd.Timestamp(table.timestamps[0])) + scores = predict(trainer, snapshot, rows, {'prediction_time': time}) + assert len(scores) == 5 + assert scores['TARGET_PRED'].isna().all() + assert scores['True_PROB'].between(0, 1).all() def test_pyfunc_returns_rows_in_input_order(fitted): diff --git a/packages/kumo-engine/test/test_scoring_defaults.py b/packages/kumo-engine/test/test_scoring_defaults.py index fd8812113..3ba5d2537 100644 --- a/packages/kumo-engine/test/test_scoring_defaults.py +++ b/packages/kumo-engine/test/test_scoring_defaults.py @@ -135,8 +135,8 @@ def test_unlisted_or_unreadable_defaults_are_ignored(fitted, tmp_path): with pytest.warns(UserWarning, match='Ignoring the scoring defaults'): ignored = kumo.Trainer.load(bundle) assert ignored.result.scoring_defaults == kumo.ScoringDefaults() - with pytest.raises(ValueError, match='must be specified'): - ignored.predict(table, snapshot=snapshot) + run = ignored.predict(table, snapshot=snapshot) + assert run.frame['TARGET_PRED'].isna().all() def test_a_fit_stores_its_defaults_in_the_bundle(fitted, tmp_path): diff --git a/packages/kumo-engine/test/test_semantic_carry_forward.py b/packages/kumo-engine/test/test_semantic_carry_forward.py index 5968ef805..81da3228e 100644 --- a/packages/kumo-engine/test/test_semantic_carry_forward.py +++ b/packages/kumo-engine/test/test_semantic_carry_forward.py @@ -129,7 +129,6 @@ def test_compile_rejects_invalid_prediction_time(snapshot, tmp_path, value, @pytest.mark.parametrize('query, threshold, message', [ ('PREDICT users.age FOR EACH users.id', .5, 'target/entity'), - ('PREDICT users.label FOR EACH users.id', None, 'must be specified'), ('PREDICT users.label FOR EACH users.id', 1.1, 'between 0 and 1'), ]) def test_prediction_admission_before_materialization(snapshot, tmp_path, query, From a44d34d49ff7be1a0fc8e8e95e840fe4808db3e0 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 09:59:40 +0000 Subject: [PATCH 030/278] Create the tracker's experiment when kumo.MLOps is built, or say why not An experiment under a Databricks folder that did not exist made every tracked fit fail to record, only with a warning under on_error='warn': result.tracking was None and nothing pointed at the folder. MlflowTracking.check, which runs when kumo.MLOps is built, now creates a missing experiment. If MLflow refuses, it raises with MLflow's error and says the parent folder must exist and how to create it. Verified: test_mlflow_tracking, test_mlflow_target. Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/integrations/mlflow.py | 20 ++++++++++++++++++- .../kumo-engine/test/test_mlflow_model.py | 8 +++++--- .../kumo-engine/test/test_mlflow_tracking.py | 19 ++++++++++++++++++ 3 files changed, 43 insertions(+), 4 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index 96c616790..774eb1bfb 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -188,11 +188,13 @@ def check(self, registry: Optional[ModelRegistryLocation] = None) -> None: uri = self.target.tracking_uri owned = frozenset(name for name in _HTTP_LIMITS if name not in os.environ) + missing = False try: with _limited_http(owned), self.target.scope(): client = self.target.client() if self.experiment_name is not None: - client.get_experiment_by_name(self.experiment_name) + missing = client.get_experiment_by_name( + self.experiment_name) is None else: client.search_experiments(max_results=1) if registry is not None: @@ -204,6 +206,22 @@ def check(self, registry: Optional[ModelRegistryLocation] = None) -> None: except Exception as error: raise RuntimeError(f'MLflow at {uri} is not usable with these ' f'credentials: {full_error(error)}') from error + if missing: + self._create_experiment(client, owned) + + def _create_experiment(self, client: Any, owned: FrozenSet[str]) -> None: + with _limited_http(owned), self.target.scope(): + try: + client.create_experiment(self.experiment_name) + except Exception as error: + if client.get_experiment_by_name(self.experiment_name): + return + raise RuntimeError( + f'Could not create the MLflow experiment ' + f'{self.experiment_name!r}: {full_error(error)}. Its ' + f'parent folder must exist: on Databricks, create it ' + f'with `databricks workspace mkdirs ` or in the ' + f'workspace UI.') from error def prepare(self) -> None: active = self._mlflow.active_run() diff --git a/packages/kumo-engine/test/test_mlflow_model.py b/packages/kumo-engine/test/test_mlflow_model.py index cb93defd8..51eb51fb7 100644 --- a/packages/kumo-engine/test/test_mlflow_model.py +++ b/packages/kumo-engine/test/test_mlflow_model.py @@ -490,12 +490,15 @@ def no_folder(self, name, *args, **kwargs): raise MlflowException(f'RESOURCE_DOES_NOT_EXIST: Parent of {name} ' f'does not exist.') - monkeypatch.setattr(MlflowClient, 'create_experiment', no_folder) table, snapshot = table_and_snapshot cfg = kumo.MLOps(str(tmp_path / 'root'), tracking=MlflowTracking('/NoSuchFolder/churn', log_system_metrics=False), on_error='raise') + warn = dataclasses.replace(cfg, on_error='warn') + monkeypatch.setattr(MlflowClient, 'get_experiment_by_name', + lambda self, name: None) + monkeypatch.setattr(MlflowClient, 'create_experiment', no_folder) trainer = kumo.Trainer(_plan()) with pytest.raises(RuntimeError, match='RESOURCE_DOES_NOT_EXIST') as raised: trainer.fit(table, snapshot=snapshot, mlops=cfg) @@ -507,10 +510,9 @@ def no_folder(self, name, *args, **kwargs): assert 'MlflowException: RESOURCE_DOES_NOT_EXIST' in off.getMessage() assert not off.exc_info - cfg = dataclasses.replace(cfg, on_error='warn') with warnings.catch_warnings(record=True) as caught: warnings.simplefilter('always') - trainer.fit(table, snapshot=snapshot, mlops=cfg) + trainer.fit(table, snapshot=snapshot, mlops=warn) reported = [str(w.message) for w in caught if 'RESOURCE_DOES_NOT_EXIST' in str(w.message)] assert reported == [ diff --git a/packages/kumo-engine/test/test_mlflow_tracking.py b/packages/kumo-engine/test/test_mlflow_tracking.py index 920e92272..709b9885f 100644 --- a/packages/kumo-engine/test/test_mlflow_tracking.py +++ b/packages/kumo-engine/test/test_mlflow_tracking.py @@ -241,6 +241,25 @@ def test_registry_model_names_are_checked_when_mlops_is_built( registry=kumo.ModelRegistryLocation(name)) +def test_building_mlops_creates_the_experiment_or_says_why_it_cannot( + tmp_path, monkeypatch): + from kumo.engine.integrations.mlflow import MlflowTracking + + uri = f'sqlite:///{tmp_path}/m.db' + kumo.MLOps(str(tmp_path), tracking=MlflowTracking( + 'churn', log_system_metrics=False, target=uri)) + assert mlflow.MlflowClient(uri).get_experiment_by_name('churn') + + def refuse(self, name, *args, **kwargs): + raise mlflow.exceptions.MlflowException('folder does not exist') + + monkeypatch.setattr(mlflow.MlflowClient, 'create_experiment', refuse) + with pytest.raises(RuntimeError, match=r"experiment '/Shared/x/churn'.*" + r"folder does not exist.*parent folder must exist"): + kumo.MLOps(str(tmp_path), tracking=MlflowTracking( + '/Shared/x/churn', log_system_metrics=False, target=uri)) + + @pytest.mark.parametrize('uri,command', [ (SAGEMAKER, 'aws sagemaker create-presigned-mlflow-tracking-server-url ' '--tracking-server-name demo --region us-west-2'), From c30bcfe0dc94d31e086989930f5b64317f08b4c0 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 12:36:41 +0000 Subject: [PATCH 031/278] Keep column types when a table export stamps its id columns prepared_artifact_file added PREDICTION_ID and BUNDLE_ID through a pandas round trip, which turns an all-null bool column into Arrow's null type. A binary run scored without a threshold (TARGET_PRED null) therefore loaded into Snowflake with TARGET_PRED as NUMBER, while a thresholded run or a baseline loaded it as BOOLEAN, and CREATE MODEL MONITOR refused the pair. The id columns are now added in Arrow, so every source column keeps its type. Verified: kumo-core test_artifact_export (new: an all-null bool and a UTC timestamp keep their types), kumo-snowflake and kumo-databricks sink tests. Co-Authored-By: Claude Opus 5.5 --- .../kumo-core/src/kumo_core/artifact_export.py | 12 +++++++++--- packages/kumo-core/tests/test_artifact_export.py | 15 +++++++++++++++ 2 files changed, 24 insertions(+), 3 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/artifact_export.py b/packages/kumo-core/src/kumo_core/artifact_export.py index 311b972bc..e438e3e79 100644 --- a/packages/kumo-core/src/kumo_core/artifact_export.py +++ b/packages/kumo-core/src/kumo_core/artifact_export.py @@ -304,10 +304,16 @@ 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() + # In Arrow: a pandas round trip turns an all-null bool column + # into Arrow's null type. + 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.array([value] * table.num_rows, + type=pa.scalar(value).type) + 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 {})) diff --git a/packages/kumo-core/tests/test_artifact_export.py b/packages/kumo-core/tests/test_artifact_export.py index b3701f00f..e74b28ff2 100644 --- a/packages/kumo-core/tests/test_artifact_export.py +++ b/packages/kumo-core/tests/test_artifact_export.py @@ -138,6 +138,21 @@ def test_file_metadata_uses_hosted_batch_bound_and_preserves_source(tmp_path): assert source.read_bytes() == before +def test_stamped_columns_keep_the_source_types(tmp_path): + import pyarrow as pa + import pyarrow.parquet as pq + from kumo_core.artifact_export import prepared_artifact_file + + source = tmp_path / 'scores.parquet' + schema = pa.schema([('ENTITY', pa.int64()), ('TARGET_PRED', pa.bool_()), + ('ANCHOR', pa.timestamp('ns', tz='UTC'))]) + pq.write_table(pa.table([[1, 2], [None, None], [None, 0]], schema=schema), + source) + with prepared_artifact_file(str(source), {'BUNDLE_ID': 'b'}) as staged: + assert pq.read_schema(staged) == schema.append( + pa.field('BUNDLE_ID', pa.string())) + + @pytest.mark.parametrize('existing_empty', [False, True]) def test_multipart_failure_and_failed_recopy_remain_retryable( tmp_path, monkeypatch, existing_empty): From 3c225163091cd6f56afa1e985621635f71e59b4d Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 16:37:20 +0000 Subject: [PATCH 032/278] Report snapshot progress through the progress events Graph.snapshot printed nothing kumo-side while it staged tables and built the index (75 s of silence on a 32k-row graph in a usability run). snapshot(callbacks=[...]) now takes the same callbacks as fit and reports snapshot_started, table_staged per table and snapshot_ended (snapshot id, reused), each with the seconds since the start and logged at INFO. mlops.md documents them, and its TrainingTable.load_from_connector examples say pquery= after #50's rename. Verified: test_fit_events (new: a snapshot's events, then a reused one). Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/docs/mlops.md | 6 +-- packages/kumo-engine/kumo/engine/_events.py | 7 +-- packages/kumo-engine/kumo/engine/graph.py | 51 ++++++++++++++++++-- packages/kumo-engine/test/test_fit_events.py | 20 ++++++++ 4 files changed, 74 insertions(+), 10 deletions(-) diff --git a/packages/kumo-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index 1a7ab98c8..e209eeaed 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -131,7 +131,7 @@ Pass the `kumo.MLOps` where a connector goes, with paths relative to the artifac snapshot = kumo.GraphSnapshot.load_from_connector(cfg, id=result.snapshot_id) trainer = kumo.Trainer.load_from_connector(cfg, path=f"models/{result.id}") champion = kumo.Trainer.load_from_connector(cfg, location=kumo.ModelRegistryLocation("users-churn", alias="champion")) -table = kumo.TrainingTable.load_from_connector(cfg, path=f"training_tables/{table_id}", pq=pquery, snapshot=snapshot) +table = kumo.TrainingTable.load_from_connector(cfg, path=f"training_tables/{table_id}", pquery=pquery, snapshot=snapshot) ``` `location=` loads the version an alias (or `version=`) names from the tracker's registry, one way on every platform. A loaded model is a `Trainer`; its `result` is the `TrainingResult` (`.id`, `.manifest`, `.metrics`, `.scoring_defaults`, `.holdout_df()`). @@ -251,7 +251,7 @@ Every artifact has one way out and one way back. A `kumo.MLOps` stands in for th | Artifact | Export | Load | |---|---|---| | Snapshot | `snapshot.export(SnapshotOutputConfig(connector, path))`, returns the location | `GraphSnapshot.load_from_connector(connector, path=location)`, `(cfg, id=...)` | -| Training table | `table.export(TrainingTableOutputConfig(connector, name))` | `TrainingTable.load_from_connector(connector, path=name, pq=pq, snapshot=snapshot)` | +| Training table | `table.export(TrainingTableOutputConfig(connector, name))` | `TrainingTable.load_from_connector(connector, path=name, pquery=pquery, snapshot=snapshot)` | | Model | `result.export(ModelOutputConfig(connector, path=... or location=...))`, or `fit(export_config=...)` | `Trainer.load_from_connector(connector, path=... or location=...)` | Loading needs a file destination; a table export is for consumers outside kumo-engine. Exporting again to a destination that holds a complete copy of the same artifact reuses it (logged at info). `ModelOutputConfig(connector, location=...)` publishes to a connector's registry outside a block. @@ -295,7 +295,7 @@ Other systems plug in through `kumo.engine.tracking.Tracking`, which receives re ### Progress events -`fit(callbacks=[...])` calls each callback with every `kumo.engine.tracking.Event`: `fit_started`, per trial `trial_started`, `epoch_ended` and `trial_ended`, then `fit_ended`. Ended events carry `params["status"]` and, on failure, `params["error"]`; a callback that raises is logged and not called again. For JSON lines: `callbacks=[lambda e: print(json.dumps(e.to_dict()))]`. +`fit(callbacks=[...])` calls each callback with every `kumo.engine.tracking.Event`: `fit_started`, per trial `trial_started`, `epoch_ended` and `trial_ended`, then `fit_ended`. Ended events carry `params["status"]` and, on failure, `params["error"]`; a callback that raises is logged and not called again. For JSON lines: `callbacks=[lambda e: print(json.dumps(e.to_dict()))]`. `graph.snapshot(callbacks=[...])` reports `snapshot_started`, `table_staged` per table and `snapshot_ended` (with `snapshot_id` and `reused`), each with `params["seconds"]` since the start, and logs each at INFO. ### Limits diff --git a/packages/kumo-engine/kumo/engine/_events.py b/packages/kumo-engine/kumo/engine/_events.py index 98e8c2394..ecf1c2a79 100644 --- a/packages/kumo-engine/kumo/engine/_events.py +++ b/packages/kumo-engine/kumo/engine/_events.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -r"""Progress events reported by `Trainer.fit`.""" +r"""Progress events reported by `Trainer.fit` and `Graph.snapshot`.""" import hashlib import logging @@ -24,7 +24,8 @@ logger = logging.getLogger(__name__) EventName = Literal['fit_started', 'trial_started', 'epoch_ended', - 'trial_ended', 'fit_ended'] + 'trial_ended', 'fit_ended', 'snapshot_started', + 'table_staged', 'snapshot_ended'] @dataclass(frozen=True) @@ -75,7 +76,7 @@ def emit(self, name: EventName, *, trial: Optional[int] = None, except Exception: logger.warning( 'Progress callback %r failed; it is disabled for the ' - 'rest of this fit', callback, exc_info=True) + 'rest of this call', callback, exc_info=True) self._callbacks.remove(callback) def event(self, name: str, **fields: Any) -> None: diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index 982ee731b..528d1c4a9 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -60,17 +60,29 @@ """ import json +import logging import threading +import time import warnings from dataclasses import dataclass, field, replace from functools import cached_property, partial from pathlib import Path -from typing import TYPE_CHECKING, Any, Dict, List, Mapping, Optional, Sequence +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + List, + Mapping, + Optional, + Sequence, +) from kumo_core.artifacts import PROVENANCE_FILE from kumo_core.pquery import NativeGraph, normalize_graph from kumoapi.graph import Edge +from kumo.engine._events import Event, EventEmitter from kumo.engine._manifest import ( RECONSTRUCTED, SnapshotManifest, @@ -95,6 +107,7 @@ SNAPSHOT_MANIFEST_FILE = PROVENANCE_FILE +logger = logging.getLogger(__name__) _warned_scratch: set[Path] = set() _warned_backfill: set[Path] = set() @@ -268,6 +281,21 @@ def report_other_graph(message: str, before: Optional[Mapping[str, Any]], stacklevel=stacklevel + 1) +class _SnapshotEvents(EventEmitter): + r"""Snapshot progress: each event goes to the callbacks and the log, + with the seconds since the snapshot started.""" + + def __init__(self, callbacks: Sequence[Callable[[Event], None]]): + super().__init__(callbacks) + self._started = time.monotonic() + + def emit(self, name: Any, **fields: Any) -> None: + params = {**fields.pop('params', {}), + 'seconds': round(time.monotonic() - self._started, 1)} + logger.info('Graph.snapshot: %s %s', name, params) + super().emit(name, params=params, **fields) + + def _store() -> Any: from kumo_core.diskgraph.snapshot import ( DiskgraphLocalConfig, @@ -671,7 +699,9 @@ def validate(self) -> None: self.native_graph def snapshot(self, *, record_sources: bool = True, - mlops: Optional[Any] = None) -> GraphSnapshot: + mlops: Optional[Any] = None, + callbacks: Sequence[Callable[[Event], None]] = () + ) -> GraphSnapshot: r"""Materialize this graph: read each table's connector, run the result through `kumo-core` to build column stats and the on-disk index. The index is stored under a directory keyed by both `self.id` @@ -688,6 +718,9 @@ def snapshot(self, *, record_sources: bool = True, record_sources: if False, the snapshot manifest records each source as connector and table name only (no bucket, path or account) and Custom SQL as a hash only (no SQL text). + callbacks: called with each progress event, + `snapshot_started`, `table_staged` (per table) and + `snapshot_ended`; each is also logged at INFO. """ self.validate() for table in self._named_tables.values(): @@ -716,9 +749,12 @@ def snapshot(self, *, record_sources: bool = True, from kumo.engine._mlops import resolve settings = resolve(mlops) + events = _SnapshotEvents(callbacks) + events.emit('snapshot_started', params={ + 'graph': graph_key, 'tables': list(self._named_tables)}) with _lock(graph_key): snapshot = self._snapshot_no_lock(graph_key, persist_dir, - record_sources) + record_sources, events) if settings is not None and settings.wants('snapshot'): from kumo.engine._mlops import raise_on_error from kumo.engine._snapshot_io import store @@ -730,7 +766,9 @@ def snapshot(self, *, record_sources: bool = True, return snapshot def _snapshot_no_lock(self, graph_key: str, persist_dir: Path, - record_sources: bool = True) -> GraphSnapshot: + record_sources: bool = True, + events: Optional['_SnapshotEvents'] = None + ) -> GraphSnapshot: r"""`snapshot()`'s body, run under its `FileLock`. Not meant to be called on its own -- `persist_dir`, unlike everything else here, isn't a pure function of `self` alone (it's `snapshot()`'s lock key, @@ -792,6 +830,8 @@ def _snapshot_no_lock(self, graph_key: str, persist_dir: Path, "limitation in kumoai/backend/local/graph.py's " "_stage_sources, which raises the same way.") table_data_versions[alias] = prepared.version + if events is not None: + events.emit('table_staged', params={'table': alias}) source_tables[alias] = source_table( connector, name, sql, prepared.version, record_sources=record_sources) @@ -848,6 +888,9 @@ def _snapshot_no_lock(self, graph_key: str, persist_dir: Path, record_sources, fresh=existing is None) _warn_ingest_data_loss(result, graph_config) + if events is not None: + events.emit('snapshot_ended', params={ + 'snapshot_id': snapshot_id, 'reused': existing is not None}) return GraphSnapshot( id=snapshot_id, index_path=Path(result.index_path), diff --git a/packages/kumo-engine/test/test_fit_events.py b/packages/kumo-engine/test/test_fit_events.py index 08a22729c..682a68a81 100644 --- a/packages/kumo-engine/test/test_fit_events.py +++ b/packages/kumo-engine/test/test_fit_events.py @@ -151,3 +151,23 @@ def test_epoch_event_without_recorded_metrics_does_not_raise(tmp_path): logger.on_train_epoch_end(None, None) assert received == [] + + +def test_a_snapshot_reports_progress_events(graph, tmp_path): + previous = get_scratch_dir() + kumo.configure(scratch_dir=tmp_path / 'scratch') + try: + events, again = [], [] + snapshot = graph.snapshot(callbacks=[events.append]) + graph.snapshot(callbacks=[again.append]) + finally: + kumo.configure(scratch_dir=previous) + + assert [event.name for event in events] == [ + 'snapshot_started', 'table_staged', 'table_staged', 'snapshot_ended'] + assert [event.params['table'] for event in events[1:3]] == [ + 'users', 'events'] + assert events[-1].params['snapshot_id'] == snapshot.id + assert events[-1].params['reused'] is False + assert again[-1].params['reused'] is True + assert all(event.params['seconds'] >= 0 for event in events) From 742d7baf7c8e7dbda8b0a56c5487e5fc74760df3 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 16:43:07 +0000 Subject: [PATCH 033/278] Keep PyG's import warning and Lightning's tips out of kumo calls MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit With torch 2.15, import kumo.engine printed PyTorch Geometric's FutureWarning about its import-time torch.jit.script calls, and every fit or predict logged Lightning's '💡 Tip' lines advertising its cloud products; users can act on neither. kumo.engine now imports torch_geometric.nn first with that one warning filtered, and fit and predict drop Lightning's Tip records while they run. Other warnings and Lightning log lines, and the user's own Lightning runs, are unchanged. Verified: test_fit_events (new: Tip lines dropped inside a kumo call only), test_optional_imports; import kumo.engine with -W error::FutureWarning on torch 2.15. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/__init__.py | 1 + packages/kumo-engine/kumo/engine/_quiet.py | 12 ++++++++ .../kumo/engine/trainer/trainer.py | 28 +++++++++++++------ packages/kumo-engine/test/test_fit_events.py | 23 +++++++++++++++ 4 files changed, 55 insertions(+), 9 deletions(-) create mode 100644 packages/kumo-engine/kumo/engine/_quiet.py diff --git a/packages/kumo-engine/kumo/engine/__init__.py b/packages/kumo-engine/kumo/engine/__init__.py index ca5a6b00f..3b94059a7 100644 --- a/packages/kumo-engine/kumo/engine/__init__.py +++ b/packages/kumo-engine/kumo/engine/__init__.py @@ -27,6 +27,7 @@ from kumoapi.plugin import ModelRegistryLocation from kumoapi.table import Column +from kumo.engine import _quiet # noqa: F401 from kumo.engine._manifest import ( ModelManifest, SnapshotManifest, diff --git a/packages/kumo-engine/kumo/engine/_quiet.py b/packages/kumo-engine/kumo/engine/_quiet.py new file mode 100644 index 000000000..3a407b1b2 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/_quiet.py @@ -0,0 +1,12 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""Imports PyTorch Geometric without the FutureWarning torch 2.15+ raises +for its import-time ``torch.jit.script`` calls, which users cannot act on.""" + +import warnings + +with warnings.catch_warnings(): + warnings.filterwarnings('ignore', r'`torch\.jit\.script` is deprecated', + FutureWarning) + import torch_geometric.nn # noqa: F401 diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 6e33eee1e..99a586f4e 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -544,7 +544,8 @@ def _stage_model_bundle(artifacts: Mapping[str, str], dest: Path) -> Path: def _without_torch_scatter_hint(method: _F) -> _F: r"""Runs ``method`` without PyTorch Geometric's GPU hint that min/max - aggregation "can be accelerated via the 'torch-scatter' package". + aggregation "can be accelerated via the 'torch-scatter' package", and + without Lightning's "Tip" log lines advertising its cloud products. kumo-engine uses PyTorch's own scatter instead of that optional extension (see ``test_wheel_uses_torch_fallbacks_without_native_extensions``), so @@ -552,18 +553,27 @@ def _without_torch_scatter_hint(method: _F) -> _F: """ @functools.wraps(method) def wrapper(*args: Any, **kwargs: Any) -> Any: - with warnings.catch_warnings(): - warnings.filterwarnings( - 'ignore', - message=r"The usage of `scatter\(reduce='\w+'\)` can be " - r"accelerated via the 'torch-scatter' package", - category=UserWarning, - ) - return method(*args, **kwargs) + tips = logging.getLogger('pytorch_lightning.utilities.rank_zero') + tips.addFilter(_no_lightning_tip) + try: + with warnings.catch_warnings(): + warnings.filterwarnings( + 'ignore', + message=r"The usage of `scatter\(reduce='\w+'\)` can be " + r"accelerated via the 'torch-scatter' package", + category=UserWarning, + ) + return method(*args, **kwargs) + finally: + tips.removeFilter(_no_lightning_tip) return cast(_F, wrapper) +def _no_lightning_tip(record: logging.LogRecord) -> bool: + return not str(record.msg).startswith('💡 Tip') + + class Trainer: r"""Trains and runs a GNN model in-process. diff --git a/packages/kumo-engine/test/test_fit_events.py b/packages/kumo-engine/test/test_fit_events.py index 682a68a81..fa34faae4 100644 --- a/packages/kumo-engine/test/test_fit_events.py +++ b/packages/kumo-engine/test/test_fit_events.py @@ -3,6 +3,7 @@ # SPDX-License-Identifier: Apache-2.0 import json +import logging import numpy as np import pandas as pd @@ -171,3 +172,25 @@ def test_a_snapshot_reports_progress_events(graph, tmp_path): assert events[-1].params['reused'] is False assert again[-1].params['reused'] is True assert all(event.params['seconds'] >= 0 for event in events) + + +def test_kumo_calls_drop_lightning_tips_and_only_those(caplog): + from pytorch_lightning.utilities.rank_zero import rank_zero_info + + from kumo.engine.trainer.trainer import _without_torch_scatter_hint + + @_without_torch_scatter_hint + def fit(): + rank_zero_info('💡 Tip: inside') + rank_zero_info('GPU available: False') + + lightning = logging.getLogger('pytorch_lightning') + lightning.addHandler(caplog.handler) # it does not propagate + try: + with caplog.at_level('INFO', logger='pytorch_lightning'): + fit() + rank_zero_info('💡 Tip: outside') + finally: + lightning.removeHandler(caplog.handler) + assert 'inside' not in caplog.text + assert 'GPU available' in caplog.text and 'outside' in caplog.text From 54159e59cf5513c4fe43060954a483478d54b7cc Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 16:45:20 +0000 Subject: [PATCH 034/278] Link Databricks runs through the workspace URL on a cluster On a cluster, result.tracking.url used the SDK's host, the regional endpoint, whose run links do not open the workspace's page. It now uses the workspace URL the notebook context names (MLflow's get_workspace_info_from_dbutils), and the SDK's host elsewhere. Verified: test_mlflow_target. Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/integrations/_mlflow_target.py | 11 +++++++++++ packages/kumo-engine/test/test_mlflow_target.py | 13 +++++++++++++ 2 files changed, 24 insertions(+) diff --git a/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py b/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py index 62ecac2aa..7f3483377 100644 --- a/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py +++ b/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py @@ -139,6 +139,17 @@ def resolve(target: Any) -> Target: def _workspace_host() -> Optional[str]: + # On a cluster the SDK's host is the regional one; the notebook context + # names the workspace's own URL, which run links need. + try: + from mlflow.utils.databricks_utils import ( + get_workspace_info_from_dbutils, + ) + host, _ = get_workspace_info_from_dbutils() + if host: + return host + except Exception: + pass try: return import_provider('databricks.sdk.config', 'databricks').Config().host diff --git a/packages/kumo-engine/test/test_mlflow_target.py b/packages/kumo-engine/test/test_mlflow_target.py index 4b9b360a8..e883cd238 100644 --- a/packages/kumo-engine/test/test_mlflow_target.py +++ b/packages/kumo-engine/test/test_mlflow_target.py @@ -114,6 +114,19 @@ def test_a_databricks_service_principal_lends_a_token_for_each_call( 'https://dbc.example/ml/experiments/7/runs/r') +def test_a_cluster_run_links_to_the_workspace_url_not_the_regional_host( + monkeypatch): + from mlflow.utils import databricks_utils + + from kumo.engine.integrations._mlflow_target import resolve + + target = resolve('databricks') + monkeypatch.setattr(databricks_utils, 'get_workspace_info_from_dbutils', + lambda: ('https://dbc-1.example', '42')) + assert target.run_url('7', 'r') == ( + 'https://dbc-1.example/ml/experiments/7/runs/r') + + def test_an_aws_access_key_signs_sagemaker_calls_and_uploads( monkeypatch, tracking_uri): boto3 = pytest.importorskip('boto3') From d7eb13330a0f749032d470701598cfc8a3606981 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 17:01:11 +0000 Subject: [PATCH 035/278] Give the prefix-scoped S3 test fake a delete and keep a URI test offline Building kumo.MLOps now writes and deletes a probe under its root. The prefix-scoped S3 fake gains delete_file (real prefix rights include s3:DeleteObject), and the snapshot_uri test, which only formats a URI, builds its MLOps with store=False so it makes no S3 call. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/test/test_mlops.py | 2 +- packages/kumo-engine/test/test_model_agnostic_mlops.py | 3 +++ 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/packages/kumo-engine/test/test_mlops.py b/packages/kumo-engine/test/test_mlops.py index 53f929086..e258f3eb8 100644 --- a/packages/kumo-engine/test/test_mlops.py +++ b/packages/kumo-engine/test/test_mlops.py @@ -76,7 +76,7 @@ def test_a_root_is_a_uri_or_a_connector_and_path(tmp_path): assert pair.wants('model') and not pair.wants('snapshot') assert not kumo.MLOps(artifact_root=str(tmp_path), store=False).kinds assert uri.snapshot_uri('ab12@cd34') == f'{tmp_path}/snapshots/ab12/cd34' - assert kumo.MLOps('s3://bucket/kumo').snapshot_uri( + assert kumo.MLOps('s3://bucket/kumo', store=False).snapshot_uri( SimpleNamespace(id='ab12@cd34')) == 's3://bucket/kumo/snapshots/ab12/cd34' diff --git a/packages/kumo-engine/test/test_model_agnostic_mlops.py b/packages/kumo-engine/test/test_model_agnostic_mlops.py index 4bfa258ad..7cd604a04 100644 --- a/packages/kumo-engine/test/test_model_agnostic_mlops.py +++ b/packages/kumo-engine/test/test_model_agnostic_mlops.py @@ -311,6 +311,9 @@ def close(self): return Writer() + def delete_file(self, path): + del self.objects[path] + def test_an_asset_is_stored_on_s3_with_prefix_scoped_rights( tmp_path, monkeypatch): From b1472c33e14e09867dea15ef1a1e26bbedfd1a95 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 18:33:29 +0000 Subject: [PATCH 036/278] Drop MLflow's request-header errors inside kumo's MLflow calls A tracked fit against Databricks from a laptop logged 'Encountered unexpected error during resolving request headers: list index out of range' about 74 times: one of MLflow's request-header providers fails off its platform, and the call proceeds without those headers. Target.scope, which wraps kumo-engine's MLflow calls, now filters that one message from mlflow.tracking.request_header.registry while it runs; other MLflow warnings, and the user's own MLflow calls, are unchanged. Verified: test_mlflow_target (new: dropped inside a kumo call only). Co-Authored-By: Claude Opus 5.5 --- .../engine/integrations/_mlflow_target.py | 38 ++++++++++++++++--- .../kumo-engine/test/test_mlflow_target.py | 21 ++++++++++ 2 files changed, 54 insertions(+), 5 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py b/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py index 7f3483377..5ce74d53e 100644 --- a/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py +++ b/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py @@ -14,6 +14,7 @@ tracking URI. """ +import logging import os import re import threading @@ -26,6 +27,32 @@ SAGEMAKER_ARN = re.compile(r'arn:aws[a-z-]*:sagemaker:([^:]*):[^:]*:' r'(mlflow-tracking-server|mlflow-app)/(.+)') _LOCK = threading.RLock() +_HEADERS = logging.getLogger('mlflow.tracking.request_header.registry') +_quiet = 0 +_QUIET_LOCK = threading.Lock() + + +def _no_header_error(record: logging.LogRecord) -> bool: + # A request-header provider that fails off its platform; the call + # proceeds without those headers. + return not record.getMessage().startswith( + 'Encountered unexpected error during resolving request headers') + + +@contextmanager +def _without_header_warnings() -> Iterator[None]: + global _quiet + with _QUIET_LOCK: + _quiet += 1 + if _quiet == 1: + _HEADERS.addFilter(_no_header_error) + try: + yield + finally: + with _QUIET_LOCK: + _quiet -= 1 + if not _quiet: + _HEADERS.removeFilter(_no_header_error) @dataclass(frozen=True) @@ -94,11 +121,12 @@ def explicit(self) -> bool: @contextmanager def scope(self) -> Iterator[None]: r"""Runs the block with this target's credentials.""" - if self._swap is None: - yield - return - with _LOCK, self._swap(): - yield + with _without_header_warnings(): + if self._swap is None: + yield + else: + with _LOCK, self._swap(): + yield def client(self) -> Any: mlflow = import_provider('mlflow', 'mlflow') diff --git a/packages/kumo-engine/test/test_mlflow_target.py b/packages/kumo-engine/test/test_mlflow_target.py index e883cd238..d052f0cb6 100644 --- a/packages/kumo-engine/test/test_mlflow_target.py +++ b/packages/kumo-engine/test/test_mlflow_target.py @@ -374,3 +374,24 @@ class Page(list): tags=held[version])) assert _registered_version(client, 'm', {'kumo_bundle_id': 'b'}) == '2' assert _registered_version(client, 'm', {'kumo_bundle_id': 'x'}) is None + + +def test_kumo_calls_drop_mlflow_request_header_errors_only(caplog): + import logging + + from kumo.engine.integrations._mlflow_target import resolve + + headers = logging.getLogger('mlflow.tracking.request_header.registry') + message = ('Encountered unexpected error during resolving request ' + 'headers: %s') + headers.addHandler(caplog.handler) # MLflow's loggers do not propagate + try: + with caplog.at_level('WARNING', logger=headers.name): + with resolve('sqlite:///x.db').scope(): + headers.warning(message, 'inside') + headers.warning('another warning') + headers.warning(message, 'outside') + finally: + headers.removeHandler(caplog.handler) + assert 'inside' not in caplog.text + assert 'another warning' in caplog.text and 'outside' in caplog.text From 31350a9df7b1ac4e3fdff5958058503b0328c52d Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 18:35:31 +0000 Subject: [PATCH 037/278] Silence Optuna's experimental-plot warnings where the searcher plots A tracked search logged 'optuna.visualization.matplotlib...plot_... is experimental' for every plot. The searcher's process-wide filter, added at construction, did not hold where analyze() ran. analyze() now plots with Optuna's ExperimentalWarning ignored locally, and the process-wide filter goes. Verified: kumo-core test_optuna_searcher (new: analyze warns nothing). Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_core/automl/optuna_searcher.py | 12 +++++++----- packages/kumo-core/tests/test_optuna_searcher.py | 13 +++++++++++++ 2 files changed, 20 insertions(+), 5 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/automl/optuna_searcher.py b/packages/kumo-core/src/kumo_core/automl/optuna_searcher.py index f5a23de69..c6e4b5fb8 100644 --- a/packages/kumo-core/src/kumo_core/automl/optuna_searcher.py +++ b/packages/kumo-core/src/kumo_core/automl/optuna_searcher.py @@ -28,11 +28,6 @@ def __init__( max_attempts: int = 50, **kwargs: Any, ) -> None: - warnings.filterwarnings( - "ignore", - message=".*plot_.*", - category=optuna.exceptions.ExperimentalWarning, - ) super().__init__(space, max_attempts) direction = "maximize" if self.tune_mode == "max" else "minimize" with warnings.catch_warnings(): @@ -90,6 +85,13 @@ def sample_value(self, key: str, value: list[Any]) -> Any: return self.current_trial.suggest_categorical(key, value) def analyze(self, path: str) -> None: + # Optuna flags its matplotlib plots as experimental on every call. + with warnings.catch_warnings(): + warnings.simplefilter("ignore", + optuna.exceptions.ExperimentalWarning) + self._analyze(path) + + def _analyze(self, path: str) -> None: import matplotlib import matplotlib.pyplot as plt from optuna.visualization.matplotlib import ( diff --git a/packages/kumo-core/tests/test_optuna_searcher.py b/packages/kumo-core/tests/test_optuna_searcher.py index 9d7f9f4f8..6dd82f8c2 100644 --- a/packages/kumo-core/tests/test_optuna_searcher.py +++ b/packages/kumo-core/tests/test_optuna_searcher.py @@ -40,3 +40,16 @@ def sampler(**kwargs: Any) -> optuna.samplers.BaseSampler: # deprecated and False is the default. assert ('warn_independent_sampling' in sampler_kwargs) == passed assert sampler_kwargs.get('warn_independent_sampling', False) is False + + +def test_analyze_emits_no_experimental_warnings(tmp_path) -> None: + searcher = TPESearcher(SPACE) + searcher.study.optimize(lambda trial: trial.suggest_float('lr', 0, 1), + n_trials=3) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter('always') + searcher.analyze(str(tmp_path)) + assert [ + str(warning.message) for warning in caught + if issubclass(warning.category, optuna.exceptions.ExperimentalWarning) + ] == [] From 821ec212ef59a66ca1f7e0d916eb8b2fe2567810 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 18:40:19 +0000 Subject: [PATCH 038/278] Give MLflow registrations the fit's metrics result.registered_model showed metrics={} for an MLflow version, while a Snowflake version carried the fit's metrics. A registration that comes back without metrics now gets the fit's, from fit and from result.export(cfg). Verified: test_mlflow_model, test_trainer_registry. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/trainer/trainer.py | 13 ++++++++++++- packages/kumo-engine/test/test_mlflow_model.py | 2 +- 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 99a586f4e..750de6629 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -375,6 +375,7 @@ def _store_and_register( registered, error = (None, None) if session is None else session.log_model( bundle_dir, manifest, snapshot, None if stored else _mlflow_registry(settings)) + registered = _with_metrics(registered, metrics) error = stored or error registry = None if stored else _connector_registry(settings) if registry is None: @@ -469,13 +470,23 @@ def _register_result(result: TrainingResult, settings: Any) -> PublishedModel: 'not set.') register = getattr(settings.tracking, 'register', None) if callable(register): - return register(result, settings.registry) + return cast(PublishedModel, _with_metrics( + register(result, settings.registry), result.metrics)) registry = cast(ModelOutputConfig, _connector_registry(settings)) _require_publisher(registry) return _publish(result.bundle_dir, result.bundle_id, result.metrics, registry) +def _with_metrics(model: Optional[PublishedModel], + metrics: Mapping[str, float]) -> Optional[PublishedModel]: + r"""``model`` with the fit's metrics, where its registry returned none + (an MLflow version, unlike a Snowflake one, carries none).""" + if model is None or model.metrics: + return model + return replace(model, metrics=dict(metrics)) + + def _mlflow_registry(settings: Any) -> Optional[ModelRegistryLocation]: if settings is None or not _logs_models(settings.tracking): return None diff --git a/packages/kumo-engine/test/test_mlflow_model.py b/packages/kumo-engine/test/test_mlflow_model.py index 51eb51fb7..2a0c0ca42 100644 --- a/packages/kumo-engine/test/test_mlflow_model.py +++ b/packages/kumo-engine/test/test_mlflow_model.py @@ -70,7 +70,7 @@ def test_fit_logs_registers_and_serves_the_model(table_and_snapshot, assert result.registered_model == PublishedModel( name='m', version=str(version.version), uri=f'models:/m/{version.version}', - metadata={'kumo_bundle_id': result.id}) + metadata={'kumo_bundle_id': result.id}, metrics=result.metrics) assert version.tags == {'kumo_bundle_id': result.id} assert trainer.result is result From 03d71a32abff62e6119a89b768d3ed389779ae56 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 18:42:52 +0000 Subject: [PATCH 039/278] Print progress events readably; drop two more dependency notes callbacks=[print] printed raw Event(...) reprs with epoch times; str(event) is now the name and its fields, e.g. 'table_staged table=users seconds=12.0'. Every fit also showed Python's 'process is multi-threaded, use of fork() may lead to deadlocks' DeprecationWarning and Lightning's 'IterableDataset has __len__ defined' UserWarning; fit and predict now drop those two, as they drop PyG's scatter hint. Verified: test_fit_events (new: str(event); the two notes dropped). Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_events.py | 7 +++++ .../kumo/engine/trainer/trainer.py | 14 +++++++-- packages/kumo-engine/test/test_fit_events.py | 29 +++++++++++++++++++ 3 files changed, 48 insertions(+), 2 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_events.py b/packages/kumo-engine/kumo/engine/_events.py index ecf1c2a79..dbbef3b83 100644 --- a/packages/kumo-engine/kumo/engine/_events.py +++ b/packages/kumo-engine/kumo/engine/_events.py @@ -42,6 +42,13 @@ class Event: def to_dict(self) -> Dict[str, Any]: return asdict(self) + def __str__(self) -> str: + fields = {'trial': self.trial, 'epoch': self.epoch, **self.metrics, + **self.params} + return ' '.join([self.name, *(f'{key}={value}' + for key, value in fields.items() + if value is not None)]) + def query_sha256(query: str) -> str: return hashlib.sha256(query.encode()).hexdigest() diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 750de6629..4f2b99b13 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -555,8 +555,11 @@ def _stage_model_bundle(artifacts: Mapping[str, str], dest: Path) -> Path: def _without_torch_scatter_hint(method: _F) -> _F: r"""Runs ``method`` without PyTorch Geometric's GPU hint that min/max - aggregation "can be accelerated via the 'torch-scatter' package", and - without Lightning's "Tip" log lines advertising its cloud products. + aggregation "can be accelerated via the 'torch-scatter' package", + Lightning's "Tip" log lines advertising its cloud products, Lightning's + note that the loader's IterableDataset defines ``__len__``, and Python's + note that forking a multi-threaded process may deadlock (the loader + workers fork before any thread that matters holds a lock). kumo-engine uses PyTorch's own scatter instead of that optional extension (see ``test_wheel_uses_torch_fallbacks_without_native_extensions``), so @@ -574,6 +577,13 @@ def wrapper(*args: Any, **kwargs: Any) -> Any: r"accelerated via the 'torch-scatter' package", category=UserWarning, ) + warnings.filterwarnings( + 'ignore', message=r'Your `IterableDataset` has `__len__`', + category=UserWarning) + warnings.filterwarnings( + 'ignore', message=r'This process .* is multi-threaded, ' + r'use of fork\(\) may lead to deadlocks', + category=DeprecationWarning) return method(*args, **kwargs) finally: tips.removeFilter(_no_lightning_tip) diff --git a/packages/kumo-engine/test/test_fit_events.py b/packages/kumo-engine/test/test_fit_events.py index fa34faae4..5b8753b27 100644 --- a/packages/kumo-engine/test/test_fit_events.py +++ b/packages/kumo-engine/test/test_fit_events.py @@ -194,3 +194,32 @@ def fit(): lightning.removeHandler(caplog.handler) assert 'inside' not in caplog.text assert 'GPU available' in caplog.text and 'outside' in caplog.text + + +def test_an_event_prints_as_its_name_and_fields(): + event = Event(name='table_staged', time=1.0, + params={'table': 'users', 'seconds': 12.0}) + assert str(event) == 'table_staged table=users seconds=12.0' + assert str(Event(name='epoch_ended', time=1.0, trial=1, epoch=2, + metrics={'val_auroc': 0.5})) == ( + 'epoch_ended trial=1 epoch=2 val_auroc=0.5') + + +def test_kumo_calls_drop_the_fork_and_iterable_dataset_notes(): + import warnings + + from kumo.engine.trainer.trainer import _without_torch_scatter_hint + + @_without_torch_scatter_hint + def fit(): + warnings.warn('This process (pid=1) is multi-threaded, use of fork() ' + 'may lead to deadlocks in the child.', + DeprecationWarning) + warnings.warn('Your `IterableDataset` has `__len__` defined.', + UserWarning) + warnings.warn('something else', UserWarning) + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter('always') + fit() + assert [str(warning.message) for warning in caught] == ['something else'] From 4bbf4cc7b635a4895d868bd4bbadcb563738d65d Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 20:21:01 +0000 Subject: [PATCH 040/278] Log metrics as train/, val/ and test/ series in trackers Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_tracking.py | 23 ++++++++++++++++--- .../kumo-engine/test/test_mlflow_tracking.py | 11 +++++---- .../test/test_snowflake_tracking.py | 4 ++-- .../test/test_tracking_contract.py | 12 ++++++---- 4 files changed, 35 insertions(+), 15 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_tracking.py b/packages/kumo-engine/kumo/engine/_tracking.py index 159c9a83a..2783bff6b 100644 --- a/packages/kumo-engine/kumo/engine/_tracking.py +++ b/packages/kumo-engine/kumo/engine/_tracking.py @@ -124,13 +124,28 @@ def _text(value: Any) -> Optional[str]: return None if value is None else str(value) +def logged_name(name: str) -> str: + r"""A metric's name in a tracker: ``val_auroc`` is ``val/auroc`` and a + trial's ``best_val_auroc`` is ``val/best_auroc``, so charts group by + split.""" + best = name.startswith('best_') + split, _, metric = name.removeprefix('best_').partition('_') + if split not in ('train', 'val', 'test') or not metric: + return name + return f"{split}/{'best_' if best else ''}{metric}" + + +def logged_names(metrics: Mapping[str, Any]) -> Dict[str, Any]: + return {logged_name(name): value for name, value in metrics.items()} + + def run_record(spec: RunSpec) -> RunRecord: params = {key: _param_value(value) for key, value in spec.params.items()} return RunRecord( name=spec.name, params=params, - metrics=_finite_floats(spec.metrics), - steps={name: tuple(history) + metrics=logged_names(_finite_floats(spec.metrics)), + steps={logged_name(name): tuple(history) for name, history in spec.steps.items()}, tags={ **_user_tags(spec.tags, params), @@ -293,8 +308,10 @@ def on_event(self, event: Event) -> None: }) events.append(event) record = run_record(trial_spec(events)) + logged = dataclasses.replace(event, + metrics=logged_names(event.metrics)) self._submit(lambda: self._call( - lambda tracker: tracker.on_trial(record, event))) + lambda tracker: tracker.on_trial(record, logged))) def log_model( self, diff --git a/packages/kumo-engine/test/test_mlflow_tracking.py b/packages/kumo-engine/test/test_mlflow_tracking.py index 709b9885f..4fd7cf0df 100644 --- a/packages/kumo-engine/test/test_mlflow_tracking.py +++ b/packages/kumo-engine/test/test_mlflow_tracking.py @@ -21,7 +21,7 @@ import kumo.engine as kumo from kumo.engine._scratch import get_scratch_dir -from kumo.engine._tracking import run_record +from kumo.engine._tracking import logged_names, run_record from kumo.engine.trainer._run_spec import fit_spec mlflow = pytest.importorskip('mlflow') @@ -134,7 +134,8 @@ def test_fit_logs_a_run_per_fit_and_per_trial(table_and_snapshot, tmp_path, for name, value in parent.data.metrics.items() if not name.startswith('system/') } - assert metrics == {**result.metrics, 'seconds': metrics['seconds']} + assert metrics == {**logged_names(result.metrics), + 'seconds': metrics['seconds']} assert 'off for the rest of this run' not in caplog.text ents, = parent.inputs.dataset_inputs assert (ents.dataset.name, ents.dataset.source_type) == ('ents', 'local') @@ -147,9 +148,9 @@ def test_fit_logs_a_run_per_fit_and_per_trial(table_and_snapshot, tmp_path, for child in children: assert child.info.status == 'FINISHED' assert child.data.tags['kumo.status'] == 'succeeded' - assert {'best_val_mae', 'seconds', - 'val_mae'} <= set(child.data.metrics) - history = client.get_metric_history(child.info.run_id, 'val_mae') + assert {'val/best_mae', 'seconds', + 'val/mae'} <= set(child.data.metrics) + history = client.get_metric_history(child.info.run_id, 'val/mae') assert sorted(metric.step for metric in history) == [1, 2] diff --git a/packages/kumo-engine/test/test_snowflake_tracking.py b/packages/kumo-engine/test/test_snowflake_tracking.py index 7cbbacec9..b982f4e03 100644 --- a/packages/kumo-engine/test/test_snowflake_tracking.py +++ b/packages/kumo-engine/test/test_snowflake_tracking.py @@ -175,8 +175,8 @@ def test_metric_and_param_limits(fake, monkeypatch): run, = fake._instance.runs.values() assert sorted({name for name, _, _ in run['metrics']}) == [ - 'test_mae', 'trial-1-attempt-1/best_val_mae', - 'trial-1-attempt-1/train_loss', 'trial-1-attempt-1/val_mae', 'val_mae' + 'test/mae', 'trial-1-attempt-1/train/loss', + 'trial-1-attempt-1/val/best_mae', 'trial-1-attempt-1/val/mae', 'val/mae' ] assert run['params']['kumo.dropped_metrics'] == '6' assert run['params']['kumo.dropped_params'] == '1' diff --git a/packages/kumo-engine/test/test_tracking_contract.py b/packages/kumo-engine/test/test_tracking_contract.py index 0d7bedbf1..9815184dd 100644 --- a/packages/kumo-engine/test/test_tracking_contract.py +++ b/packages/kumo-engine/test/test_tracking_contract.py @@ -267,7 +267,9 @@ def observe(): names: Dict[Optional[str], Dict[str, Dict[str, Any]]] = {} def bucket(key): - prefix, _, name = key.rpartition('/') + prefix, _, name = key.partition('/') + if not prefix.startswith('trial-'): + prefix, name = '', key return names.setdefault(prefix or None, { 'params': {}, 'tags': {}, 'metrics': {}, 'steps': {} }), name @@ -319,7 +321,7 @@ def test_records_of_the_recorded_fit(): spoofed = run_record(fit_spec(CONTEXT, tags={ 'kumo.bundle_id': 'x', 'training_job.tune_metric': 'y', 'a': 1})) assert spoofed.tags.keys() - fit.tags.keys() == {'a'} - assert fit.metrics == {'val_mae': 0.3, 'test_mae': 0.32, 'seconds': 9.0} + assert fit.metrics == {'val/mae': 0.3, 'test/mae': 0.32, 'seconds': 9.0} assert fit.artifacts['query.pql'] == QUERY assert json.loads(fit.artifacts['model_plan.json']) == CONTEXT.model_plan assert fit.inputs == INPUTS @@ -334,9 +336,9 @@ def test_records_of_the_recorded_fit(): }) retry = records['trial-2-attempt-2'] assert retry.params['attempt'] == '2' - assert retry.metrics == {'best_val_mae': 0.3, 'seconds': 4.0} - assert retry.steps == {'train_loss': ((1, 0.6), (2, 0.4)), - 'val_mae': ((1, 0.5), (2, 0.3))} + assert retry.metrics == {'val/best_mae': 0.3, 'seconds': 4.0} + assert retry.steps == {'train/loss': ((1, 0.6), (2, 0.4)), + 'val/mae': ((1, 0.5), (2, 0.3))} def test_adapter_records_the_same_runs(adapter, tmp_path): From 52879e5b35a0e1f718f7bb71a656572b24738985 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 20:27:42 +0000 Subject: [PATCH 041/278] Chart a fit's own epochs, or a search's progress and selected trial, on the fit run A single-trial fit logs train/ and val/ on the fit run as each epoch ends. A search logs search/best_val_ after each trial, then the selected trial's epochs once the fit succeeds. The fit run's summary keeps the test metrics; trial runs are unchanged. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_tracking.py | 62 ++++++++++++++++++- .../kumo/engine/integrations/mlflow.py | 12 +++- .../kumo/engine/integrations/snowflake.py | 13 +++- .../kumo/engine/trainer/_run_spec.py | 7 ++- .../kumo-engine/test/test_mlflow_tracking.py | 15 ++++- .../test/test_snowflake_tracking.py | 6 +- .../test/test_tracking_contract.py | 38 ++++++++++-- 7 files changed, 137 insertions(+), 16 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_tracking.py b/packages/kumo-engine/kumo/engine/_tracking.py index 2783bff6b..7f0cf98aa 100644 --- a/packages/kumo-engine/kumo/engine/_tracking.py +++ b/packages/kumo-engine/kumo/engine/_tracking.py @@ -139,6 +139,17 @@ def logged_names(metrics: Mapping[str, Any]) -> Dict[str, Any]: return {logged_name(name): value for name, value in metrics.items()} +def by_step(steps: Mapping[str, Sequence[Tuple[int, float]]] + ) -> List[Tuple[int, Dict[str, float]]]: + r"""Per-step histories regrouped as each step's metrics, in step + order.""" + grouped: Dict[int, Dict[str, float]] = {} + for name, history in steps.items(): + for step, value in history: + grouped.setdefault(step, {})[name] = value + return sorted(grouped.items()) + + def run_record(spec: RunSpec) -> RunRecord: params = {key: _param_value(value) for key, value in spec.params.items()} return RunRecord( @@ -211,7 +222,9 @@ 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); ``log_model(bundle_dir, manifest, + ``log_metrics(metrics, step)``, series of the fit run as they come (a + single trial's epochs, or a search's best validation metric after each + trial); ``tags`` (extra fit-run tags); ``log_model(bundle_dir, manifest, 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 @@ -232,6 +245,41 @@ def finish(self, record: Optional[RunRecord], ... +@dataclass +class _Search: + r"""The fit run's own series: a single trial's epochs as they end, or, + for a search, the best validation ``metric`` after each trial; then + `best`, the selected trial, whose epochs the run records at the end.""" + + single: bool + metric: str + lower: bool + best: Optional[Tuple[float, int]] = None + trials: int = 0 + + @classmethod + def of(cls, params: Mapping[str, Any]) -> '_Search': + from kumo_core.evaluation.suite import MetricStore + + metric = f"val_{params.get('training_job.tune_metric', 'loss')}" + return cls(single=params.get('training_job.num_experiments') == 1, + metric=metric, + lower=MetricStore.tune_mode(metric) == 'min') + + def on(self, event: Event) -> Optional[Tuple[Dict[str, float], int]]: + if event.name == 'epoch_ended' and self.single: + return dict(event.metrics), event.epoch or 0 + value = event.metrics.get(self.metric) + if (event.name != 'trial_ended' or self.single or value is None + or event.params.get('status') != 'succeeded'): + return None + if self.best is None or (value < self.best[0] if self.lower else + value > self.best[0]): + self.best = (value, event.trial or 0) + self.trials += 1 + return {f'search/best_{self.metric}': self.best[0]}, self.trials + + class RegistrationError(Exception): r"""Raised by a tracker's ``log_model`` from the error that registering the logged model raised.""" @@ -244,6 +292,7 @@ def __init__(self, tracking: Optional[Tracking], self._warn = warn self.spec: Optional[RunSpec] = None self._trials: Dict[int, List[Event]] = {} + self._search: Optional[_Search] = None self._jobs: 'queue.Queue[Optional[Callable[[], Any]]]' = queue.Queue( QUEUE_SIZE) self._worker: Optional[threading.Thread] = None @@ -267,6 +316,7 @@ def start(self, spec: Callable[[], RunSpec]) -> None: self._tracker = None return self.spec = value + self._search = _Search.of(value.params) self._call(lambda tracker: getattr(tracker, 'prepare', _none)()) self._worker = threading.Thread(target=self._work, name='kumo-tracking', daemon=True) @@ -312,6 +362,12 @@ def on_event(self, event: Event) -> None: metrics=logged_names(event.metrics)) self._submit(lambda: self._call( lambda tracker: tracker.on_trial(record, logged))) + series = None if self._search is None else self._search.on(event) + if series is not None: + metrics, step = logged_names(series[0]), series[1] + self._submit(lambda: self._call( + lambda tracker: getattr(tracker, 'log_metrics', _none)( + metrics, step))) def log_model( self, @@ -384,6 +440,10 @@ def finish(self, spec: Optional[Callable[[], RunSpec]], return try: final = self.spec if spec is None else spec() + selected = None if self._search is None else self._search.best + if error is None and selected is not None: + final = dataclasses.replace( + final, steps=trial_spec(self._trials[selected[1]]).steps) if error is not None and final.error is None: final = dataclasses.replace(final, error=format_error(error)) record: Optional[RunRecord] = run_record(final) diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index 774eb1bfb..bb6b28659 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -37,7 +37,12 @@ from kumo.engine._events import Event from kumo.engine._mlops import full_error -from kumo.engine._tracking import DROPPED_PARAMS, RegistrationError, RunRecord +from kumo.engine._tracking import ( + DROPPED_PARAMS, + RegistrationError, + RunRecord, + by_step, +) from kumo.engine._version import __version__ from kumo.engine.connector._extras import import_provider from kumo.engine.integrations._mlflow_target import ( @@ -403,6 +408,8 @@ def finish(self, record: Optional[RunRecord], and record.tags.get('kumo.kind') == 'fit') try: if record is not None and self._parent is not None: + for epoch, metrics in by_step(record.steps): + self._log_metrics(self._parent, metrics, epoch) self._log_metrics(self._parent, record.metrics, model_id=self._model_id) self._log_tags(self._parent, record.tags) @@ -422,6 +429,9 @@ def finish(self, record: Optional[RunRecord], if fit_run and hint: warnings.warn(hint, UserWarning, stacklevel=2) + def log_metrics(self, metrics: Mapping[str, float], step: int) -> None: + self._log_metrics(self._parent, metrics, step) + def _log_system_metrics(self, metrics: Mapping[str, float], step: int) -> None: self._log_metrics(self._parent, metrics, step) diff --git a/packages/kumo-engine/kumo/engine/integrations/snowflake.py b/packages/kumo-engine/kumo/engine/integrations/snowflake.py index 78b3e23b6..382c58f66 100644 --- a/packages/kumo-engine/kumo/engine/integrations/snowflake.py +++ b/packages/kumo-engine/kumo/engine/integrations/snowflake.py @@ -16,7 +16,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, + by_step, +) from kumo.engine.connector._extras import import_provider logger = logging.getLogger(__name__) @@ -190,6 +195,8 @@ def finish(self, record: Optional[RunRecord], status = 'FAILED' try: if self._run is not None and record is not None: + for step, metrics in by_step(record.steps): + self.log_metrics(metrics, step) self._log_metrics(record.metrics, _MAX_METRICS) self._log_params({k: v for k, v in record.tags.items() @@ -209,6 +216,10 @@ def finish(self, record: Optional[RunRecord], self._run_context.close() self._reset() + def log_metrics(self, metrics: Mapping[str, float], step: int) -> None: + if self._run is not None: + self._log_metrics(metrics, _EPOCH_METRICS, step=step) + def _new_run_name(self, tracking: Any, experiment: Any) -> Optional[str]: if self.run_name is None: diff --git a/packages/kumo-engine/kumo/engine/trainer/_run_spec.py b/packages/kumo-engine/kumo/engine/trainer/_run_spec.py index 1a916aad4..81e79918d 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_run_spec.py +++ b/packages/kumo-engine/kumo/engine/trainer/_run_spec.py @@ -41,14 +41,17 @@ def fit_spec(context: FitContext, *, def fitted_spec(spec: RunSpec, result: Any, seconds: Optional[float] = None) -> RunSpec: r"""``spec`` once the fit returned ``result``, a `TrainingResult`, - after ``seconds``.""" + after ``seconds``: its test metrics, since the run's train and + validation series show the selected trial's epochs.""" artifacts = dict(spec.artifacts) provenance = Path(result.bundle_dir) / PROVENANCE_FILE if provenance.is_file(): artifacts[PROVENANCE_FILE] = provenance.read_text(encoding='utf-8') + tested = {name: value for name, value in result.metrics.items() + if not name.startswith(('train_', 'val_'))} return replace( spec, metrics={ - **result.metrics, **({} if seconds is None else { + **tested, **({} if seconds is None else { 'seconds': seconds }) }, parents={'bundle_id': result.id}, artifacts=artifacts) diff --git a/packages/kumo-engine/test/test_mlflow_tracking.py b/packages/kumo-engine/test/test_mlflow_tracking.py index 4fd7cf0df..e68a2d3b8 100644 --- a/packages/kumo-engine/test/test_mlflow_tracking.py +++ b/packages/kumo-engine/test/test_mlflow_tracking.py @@ -134,8 +134,12 @@ def test_fit_logs_a_run_per_fit_and_per_trial(table_and_snapshot, tmp_path, for name, value in parent.data.metrics.items() if not name.startswith('system/') } - assert metrics == {**logged_names(result.metrics), - 'seconds': metrics['seconds']} + tested = {name: value for name, value in logged_names( + result.metrics).items() if name.startswith('test/')} + assert set(metrics) == {*tested, 'seconds', 'search/best_val_mae', + *(name for name in metrics + if name.startswith(('train/', 'val/')))} + assert {name: metrics[name] for name in tested} == tested assert 'off for the rest of this run' not in caplog.text ents, = parent.inputs.dataset_inputs assert (ents.dataset.name, ents.dataset.source_type) == ('ents', 'local') @@ -152,6 +156,13 @@ def test_fit_logs_a_run_per_fit_and_per_trial(table_and_snapshot, tmp_path, 'val/mae'} <= set(child.data.metrics) history = client.get_metric_history(child.info.run_id, 'val/mae') assert sorted(metric.step for metric in history) == [1, 2] + selected = min(children, key=lambda run: run.data.metrics['val/best_mae']) + for name in ('train/loss', 'val/mae'): + curve, trial = (sorted((metric.step, metric.value) + for metric in client.get_metric_history( + run.info.run_id, name)) + for run in (parent, selected)) + assert curve == trial and len(curve) == 2 def test_failed_fit_marks_the_run_failed(table_and_snapshot, tmp_path, diff --git a/packages/kumo-engine/test/test_snowflake_tracking.py b/packages/kumo-engine/test/test_snowflake_tracking.py index b982f4e03..f2c7786ec 100644 --- a/packages/kumo-engine/test/test_snowflake_tracking.py +++ b/packages/kumo-engine/test/test_snowflake_tracking.py @@ -175,10 +175,10 @@ def test_metric_and_param_limits(fake, monkeypatch): run, = fake._instance.runs.values() assert sorted({name for name, _, _ in run['metrics']}) == [ - 'test/mae', 'trial-1-attempt-1/train/loss', - 'trial-1-attempt-1/val/best_mae', 'trial-1-attempt-1/val/mae', 'val/mae' + 'seconds', 'test/mae', 'trial-1-attempt-1/train/loss', + 'trial-1-attempt-1/val/best_mae', 'trial-1-attempt-1/val/mae' ] - assert run['params']['kumo.dropped_metrics'] == '6' + assert run['params']['kumo.dropped_metrics'] == '8' assert run['params']['kumo.dropped_params'] == '1' assert 'note' not in run['params'] assert run['params']['level'] == '2' diff --git a/packages/kumo-engine/test/test_tracking_contract.py b/packages/kumo-engine/test/test_tracking_contract.py index 9815184dd..19c8ce2cf 100644 --- a/packages/kumo-engine/test/test_tracking_contract.py +++ b/packages/kumo-engine/test/test_tracking_contract.py @@ -76,6 +76,10 @@ def _epochs(trial, *values): _event('fit_ended', metrics={'val_mae': 0.3, 'test_mae': 0.32, 'seconds': 9.0}, status='succeeded'), ] +# The fit run's own series: the best val_mae after each trial that +# succeeded and, once the fit succeeds, the selected trial 2's epochs. +SEARCH = {'search/best_val_mae': ((1, 0.35), (2, 0.3))} +SELECTED = {'train_loss': ((1, 0.6), (2, 0.4)), 'val_mae': ((1, 0.5), (2, 0.3))} BUNDLE = 'b' * 64 RESULT = SimpleNamespace(metrics={'val_mae': 0.3, 'test_mae': 0.32}, id=BUNDLE, bundle_dir=Path('no-bundle')) @@ -100,6 +104,10 @@ def on_trial(self, record, event): if self.tracker is not None: self.tracker.on_trial(record, event) + def log_metrics(self, metrics, step): + self.calls.append(('log_metrics', (metrics, step))) + getattr(self.tracker, 'log_metrics', lambda *_: None)(metrics, step) + def finish(self, record, error): self.calls.append(('finish', record)) if self.tracker is not None: @@ -112,13 +120,13 @@ def join_tracking_threads(): thread.join(30) -def replay(tracker=None, events=EVENTS, result=RESULT, - error=None) -> Dict[str, RunRecord]: +def replay(tracker=None, events=EVENTS, result=RESULT, error=None, + context=CONTEXT, series=None) -> Dict[str, RunRecord]: r"""Feeds the recorded fit to `tracker`; returns the records expected - of it.""" + of it, with the fit run's ``series`` (default: the search's).""" recorder = _Recorder(tracker) session = TrackingSession(recorder) - session.start(lambda: fit_spec(CONTEXT)) + session.start(lambda: fit_spec(context)) for event in events: session.on_event(event) session.finish( @@ -128,9 +136,11 @@ def replay(tracker=None, events=EVENTS, result=RESULT, record.name: record for hook, record in recorder.calls if hook == 'on_trial' } - spec = fit_spec(CONTEXT, tags=USER_TAGS) + spec = fit_spec(context, tags=USER_TAGS) if result is not None: spec = fitted_spec(spec, result, events[-1].metrics.get('seconds')) + spec = replace(spec, steps=series if series is not None else { + **SEARCH, **({} if error else SELECTED)}) if 'error' in events[-1].params: spec = replace(spec, error=str(events[-1].params['error'])) records[FIT] = run_record(spec) @@ -321,7 +331,10 @@ def test_records_of_the_recorded_fit(): spoofed = run_record(fit_spec(CONTEXT, tags={ 'kumo.bundle_id': 'x', 'training_job.tune_metric': 'y', 'a': 1})) assert spoofed.tags.keys() - fit.tags.keys() == {'a'} - assert fit.metrics == {'val/mae': 0.3, 'test/mae': 0.32, 'seconds': 9.0} + assert fit.metrics == {'test/mae': 0.32, 'seconds': 9.0} + assert fit.steps == {'search/best_val_mae': ((1, 0.35), (2, 0.3)), + 'train/loss': ((1, 0.6), (2, 0.4)), + 'val/mae': ((1, 0.5), (2, 0.3))} assert fit.artifacts['query.pql'] == QUERY assert json.loads(fit.artifacts['model_plan.json']) == CONTEXT.model_plan assert fit.inputs == INPUTS @@ -351,6 +364,19 @@ def test_adapter_records_the_same_runs(adapter, tmp_path): assert gaps(expected, adapter.observe(), adapter) == [] +def test_adapter_records_a_single_trial_fit_epoch_by_epoch(adapter): + plan = {**CONTEXT.model_plan, + 'training_job': {'tune_metric': 'mae', 'num_experiments': 1}} + events = [EVENTS[0], *_epochs(1, (0.5, 0.4), (0.3, 0.35)), EVENTS[3], + EVENTS[-1]] + expected = replay(adapter.make(), events, + context=replace(CONTEXT, model_plan=plan), series={ + 'train_loss': ((1, 0.5), (2, 0.3)), + 'val_mae': ((1, 0.4), (2, 0.35))}) + + assert gaps(expected, adapter.observe(), adapter) == [] + + def test_a_raising_tracker_is_ended_and_not_called_again(caplog): class RaisesOnTrial(_Recorder): def on_trial(self, record, event): From 68431b84dac9578a9afa3be94a9a5885bafc6dda Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 20:36:05 +0000 Subject: [PATCH 042/278] Name fit runs after their model and tag task type, snapshot and registered version Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/integrations/mlflow.py | 9 +++--- .../kumo/engine/integrations/snowflake.py | 20 +++++++------ .../kumo/engine/trainer/_run_spec.py | 28 ++++++++++++++++--- .../kumo/engine/trainer/trainer.py | 14 ++++++++-- .../kumo-engine/test/test_mlflow_tracking.py | 7 +++-- .../test/test_snowflake_tracking.py | 4 +-- .../test/test_tracking_contract.py | 25 ++++++++++++----- 7 files changed, 77 insertions(+), 30 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index bb6b28659..9b763e9e1 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -92,8 +92,8 @@ class MlflowTracking: experiment_name: experiment for the fit run, created if missing. `None` uses MLflow's current experiment; required to record runs at a ``target``. - run_name: name of the fit run. `None` names it after its UTC start - time, as ``fit-20261027T061838Z``. + run_name: name of the fit run. `None` names it after its registry + model and UTC start, as ``fit-users-churn-20261027-0618``. log_system_metrics: sample CPU, memory, disk, network and (with ``nvidia-ml-py``) GPU usage on the fit run. tags: extra tags for the fit run. @@ -329,8 +329,9 @@ def start(self, record: RunRecord) -> None: } run = client.create_run( experiment_id or '0', - run_name=record.name or self.run_name - or time.strftime('fit-%Y%m%dT%H%M%SZ', time.gmtime()), + run_name=(self.run_name if self.run_name + and record.tags.get('kumo.kind') == 'fit' else + record.name), tags=context.resolve_tags({**record.tags, **parent})) self._parent = self.run_id = run.info.run_id self._experiment_id = run.info.experiment_id diff --git a/packages/kumo-engine/kumo/engine/integrations/snowflake.py b/packages/kumo-engine/kumo/engine/integrations/snowflake.py index 382c58f66..ea3c6eeb8 100644 --- a/packages/kumo-engine/kumo/engine/integrations/snowflake.py +++ b/packages/kumo-engine/kumo/engine/integrations/snowflake.py @@ -51,7 +51,8 @@ class SnowflakeTracking: database: `None` uses the session's current database. schema_name: `None` uses the session's current schema, or ``PUBLIC`` when `database` is given or the session has none. - run_name: name of the fit run. `None` lets Snowflake pick one; a + run_name: name of the fit run. `None` names it after its registry + model and UTC start, as ``"fit-USERS_CHURN-20261027-0618"``; a taken name gets a random suffix. tags: extra tags for the fit run. native_model_lineage: keep the run open as a context for the whole @@ -159,7 +160,7 @@ def start(self, record: RunRecord) -> None: database_name=database, schema_name=schema or 'PUBLIC') self._run = tracking.start_run( - self._new_run_name(tracking, experiment.name)) + self._new_run_name(tracking, experiment.name, record.name)) self.run_id = str(self._run.name) if self.native_model_lineage: self._run_context.enter_context(self._run) @@ -220,9 +221,10 @@ def log_metrics(self, metrics: Mapping[str, float], step: int) -> None: if self._run is not None: self._log_metrics(metrics, _EPOCH_METRICS, step=step) - def _new_run_name(self, tracking: Any, - experiment: Any) -> Optional[str]: - if self.run_name is None: + def _new_run_name(self, tracking: Any, experiment: Any, + default: Optional[str]) -> Optional[str]: + wanted = self.run_name or (default and f'"{default}"') + if wanted is None: return None # start_run(name) resumes, or raises on, any run matching name as a # SQL LIKE pattern and cannot create an existing run, so pick a name @@ -233,16 +235,16 @@ def _new_run_name(self, tracking: Any, experiment_name=experiment) ] short = uuid.uuid4().hex - for name in (self.run_name, _suffixed(self.run_name, short[:8]), - _suffixed(self.run_name, short)): + for name in (wanted, _suffixed(wanted, short[:8]), + _suffixed(wanted, short)): if not _taken(name, existing): break else: name = None - if name != self.run_name: + if name != wanted: logger.warning( 'Snowflake run %s already exists in experiment %s; recording ' - 'this fit as %s', self.run_name, self.experiment_name, + 'this fit as %s', wanted, self.experiment_name, name or 'a generated run name') return name diff --git a/packages/kumo-engine/kumo/engine/trainer/_run_spec.py b/packages/kumo-engine/kumo/engine/trainer/_run_spec.py index 81e79918d..17b464ff3 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_run_spec.py +++ b/packages/kumo-engine/kumo/engine/trainer/_run_spec.py @@ -5,6 +5,7 @@ import json from dataclasses import dataclass, field, replace +from datetime import datetime, timezone from pathlib import Path from typing import Any, Mapping, Optional @@ -21,13 +22,22 @@ class FitContext: model_plan: Mapping[str, Any] inputs: Mapping[str, Mapping[str, Optional[str]]] = field( default_factory=dict) + started: Optional[datetime] = None + snapshot_id: Optional[str] = None + task_type: Optional[str] = None + model_name: Optional[str] = None def fit_spec(context: FitContext, *, tags: Optional[Mapping[str, Any]] = None) -> RunSpec: + started = context.started or datetime.now(timezone.utc) return RunSpec( kind='fit', + name='-'.join(filter(None, ('fit', context.model_name, + started.strftime('%Y%m%d-%H%M')))), params=dict(_flatten(context.model_plan)), + parents={'snapshot_id': context.snapshot_id, + 'task_type': context.task_type}, tags=dict(tags or {}), inputs=dict(context.inputs), artifacts={ @@ -38,11 +48,12 @@ def fit_spec(context: FitContext, *, ) -def fitted_spec(spec: RunSpec, result: Any, - seconds: Optional[float] = None) -> RunSpec: +def fitted_spec(spec: RunSpec, result: Any, seconds: Optional[float] = None, + alias: Optional[str] = None) -> RunSpec: r"""``spec`` once the fit returned ``result``, a `TrainingResult`, after ``seconds``: its test metrics, since the run's train and - validation series show the selected trial's epochs.""" + validation series show the selected trial's epochs, and the version it + registered with ``alias``.""" artifacts = dict(spec.artifacts) provenance = Path(result.bundle_dir) / PROVENANCE_FILE if provenance.is_file(): @@ -54,4 +65,13 @@ def fitted_spec(spec: RunSpec, result: Any, **tested, **({} if seconds is None else { 'seconds': seconds }) - }, parents={'bundle_id': result.id}, artifacts=artifacts) + }, parents={**spec.parents, 'bundle_id': result.id, + **_registered(result.registered_model, alias)}, + artifacts=artifacts) + + +def _registered(model: Any, alias: Optional[str]) -> Mapping[str, Any]: + if model is None: + return {} + return {'model_name': model.name, 'model_version': model.version, + 'model_alias': alias} diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 4f2b99b13..e6faf7044 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -748,7 +748,9 @@ def fit( scoring_defaults=scoring_defaults, warm_start=warm_start, run_id=run_id, events=events, tracking=session, - deps=deps) + deps=deps, model_name=None if ( + settings is None or settings.registry is + None) else settings.registry.name) except BaseException as error: events.emit('fit_ended', params={ 'status': 'failed', @@ -769,8 +771,11 @@ def fit( if registered is not None: result = replace(result, registered_model=registered) self._training_result = result + alias = None if settings is None or settings.registry is None else ( + settings.registry.alias) session.finish( - lambda: fitted_spec(cast(RunSpec, session.spec), result, seconds), + lambda: fitted_spec(cast(RunSpec, session.spec), result, seconds, + alias), error) raise_on_error(settings, session.error or error, 'Trainer.fit (its model is in trainer.result)') @@ -788,6 +793,7 @@ def _fit( events: EventEmitter, tracking: TrackingSession, deps: Dependencies, + model_name: Optional[str] = None, ) -> TrainingResult: from kumo_core.automl.runtime import ( inherit_warm_start_model_plan, @@ -931,6 +937,10 @@ def _fit( query=train_table.query, model_plan=to_json_dict(materialized.model_plan), inputs=fit_inputs(train_table.manifest), + started=started_at, + snapshot_id=snapshot.id, + task_type=materialized.task.task_type.value, + model_name=model_name, ))) training_input = ResolvedTrainingInput( task=materialized.task, diff --git a/packages/kumo-engine/test/test_mlflow_tracking.py b/packages/kumo-engine/test/test_mlflow_tracking.py index e68a2d3b8..21574a3dc 100644 --- a/packages/kumo-engine/test/test_mlflow_tracking.py +++ b/packages/kumo-engine/test/test_mlflow_tracking.py @@ -107,7 +107,7 @@ def test_fit_logs_a_run_per_fit_and_per_trial(table_and_snapshot, tmp_path, if run.data.tags.get('mlflow.parentRunId') == outer.info.run_id ] assert tracking.run_id == parent.info.run_id - assert re.fullmatch(r'fit-\d{8}T\d{6}Z', parent.info.run_name) + assert re.fullmatch(r'fit-\d{8}-\d{4}', parent.info.run_name) children = sorted( (run for run in runs if run.data.tags.get('mlflow.parentRunId') == parent.info.run_id), @@ -116,7 +116,10 @@ def test_fit_logs_a_run_per_fit_and_per_trial(table_and_snapshot, tmp_path, assert client.get_experiment(parent.info.experiment_id).name == 'ents' assert parent.info.status == 'FINISHED' assert {k for k in parent.data.tags if k.startswith('kumo.')} == { - 'kumo.kind', 'kumo.bundle_id', 'kumo.dropped_params'} + 'kumo.kind', 'kumo.bundle_id', 'kumo.dropped_params', + 'kumo.snapshot_id', 'kumo.task_type'} + assert parent.data.tags['kumo.snapshot_id'] == snapshot.id + assert parent.data.tags['kumo.task_type'] == 'regression' assert parent.data.tags['kumo.bundle_id'] == result.id stored = tmp_path / 'root' / 'models' / result.id / PROVENANCE_FILE assert json.loads(stored.read_text())['tracking'] == { diff --git a/packages/kumo-engine/test/test_snowflake_tracking.py b/packages/kumo-engine/test/test_snowflake_tracking.py index f2c7786ec..bca68d51c 100644 --- a/packages/kumo-engine/test/test_snowflake_tracking.py +++ b/packages/kumo-engine/test/test_snowflake_tracking.py @@ -169,8 +169,8 @@ def test_metric_and_param_limits(fake, monkeypatch): monkeypatch.setattr(snowflake, '_EPOCH_METRICS', 2) monkeypatch.setattr(snowflake, '_SUMMARY_METRICS', 3) monkeypatch.setattr(snowflake, '_MAX_METRICS', 5) - result = type(RESULT)(metrics={**RESULT.metrics, 'r2': float('nan')}, - id=RESULT.id, bundle_dir=RESULT.bundle_dir) + result = type(RESULT)(**{**vars(RESULT), 'metrics': { + **RESULT.metrics, 'r2': float('nan')}}) _replay(_tracking(tags={'note': 'a$$b', 'level': 2}), result) run, = fake._instance.runs.values() diff --git a/packages/kumo-engine/test/test_tracking_contract.py b/packages/kumo-engine/test/test_tracking_contract.py index 19c8ce2cf..cca628510 100644 --- a/packages/kumo-engine/test/test_tracking_contract.py +++ b/packages/kumo-engine/test/test_tracking_contract.py @@ -7,6 +7,7 @@ import threading import time from dataclasses import dataclass, field, replace +from datetime import datetime, timezone from pathlib import Path from types import SimpleNamespace from typing import Any, Callable, Dict, List, Mapping, Optional, Tuple @@ -35,8 +36,10 @@ 'users': {'connector': 'SNOWFLAKE', 'location': '"DB"."SCH"."USERS"', 'version': '1790000000', 'version_kind': 'timestamp'}, } +SNAPSHOT = f"{'s' * 64}@{'v' * 16}" CONTEXT = FitContext( - query=QUERY, inputs=INPUTS, + query=QUERY, inputs=INPUTS, snapshot_id=SNAPSHOT, task_type='regression', + started=datetime(2026, 10, 29, 6, 18, tzinfo=timezone.utc), model_plan={ 'optimization': {'base_lr': [0.01, 0.001], 'max_epochs': 2}, 'training_job': {'tune_metric': 'mae'}, @@ -82,7 +85,8 @@ def _epochs(trial, *values): SELECTED = {'train_loss': ((1, 0.6), (2, 0.4)), 'val_mae': ((1, 0.5), (2, 0.3))} BUNDLE = 'b' * 64 RESULT = SimpleNamespace(metrics={'val_mae': 0.3, 'test_mae': 0.32}, - id=BUNDLE, bundle_dir=Path('no-bundle')) + id=BUNDLE, bundle_dir=Path('no-bundle'), + registered_model=None) class _Recorder: @@ -273,7 +277,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): @@ -298,7 +302,7 @@ def bucket(key): (step, value)) for prefix, runs in names.items(): observed[prefix or FIT] = Observed( - name=prefix, + name=prefix or run_name, artifacts=run['artifacts'] if prefix is None else {}, **runs) return observed @@ -318,8 +322,10 @@ def test_records_of_the_recorded_fit(): records = replay() fit = records[FIT] + assert fit.name == 'fit-20261029-0618' assert fit.tags == {'team': 'ents', 'kumo.kind': 'fit', - 'kumo.bundle_id': BUNDLE} + 'kumo.bundle_id': BUNDLE, 'kumo.snapshot_id': SNAPSHOT, + 'kumo.task_type': 'regression'} assert fit.params == { 'optimization.base_lr': '[0.01, 0.001]', 'optimization.max_epochs': '2', @@ -335,6 +341,12 @@ def test_records_of_the_recorded_fit(): assert fit.steps == {'search/best_val_mae': ((1, 0.35), (2, 0.3)), 'train/loss': ((1, 0.6), (2, 0.4)), 'val/mae': ((1, 0.5), (2, 0.3))} + registered = SimpleNamespace(**{**vars(RESULT), 'registered_model': + SimpleNamespace(name='ents', version='3')}) + tags = run_record(fitted_spec(fit_spec(CONTEXT), registered, + alias='candidate')).tags + assert tags.items() >= {'kumo.model_name': 'ents', 'kumo.model_version': + '3', 'kumo.model_alias': 'candidate'}.items() assert fit.artifacts['query.pql'] == QUERY assert json.loads(fit.artifacts['model_plan.json']) == CONTEXT.model_plan assert fit.inputs == INPUTS @@ -356,8 +368,7 @@ def test_records_of_the_recorded_fit(): def test_adapter_records_the_same_runs(adapter, tmp_path): (tmp_path / PROVENANCE_FILE).write_text('{"bundle": "b"}') - result = SimpleNamespace(metrics=RESULT.metrics, id=BUNDLE, - bundle_dir=tmp_path) + result = SimpleNamespace(**{**vars(RESULT), 'bundle_dir': tmp_path}) expected = replay(adapter.make(), result=result) assert expected[FIT].artifacts[PROVENANCE_FILE] == '{"bundle": "b"}' From f448e57e1906f62ed4005dc7d2db4501d1a90b0c Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 20:39:08 +0000 Subject: [PATCH 043/278] Tag fit runs with the tune metric, engine version and selected trial Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_tracking.py | 3 ++- packages/kumo-engine/kumo/engine/trainer/_run_spec.py | 6 +++++- packages/kumo-engine/test/test_mlflow_tracking.py | 3 ++- packages/kumo-engine/test/test_tracking_contract.py | 8 +++++++- 4 files changed, 16 insertions(+), 4 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_tracking.py b/packages/kumo-engine/kumo/engine/_tracking.py index 7f0cf98aa..0334d341f 100644 --- a/packages/kumo-engine/kumo/engine/_tracking.py +++ b/packages/kumo-engine/kumo/engine/_tracking.py @@ -443,7 +443,8 @@ def finish(self, spec: Optional[Callable[[], RunSpec]], selected = None if self._search is None else self._search.best if error is None and selected is not None: final = dataclasses.replace( - final, steps=trial_spec(self._trials[selected[1]]).steps) + final, steps=trial_spec(self._trials[selected[1]]).steps, + parents={**final.parents, 'best_trial': selected[1]}) if error is not None and final.error is None: final = dataclasses.replace(final, error=format_error(error)) record: Optional[RunRecord] = run_record(final) diff --git a/packages/kumo-engine/kumo/engine/trainer/_run_spec.py b/packages/kumo-engine/kumo/engine/trainer/_run_spec.py index 17b464ff3..0048053fd 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_run_spec.py +++ b/packages/kumo-engine/kumo/engine/trainer/_run_spec.py @@ -12,6 +12,7 @@ from kumo_core.artifacts import PROVENANCE_FILE from kumo.engine._tracking import RunSpec, _flatten +from kumo.engine._version import __version__ @dataclass(frozen=True) @@ -37,7 +38,10 @@ def fit_spec(context: FitContext, *, started.strftime('%Y%m%d-%H%M')))), params=dict(_flatten(context.model_plan)), parents={'snapshot_id': context.snapshot_id, - 'task_type': context.task_type}, + 'task_type': context.task_type, + 'tune_metric': (context.model_plan.get('training_job') + or {}).get('tune_metric'), + 'engine_version': __version__}, tags=dict(tags or {}), inputs=dict(context.inputs), artifacts={ diff --git a/packages/kumo-engine/test/test_mlflow_tracking.py b/packages/kumo-engine/test/test_mlflow_tracking.py index 21574a3dc..486fe47c9 100644 --- a/packages/kumo-engine/test/test_mlflow_tracking.py +++ b/packages/kumo-engine/test/test_mlflow_tracking.py @@ -117,7 +117,8 @@ def test_fit_logs_a_run_per_fit_and_per_trial(table_and_snapshot, tmp_path, assert parent.info.status == 'FINISHED' assert {k for k in parent.data.tags if k.startswith('kumo.')} == { 'kumo.kind', 'kumo.bundle_id', 'kumo.dropped_params', - 'kumo.snapshot_id', 'kumo.task_type'} + 'kumo.snapshot_id', 'kumo.task_type', 'kumo.tune_metric', + 'kumo.engine_version', 'kumo.best_trial'} assert parent.data.tags['kumo.snapshot_id'] == snapshot.id assert parent.data.tags['kumo.task_type'] == 'regression' assert parent.data.tags['kumo.bundle_id'] == result.id diff --git a/packages/kumo-engine/test/test_tracking_contract.py b/packages/kumo-engine/test/test_tracking_contract.py index cca628510..be8a96e23 100644 --- a/packages/kumo-engine/test/test_tracking_contract.py +++ b/packages/kumo-engine/test/test_tracking_contract.py @@ -21,6 +21,7 @@ TrackingSession, run_record, ) +from kumo.engine._version import __version__ from kumo.engine.tracking import Event, RunRecord from kumo.engine.trainer._run_spec import FitContext, fit_spec, fitted_spec @@ -145,6 +146,8 @@ def replay(tracker=None, events=EVENTS, result=RESULT, error=None, spec = fitted_spec(spec, result, events[-1].metrics.get('seconds')) spec = replace(spec, steps=series if series is not None else { **SEARCH, **({} if error else SELECTED)}) + if series is None and not error: + spec = replace(spec, parents={**spec.parents, 'best_trial': 2}) if 'error' in events[-1].params: spec = replace(spec, error=str(events[-1].params['error'])) records[FIT] = run_record(spec) @@ -325,7 +328,10 @@ def test_records_of_the_recorded_fit(): assert fit.name == 'fit-20261029-0618' assert fit.tags == {'team': 'ents', 'kumo.kind': 'fit', 'kumo.bundle_id': BUNDLE, 'kumo.snapshot_id': SNAPSHOT, - 'kumo.task_type': 'regression'} + 'kumo.task_type': 'regression', + 'kumo.tune_metric': 'mae', + 'kumo.engine_version': __version__, + 'kumo.best_trial': '2'} assert fit.params == { 'optimization.base_lr': '[0.01, 0.001]', 'optimization.max_epochs': '2', From 6a19e2316159793a55111f9755e439bf5693043f Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 20:39:50 +0000 Subject: [PATCH 044/278] Report each epoch's seconds as train_epoch_seconds Co-Authored-By: Claude Opus 5.5 --- .../kumo-core/src/kumo_core/callbacks/train_progress_logger.py | 3 ++- packages/kumo-engine/test/test_fit_events.py | 1 + 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/packages/kumo-core/src/kumo_core/callbacks/train_progress_logger.py b/packages/kumo-core/src/kumo_core/callbacks/train_progress_logger.py index abdbd7803..8d17d417a 100644 --- a/packages/kumo-core/src/kumo_core/callbacks/train_progress_logger.py +++ b/packages/kumo-core/src/kumo_core/callbacks/train_progress_logger.py @@ -197,7 +197,8 @@ def on_train_epoch_end(self, pl_trainer, pl_module): self.on_event('epoch_ended', trial=self.trial, epoch=epoch, metrics={ **metrics.train_metrics, - **metrics.validation_metrics + **metrics.validation_metrics, + 'train_epoch_seconds': metrics.time_taken, }) except Exception: logger.warning('Could not report epoch_ended', diff --git a/packages/kumo-engine/test/test_fit_events.py b/packages/kumo-engine/test/test_fit_events.py index 5b8753b27..196b9e077 100644 --- a/packages/kumo-engine/test/test_fit_events.py +++ b/packages/kumo-engine/test/test_fit_events.py @@ -83,6 +83,7 @@ def broken(event): assert [(e.trial, e.epoch) for e in events if e.name == 'epoch_ended' ] == [(1, 1), (1, 2), (2, 1), (2, 2)] assert any(name.startswith('val_') for name in events[2].metrics) + assert events[2].metrics['train_epoch_seconds'] > 0 assert events[1].params['attempt'] == 1 assert events[4].params == {'status': 'succeeded', 'attempt': 1} assert {'val_mae', 'seconds'} <= set(events[4].metrics) From b701936e8774eaa0707cd69328eb5b4e58217f7a Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 21:12:05 +0000 Subject: [PATCH 045/278] Log data, graph and dataset summaries on the fit run through per-tracker sinks kumo_core.summaries holds small typed results (Scalar, Series, Curve, Histogram, Table, Confusion, Dataset) that know nothing about trackers; each tracker maps the types to its own calls. A fit run gets each split's rows and label balance, the snapshot's node and edge counts, a table of column statistics, and one dataset per source table with its schema and profile. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-core/src/kumo_core/summaries.py | 81 +++++++++++ .../kumo/engine/_data_summaries.py | 137 ++++++++++++++++++ packages/kumo-engine/kumo/engine/_tracking.py | 9 +- .../kumo/engine/integrations/mlflow.py | 117 ++++++++++++++- .../kumo/engine/integrations/snowflake.py | 82 ++++++++++- .../kumo/engine/trainer/_run_spec.py | 12 +- .../kumo/engine/trainer/trainer.py | 16 +- .../kumo-engine/test/test_data_summaries.py | 85 +++++++++++ .../kumo-engine/test/test_mlflow_model.py | 14 +- .../kumo-engine/test/test_mlflow_tracking.py | 19 ++- .../test/test_tracking_contract.py | 73 +++++++++- 11 files changed, 613 insertions(+), 32 deletions(-) create mode 100644 packages/kumo-core/src/kumo_core/summaries.py create mode 100644 packages/kumo-engine/kumo/engine/_data_summaries.py create mode 100644 packages/kumo-engine/test/test_data_summaries.py diff --git a/packages/kumo-core/src/kumo_core/summaries.py b/packages/kumo-core/src/kumo_core/summaries.py new file mode 100644 index 000000000..629a24bd7 --- /dev/null +++ b/packages/kumo-core/src/kumo_core/summaries.py @@ -0,0 +1,81 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""Small typed results of an evaluation or a data summary, named +``
/`` (``test/roc``, ``pred/True_PROB/hist``, ``graph/rows``). + +They know nothing about experiment trackers: each tracker maps these types to +its own calls, so another rendering (a figure, a UI) is one more mapping. +""" + +from dataclasses import dataclass +from typing import Any, Mapping, Optional, Tuple, Union + + +@dataclass(frozen=True) +class Scalar: + name: str + value: float + + +@dataclass(frozen=True) +class Series: + r"""``(step, value)`` points; ``x_label`` says what a step counts.""" + + name: str + points: Tuple[Tuple[int, float], ...] + x_label: str = 'step' + + +@dataclass(frozen=True) +class Curve: + r"""``y`` against ``x``, two series over the same steps, such as ROC's + true-positive against false-positive rate over thresholds.""" + + name: str + x: Series + y: Series + + +@dataclass(frozen=True) +class Histogram: + r"""``counts[i]`` values in ``[edges[i - 1], edges[i])``: ``counts[0]`` + below ``edges[0]`` and ``counts[-1]`` at or above ``edges[-1]``.""" + + name: str + edges: Tuple[float, ...] + counts: Tuple[float, ...] + + +@dataclass(frozen=True) +class Table: + name: str + columns: Tuple[str, ...] + rows: Tuple[Tuple[Any, ...], ...] + + +@dataclass(frozen=True) +class Confusion: + r"""``counts[i][j]`` rows of actual ``labels[i]`` predicted + ``labels[j]``.""" + + name: str + labels: Tuple[str, ...] + counts: Tuple[Tuple[int, ...], ...] + + +@dataclass(frozen=True) +class Dataset: + r"""A source table: ``source`` its ``connector``, ``location``, + ``version`` and ``version_kind``; ``schema`` each column's ``stype`` and + ``dtype``; ``profile`` its rows and column statistics; ``context`` + ``training`` or ``scoring``.""" + + name: str + source: Mapping[str, Optional[str]] + schema: Mapping[str, Mapping[str, str]] + profile: Mapping[str, Any] + context: str + + +Summary = Union[Scalar, Series, Curve, Histogram, Table, Confusion, Dataset] diff --git a/packages/kumo-engine/kumo/engine/_data_summaries.py b/packages/kumo-engine/kumo/engine/_data_summaries.py new file mode 100644 index 000000000..fd115e793 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/_data_summaries.py @@ -0,0 +1,137 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""Summaries of a training table's splits and of a graph snapshot, read +from the statistics their builds wrote (``split_stats.json``, +``col_stats.json``, ``ingest_stats.json``), with no pass over the data.""" + +import json +import logging +from pathlib import Path +from typing import Any, Dict, List, Mapping, Optional + +from kumo_core.instrumentation import format_error +from kumo_core.summaries import Dataset, Scalar, Summary, Table + +logger = logging.getLogger(__name__) + +_SPLITS = ('train', 'val', 'test') +_LABEL_STATS = {'target_mean': 'label_mean', 'target_median': 'label_q50', + 'target_min': 'label_min', 'target_max': 'label_max'} +_COLUMN_STATS = ('count', 'null_rate', 'num_unique', 'min', 'max', 'mean', + 'std', 'q25', 'q50', 'q75', 'avg_tokens', 'max_tokens') +_TOP_CATEGORIES = 10 + + +def split_summaries(staging_dir: Path) -> List[Summary]: + r"""``data//rows`` and ``entities`` per split, and its label + balance: ``positive_rate``, or the label's mean, median, min and max.""" + stats = _read(staging_dir / 'split_stats.json') + summaries: List[Summary] = [] + for split in _SPLITS: + table = stats.get(f'{split}_table_stats') + if not isinstance(table, dict): + continue + prefix = f'data/{split}' + for key, name in (('num_train_rows', 'rows'), + ('num_entities', 'entities')): + if _number(table.get(key)) is not None: + summaries.append(Scalar(f'{prefix}/{name}', table[key])) + target = table.get('target_distribution') or {} + counts = target.get('count') or {} + if sum(counts.values()): + summaries.append(Scalar(f'{prefix}/positive_rate', + counts.get('true', 0) / + sum(counts.values()))) + summaries += [ + Scalar(f'{prefix}/{name}', target[key]) + for key, name in _LABEL_STATS.items() + if _number(target.get(key)) is not None + ] + return summaries + + +def snapshot_summaries(index: Path, context: str) -> List[Summary]: + r"""``graph/rows/``, ``graph/edges/``, ``graph/tables`` and + ``graph/edge_types``; ``graph/col_stats``, a row per column; and a + `Dataset` per source table (its location, version, schema and profile) + in ``context``, ``training`` or ``scoring``.""" + ingest = _read(index / 'ingest_stats.json').get('stats') or {} + columns: Dict[str, Dict[str, Any]] = _read(index / 'col_stats.json') + sources = (_read(index / '_provenance.json').get('parents') + or {}).get('source_tables') or {} + config = _read(index / 'diskgraph_config.json').get('tables') or [] + rows = dict(ingest.get('node_counts') or {}) + edges = {name: count + for name, count in (ingest.get('edge_counts') or {}).items() + if not name.endswith('_reverse')} + summaries: List[Summary] = [ + *(Scalar(f'graph/rows/{table}', count) + for table, count in sorted(rows.items())), + *(Scalar(f'graph/edges/{edge}', count) + for edge, count in sorted(edges.items())), + Scalar('graph/tables', len(rows)), + Scalar('graph/edge_types', len(edges)), + ] + profiles = { + table: {column: _profile(stats) for column, stats in table_stats.items()} + for table, table_stats in sorted(columns.items()) + } + summaries.append(Table('graph/col_stats', + ('table', 'column', 'stype', *_COLUMN_STATS, 'top'), + tuple((table, column, *profile.values()) + for table, table_profiles in profiles.items() + for column, profile in table_profiles.items()))) + for table in config: + name = table.get('kumo_source_table_name') or table.get('name') + if name not in sources: + continue + summaries.append(Dataset( + name=name, source=sources[name], context=context, + schema={column['name']: {'stype': column.get('semantic_type'), + 'dtype': column.get('source_dtype')} + for column in table.get('columns') or ()}, + profile={'rows': rows.get(table.get('name')), + 'columns': profiles.get(table.get('name'), {})})) + return summaries + + +def table_source(manifest: Any) -> List[Summary]: + r"""A `Dataset` for a training table read from its own source table, + which a snapshot's sources do not cover.""" + from kumo.engine._tracking import fit_inputs + + source = fit_inputs(manifest).get('training_table') + return [] if source is None else [ + Dataset('training_table', source, {}, {}, 'training')] + + +def _profile(stats: Mapping[str, Any]) -> Dict[str, Any]: + count = _number(stats.get('count')) + profile: Dict[str, Any] = {'stype': stats.get('semantic_type')} + for name in _COLUMN_STATS: + value = (stats.get('null_count', 0) / count if count else None) \ + if name == 'null_rate' else _number(stats.get(name)) + profile[name] = value + top = stats.get('category_counts') or () + profile['top'] = ', '.join(f'{value}: {n}' + for value, n in top[:_TOP_CATEGORIES]) or None + return profile + + +def _number(value: Any) -> Optional[float]: + if isinstance(value, bool) or not isinstance(value, (int, float)): + return None + return float(value) + + +def _read(path: Path) -> Dict[str, Any]: + try: + value = json.loads(path.read_text()) + except FileNotFoundError: + return {} + except (OSError, ValueError) as error: + logger.warning('Could not read %s for the run summary: %s', path, + format_error(error)) + return {} + return value if isinstance(value, dict) else {} diff --git a/packages/kumo-engine/kumo/engine/_tracking.py b/packages/kumo-engine/kumo/engine/_tracking.py index 0334d341f..959886f7f 100644 --- a/packages/kumo-engine/kumo/engine/_tracking.py +++ b/packages/kumo-engine/kumo/engine/_tracking.py @@ -33,6 +33,7 @@ from kumo.engine._mlops import full_error, warn_failure if TYPE_CHECKING: + from kumo_core.summaries import Summary from kumoapi.plugin import ModelRegistryLocation, PublishedModel from kumo.engine._manifest import ModelManifest, TrackingRef @@ -51,9 +52,10 @@ @dataclass(frozen=True) class RunRecord: r"""One tracked run; ``steps`` are per-epoch ``(epoch, value)`` - histories, ``artifacts`` map file names to text and ``inputs`` map + histories, ``artifacts`` map file names to text, ``inputs`` map source table names to their ``connector``, ``location``, ``version`` and - ``version_kind``.""" + ``version_kind``, and ``summaries`` are `kumo_core.summaries` results, + which each tracker logs in its own way.""" name: Optional[str] = None params: Mapping[str, str] = field(default_factory=dict) @@ -64,6 +66,7 @@ class RunRecord: artifacts: Mapping[str, str] = field(default_factory=dict) inputs: Mapping[str, Mapping[str, Optional[str]]] = field( default_factory=dict) + summaries: Tuple['Summary', ...] = () @dataclass(frozen=True) @@ -83,6 +86,7 @@ class RunSpec: default_factory=dict) parents: Mapping[str, Optional[str]] = field(default_factory=dict) artifacts: Mapping[str, str] = field(default_factory=dict) + summaries: Sequence['Summary'] = () status: Optional[str] = None error: Optional[str] = None @@ -168,6 +172,7 @@ def run_record(spec: RunSpec) -> RunRecord: }, artifacts=dict(spec.artifacts), inputs=dict(spec.inputs), + summaries=tuple(spec.summaries), ) diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index 9b763e9e1..f9300b970 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -27,12 +27,23 @@ Literal, Mapping, Optional, + Sequence, Set, Tuple, ) from urllib.parse import quote from kumo_core.instrumentation import format_error +from kumo_core.summaries import ( + Confusion, + Curve, + Dataset, + Histogram, + Scalar, + Series, + Summary, + Table, +) from kumoapi.plugin import BUNDLE_ID_KEY, ModelRegistryLocation, PublishedModel from kumo.engine._events import Event @@ -340,8 +351,9 @@ def start(self, record: RunRecord) -> None: _SystemMetrics(self._log_system_metrics)) self._log_params(self._parent, record.params) self._log_artifacts(record.artifacts) - with self.target.scope(): - _log_inputs(self._mlflow, client, self._parent, record.inputs) + if record.tags.get('kumo.kind') != 'fit': # a fit logs Datasets + with self.target.scope(): + _log_inputs(self._mlflow, client, self._parent, record.inputs) def on_trial(self, record: RunRecord, event: Event) -> None: key = (event.trial or 0, int(event.params['attempt'])) @@ -415,6 +427,7 @@ def finish(self, record: Optional[RunRecord], model_id=self._model_id) self._log_tags(self._parent, record.tags) self._log_artifacts(record.artifacts) + self.log_summaries(record.summaries) if error is None: status = 'FINISHED' finally: @@ -437,6 +450,86 @@ def _log_system_metrics(self, metrics: Mapping[str, float], step: int) -> None: self._log_metrics(self._parent, metrics, step) + def log_summaries(self, summaries: Sequence[Summary]) -> None: + r"""Logs each summary on the current run through its type's call in + `_SINKS`; a test scalar also goes on the logged model.""" + for summary in summaries: + _SINKS[type(summary)](self, summary) + + def _log_scalar(self, summary: Scalar) -> None: + model_id = self._model_id if summary.name.startswith('test/') else None + self._log_metrics(self._parent, {summary.name: summary.value}, + model_id=model_id) + + def _log_series(self, summary: Series) -> None: + self._log_points({summary.name: summary.points}) + + def _log_curve(self, summary: Curve) -> None: + self._log_points({summary.x.name: summary.x.points, + summary.y.name: summary.y.points}) + + def _log_histogram(self, summary: Histogram) -> None: + self._log_points({ + f'{summary.name}/count': tuple(enumerate(summary.counts)), + f'{summary.name}/lower_edge': tuple(enumerate(summary.edges, 1)), + }) + + def _log_table(self, summary: Table) -> None: + path = f'{summary.name}.json' + self._log_artifacts({path: json.dumps({ + 'columns': list(summary.columns), + 'data': [list(row) for row in summary.rows], + }, default=str)}) + with self.target.scope(): + tags = self._client.get_run(self._parent).data.tags + logged = json.loads(tags.get('mlflow.loggedArtifacts', '[]')) + entry = {'path': path, 'type': 'table'} + if entry not in logged: + self._client.set_tag(self._parent, 'mlflow.loggedArtifacts', + json.dumps([*logged, entry])) + + def _log_confusion(self, summary: Confusion) -> None: + self._log_table(Table( + summary.name, ('actual', *summary.labels), + tuple((label, *row) + for label, row in zip(summary.labels, summary.counts)))) + + def _log_dataset(self, summary: Dataset) -> None: + source = _dataset_source(summary.source) + if source is None: + return + entities = self._mlflow.entities + version = summary.source.get('version') or '' + dataset = entities.Dataset( + name=summary.name, + digest=version[:36] or hashlib.sha256(json.dumps( + dict(summary.source), sort_keys=True).encode()).hexdigest()[:8], + source_type=source._get_source_type(), + source=source.to_json(), + schema=json.dumps({'mlflow_colspec': [ + {'name': name, 'type': _COLUMN_TYPES.get(column['dtype'], + 'string')} + for name, column in summary.schema.items()]}), + profile=json.dumps(summary.profile, default=str)) + with self.target.scope(): + self._client.log_inputs(self._parent, datasets=[ + entities.DatasetInput(dataset, tags=[entities.InputTag( + 'mlflow.data.context', summary.context)])]) + + def _log_points(self, series: Mapping[str, Sequence[Tuple[int, float]]] + ) -> None: + entities = self._mlflow.entities + timestamp = int(time.time() * 1000) + metrics = [ + entities.Metric(metric_key(name), float(value), timestamp, step) + for name, points in series.items() for step, value in points + ] + with self.target.scope(): + for start in range(0, len(metrics), _BATCH_METRICS): + self._client.log_batch( + self._parent, + metrics=metrics[start:start + _BATCH_METRICS]) + def _log_artifacts(self, artifacts: Mapping[str, str]) -> None: for name, text in artifacts.items(): if name in self._artifacts: @@ -449,9 +542,10 @@ def _log_artifacts(self, artifacts: Mapping[str, str]) -> None: self._client.log_text(self._parent, text, name) else: with tempfile.TemporaryDirectory() as tmp: - path = Path(tmp) / name + path = Path(tmp) / posixpath.basename(name) path.write_text(text) - repository.log_artifact(str(path)) + repository.log_artifact( + str(path), posixpath.dirname(name) or None) self._artifacts.add(name) def _log_tags(self, run_id: str, tags: Mapping[str, str]) -> None: @@ -491,6 +585,21 @@ def _log_metrics(self, run_id: Optional[str], ]) +_SINKS: Dict[type, Callable[[MlflowTracking, Any], None]] = { + Scalar: MlflowTracking._log_scalar, + Series: MlflowTracking._log_series, + Curve: MlflowTracking._log_curve, + Histogram: MlflowTracking._log_histogram, + Table: MlflowTracking._log_table, + Confusion: MlflowTracking._log_confusion, + Dataset: MlflowTracking._log_dataset, +} +_BATCH_METRICS = 1000 # MLflow's per-request limit +_COLUMN_TYPES = {'int64': 'long', 'int32': 'integer', 'float64': 'double', + 'float32': 'float', 'bool': 'boolean', 'date': 'datetime', + 'time': 'datetime'} + + class _SystemMetrics: r"""Samples MLflow's CPU, disk, network and GPU monitors every ``interval`` seconds on its own thread and hands each aggregate to diff --git a/packages/kumo-engine/kumo/engine/integrations/snowflake.py b/packages/kumo-engine/kumo/engine/integrations/snowflake.py index ea3c6eeb8..5ffc17a6d 100644 --- a/packages/kumo-engine/kumo/engine/integrations/snowflake.py +++ b/packages/kumo-engine/kumo/engine/integrations/snowflake.py @@ -4,16 +4,38 @@ r"""Snowflake Experiment Tracking for `Trainer.fit` (``pip install 'kumo-engine[snowflake]'``).""" +import json import logging +import posixpath import re import tempfile import uuid from contextlib import ExitStack from pathlib import Path -from typing import Any, Dict, List, Mapping, Optional, Set, Tuple +from typing import ( + Any, + Callable, + Dict, + List, + Mapping, + Optional, + Sequence, + Set, + Tuple, +) from kumo_core.instrumentation import format_error from kumo_core.sources.sql import resolve_snowflake_identifier +from kumo_core.summaries import ( + Confusion, + Curve, + Dataset, + Histogram, + Scalar, + Series, + Summary, + Table, +) from kumo.engine._events import Event from kumo.engine._tracking import ( @@ -203,6 +225,7 @@ def finish(self, record: Optional[RunRecord], for k, v in record.tags.items() if k not in self._start_tags}, _MAX_PARAMS) self._log_artifacts(record.artifacts) + self.log_summaries(record.summaries) self._log_dropped() if error is None: status = 'FINISHED' @@ -248,14 +271,56 @@ def _new_run_name(self, tracking: Any, experiment: Any, name or 'a generated run name') return name + def log_summaries(self, summaries: Sequence[Summary]) -> None: + r"""Logs each summary on the run through its type's call in + `_SINKS`: values as metrics within the name budget, tables as JSON + artifacts; Snowflake has no datasets.""" + if self._run is not None: + for summary in summaries: + _SINKS[type(summary)](self, summary) + + def _log_scalar(self, summary: Scalar) -> None: + self._log_metrics({summary.name: summary.value}, _MAX_METRICS) + + def _log_series(self, summary: Series) -> None: + self._log_points({summary.name: summary.points}) + + def _log_curve(self, summary: Curve) -> None: + self._log_points({summary.x.name: summary.x.points, + summary.y.name: summary.y.points}) + + def _log_histogram(self, summary: Histogram) -> None: + self._log_points({ + f'{summary.name}/count': tuple(enumerate(summary.counts)), + f'{summary.name}/lower_edge': tuple(enumerate(summary.edges, 1)), + }) + + def _log_table(self, summary: Table) -> None: + self._log_artifacts({f'{summary.name}.json': json.dumps({ + 'columns': list(summary.columns), + 'data': [list(row) for row in summary.rows], + }, default=str)}) + + def _log_confusion(self, summary: Confusion) -> None: + self._log_table(Table( + summary.name, ('actual', *summary.labels), + tuple((label, *row) + for label, row in zip(summary.labels, summary.counts)))) + + def _log_points(self, series: Mapping[str, Sequence[Tuple[int, float]]] + ) -> None: + for step, metrics in by_step(series): + self._log_metrics(metrics, _SUMMARY_METRICS, step=step) + def _log_artifacts(self, artifacts: Mapping[str, str]) -> None: with tempfile.TemporaryDirectory() as directory: for name, text in artifacts.items(): if name in self._artifacts: continue - path = Path(directory) / name + path = Path(directory) / posixpath.basename(name) path.write_text(text, encoding='utf-8') - self._tracking.log_artifact(str(path)) + self._tracking.log_artifact( + str(path), posixpath.dirname(name) or None) self._artifacts.add(name) def _log_dropped(self) -> None: @@ -296,6 +361,17 @@ def _log_metrics(self, metrics: Mapping[str, float], limit: int, self._tracking.log_metrics(kept, step=step) +_SINKS: Dict[type, Callable[[SnowflakeTracking, Any], None]] = { + Scalar: SnowflakeTracking._log_scalar, + Series: SnowflakeTracking._log_series, + Curve: SnowflakeTracking._log_curve, + Histogram: SnowflakeTracking._log_histogram, + Table: SnowflakeTracking._log_table, + Confusion: SnowflakeTracking._log_confusion, + Dataset: lambda tracker, summary: None, +} + + def _suffixed(name: str, suffix: str) -> str: if len(name) > 1 and name[0] == name[-1] == '"': return f'{name[:-1]}_{suffix}"' diff --git a/packages/kumo-engine/kumo/engine/trainer/_run_spec.py b/packages/kumo-engine/kumo/engine/trainer/_run_spec.py index 0048053fd..f93bed7c2 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_run_spec.py +++ b/packages/kumo-engine/kumo/engine/trainer/_run_spec.py @@ -7,9 +7,10 @@ from dataclasses import dataclass, field, replace from datetime import datetime, timezone from pathlib import Path -from typing import Any, Mapping, Optional +from typing import Any, Mapping, Optional, Sequence from kumo_core.artifacts import PROVENANCE_FILE +from kumo_core.summaries import Summary from kumo.engine._tracking import RunSpec, _flatten from kumo.engine._version import __version__ @@ -53,11 +54,12 @@ def fit_spec(context: FitContext, *, def fitted_spec(spec: RunSpec, result: Any, seconds: Optional[float] = None, - alias: Optional[str] = None) -> RunSpec: + alias: Optional[str] = None, + summaries: Sequence[Summary] = ()) -> RunSpec: r"""``spec`` once the fit returned ``result``, a `TrainingResult`, after ``seconds``: its test metrics, since the run's train and - validation series show the selected trial's epochs, and the version it - registered with ``alias``.""" + validation series show the selected trial's epochs, the version it + registered with ``alias``, and ``summaries``.""" artifacts = dict(spec.artifacts) provenance = Path(result.bundle_dir) / PROVENANCE_FILE if provenance.is_file(): @@ -71,7 +73,7 @@ def fitted_spec(spec: RunSpec, result: Any, seconds: Optional[float] = None, }) }, parents={**spec.parents, 'bundle_id': result.id, **_registered(result.registered_model, alias)}, - artifacts=artifacts) + artifacts=artifacts, summaries=summaries) def _registered(model: Any, alias: Optional[str]) -> Mapping[str, Any]: diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index e6faf7044..b901b4bf9 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -57,6 +57,11 @@ from typing_extensions import Self from kumo.engine import _deps as dependency +from kumo.engine._data_summaries import ( + snapshot_summaries, + split_summaries, + table_source, +) from kumo.engine._deps import Dependencies from kumo.engine._embeddings import ( no_encoder_inference_warnings, @@ -750,7 +755,7 @@ def fit( events=events, tracking=session, deps=deps, model_name=None if ( settings is None or settings.registry is - None) else settings.registry.name) + None) else settings.registry.model_name) except BaseException as error: events.emit('fit_ended', params={ 'status': 'failed', @@ -774,9 +779,12 @@ def fit( alias = None if settings is None or settings.registry is None else ( settings.registry.alias) session.finish( - lambda: fitted_spec(cast(RunSpec, session.spec), result, seconds, - alias), - error) + lambda: fitted_spec( + cast(RunSpec, session.spec), result, seconds, alias, ( + *snapshot_summaries(Path(snapshot.index_path), + 'training'), + *table_source(train_table.manifest), + *split_summaries(train_table.staging_dir))), error) raise_on_error(settings, session.error or error, 'Trainer.fit (its model is in trainer.result)') return result diff --git a/packages/kumo-engine/test/test_data_summaries.py b/packages/kumo-engine/test/test_data_summaries.py new file mode 100644 index 000000000..71ac424c7 --- /dev/null +++ b/packages/kumo-engine/test/test_data_summaries.py @@ -0,0 +1,85 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +from pathlib import Path + +import numpy as np +import pandas as pd +import pytest +from kumo_core.summaries import Dataset, Scalar, Table + +import kumo.engine as kumo +from kumo.engine._data_summaries import snapshot_summaries, split_summaries +from kumo.engine._scratch import get_scratch_dir +from kumo.engine.pquery import PredictiveQuery + + +@pytest.fixture +def scratch(tmp_path): + previous = get_scratch_dir() + kumo.configure(scratch_dir=tmp_path / 'scratch') + yield tmp_path + kumo.configure(scratch_dir=previous) + + +def test_summaries_match_the_snapshot_and_training_table(scratch): + rng = np.random.default_rng(0) + data = scratch / 'data' + data.mkdir() + users = pd.DataFrame({'USER_ID': np.arange(300), + 'AGE': rng.normal(40, 10, 300)}) + users.loc[0, 'AGE'] = None + orders = pd.DataFrame({ + 'ORDER_ID': np.arange(6000), + 'USER_ID': rng.integers(0, 300, 6000), + 'ORDER_TIME': pd.Timestamp('2024-01-01') + pd.to_timedelta( + rng.integers(0, 120, 6000), unit='D'), + }) + users.to_parquet(data / 'users.parquet') + orders.to_parquet(data / 'orders.parquet') + connector = kumo.FileConnector(root_dir=str(data)) + graph = kumo.Graph( + tables={ + 'users': kumo.Table.from_source_table( + connector['users'], primary_key='USER_ID'), + 'orders': kumo.Table.from_source_table( + connector['orders'], primary_key='ORDER_ID', + time_column='ORDER_TIME').with_column('USER_ID', stype='ID'), + }, + edges=[kumo.Edge(src_table='orders', fkey='USER_ID', + dst_table='users')]) + snapshot = graph.snapshot() + table = PredictiveQuery( + graph, 'PREDICT COUNT(orders.*, 0, 7, days) > 0 FOR EACH users.USER_ID' + ).generate_training_table(snapshot) + + splits = {summary.name: summary.value + for summary in split_summaries(table.staging_dir)} + for split in ('train', 'val', 'test'): + seeds = pd.read_parquet(table.staging_dir / f'{split}_seeds') + assert splits[f'data/{split}/rows'] == len(seeds) + assert splits[f'data/{split}/positive_rate'] == pytest.approx( + seeds['TARGET1'].mean()) + + summaries = snapshot_summaries(Path(snapshot.index_path), 'training') + assert {s.name: s.value for s in summaries if isinstance(s, Scalar)} == { + 'graph/rows/users': 300, 'graph/rows/orders': 6000, + 'graph/edges/orders_USER_ID_to_users_USER_ID': 6000, + 'graph/tables': 2, 'graph/edge_types': 1} + stats, = (s for s in summaries if isinstance(s, Table)) + age = next(dict(zip(stats.columns, row)) for row in stats.rows + if row[:2] == ('users', 'AGE')) + assert (age['stype'], age['count'], age['null_rate'], age['min'], + age['max']) == ('numerical', 300, 1 / 300, users['AGE'].min(), + users['AGE'].max()) + assert age['mean'] == pytest.approx(users['AGE'].mean()) + datasets = {s.name: s for s in summaries if isinstance(s, Dataset)} + assert set(datasets) == {'users', 'orders'} + assert datasets['orders'].schema == { + 'ORDER_ID': {'stype': 'id', 'dtype': 'int64'}, + 'USER_ID': {'stype': 'id', 'dtype': 'int64'}, + 'ORDER_TIME': {'stype': 'timestamp', 'dtype': 'date'}} + assert datasets['orders'].profile['rows'] == 6000 + assert datasets['users'].source['location'] == str( + (data / 'users.parquet').resolve()) + assert datasets['users'].context == 'training' diff --git a/packages/kumo-engine/test/test_mlflow_model.py b/packages/kumo-engine/test/test_mlflow_model.py index 2a0c0ca42..1aaa2a192 100644 --- a/packages/kumo-engine/test/test_mlflow_model.py +++ b/packages/kumo-engine/test/test_mlflow_model.py @@ -20,6 +20,7 @@ import kumo.engine as kumo from kumo.engine._scratch import get_scratch_dir +from kumo.engine._tracking import logged_names mlflow = pytest.importorskip('mlflow') @@ -59,14 +60,15 @@ def test_fit_logs_registers_and_serves_the_model(table_and_snapshot, model, = mlflow.search_logged_models(output_format='list') assert model.source_run_id == parent.info.run_id metrics = {metric.key: metric.value for metric in model.metrics} - assert metrics == {**result.metrics, 'seconds': metrics['seconds']} + assert metrics == { + **{name: value for name, value in logged_names(result.metrics).items() + if name.startswith('test/')}, 'seconds': metrics['seconds']} version = mlflow.MlflowClient().get_model_version_by_alias('m', 'champion') assert version.source == model.model_uri - assert {k: v for k, v in parent.data.tags.items() - if k.startswith('kumo.')} == { - 'kumo.kind': 'fit', - 'kumo.bundle_id': result.id, - } + assert parent.data.tags.items() >= { + 'kumo.kind': 'fit', 'kumo.bundle_id': result.id, + 'kumo.model_name': 'm', 'kumo.model_version': str(version.version), + 'kumo.model_alias': 'champion'}.items() assert result.registered_model == PublishedModel( name='m', version=str(version.version), uri=f'models:/m/{version.version}', diff --git a/packages/kumo-engine/test/test_mlflow_tracking.py b/packages/kumo-engine/test/test_mlflow_tracking.py index 486fe47c9..d2c260f75 100644 --- a/packages/kumo-engine/test/test_mlflow_tracking.py +++ b/packages/kumo-engine/test/test_mlflow_tracking.py @@ -15,9 +15,15 @@ import pandas as pd import pytest from kumo_core.artifacts import PROVENANCE_FILE +from kumo_core.summaries import Dataset from kumoapi.encoder import Numerical from kumoapi.model_plan import ModelPlan -from test_tracking_contract import CONTEXT, join_tracking_threads, replay +from test_tracking_contract import ( + CONTEXT, + INPUTS, + join_tracking_threads, + replay, +) import kumo.engine as kumo from kumo.engine._scratch import get_scratch_dir @@ -140,9 +146,10 @@ def test_fit_logs_a_run_per_fit_and_per_trial(table_and_snapshot, tmp_path, } tested = {name: value for name, value in logged_names( result.metrics).items() if name.startswith('test/')} - assert set(metrics) == {*tested, 'seconds', 'search/best_val_mae', - *(name for name in metrics - if name.startswith(('train/', 'val/')))} + assert {name for name in metrics if not name.startswith(( + 'train/', 'val/', 'data/', 'graph/'))} == { + *tested, 'seconds', 'search/best_val_mae'} + assert metrics['graph/rows/ents'] == 400 assert {name: metrics[name] for name in tested} == tested assert 'off for the rest of this run' not in caplog.text ents, = parent.inputs.dataset_inputs @@ -360,7 +367,9 @@ def test_source_tables_are_logged_as_native_inputs( monkeypatch.setattr(DeltaDatasetSource, '_lookup_table_id', lambda self, name: 'table-id') - replay(MlflowTracking(log_system_metrics=False)) + replay(MlflowTracking(log_system_metrics=False), summaries=[ + Dataset(name, source, {}, {}, 'training') + for name, source in INPUTS.items()]) run, = [run for run in _runs(mlflow.MlflowClient()) if run.data.tags['kumo.kind'] == 'fit'] diff --git a/packages/kumo-engine/test/test_tracking_contract.py b/packages/kumo-engine/test/test_tracking_contract.py index be8a96e23..9fa6363e3 100644 --- a/packages/kumo-engine/test/test_tracking_contract.py +++ b/packages/kumo-engine/test/test_tracking_contract.py @@ -14,6 +14,15 @@ import pytest from kumo_core.artifacts import PROVENANCE_FILE +from kumo_core.summaries import ( + Confusion, + Curve, + Dataset, + Histogram, + Scalar, + Series, + Table, +) from kumo.engine._tracking import ( DROPPED_METRICS, @@ -126,9 +135,11 @@ def join_tracking_threads(): def replay(tracker=None, events=EVENTS, result=RESULT, error=None, - context=CONTEXT, series=None) -> Dict[str, RunRecord]: + context=CONTEXT, series=None, + summaries=()) -> Dict[str, RunRecord]: r"""Feeds the recorded fit to `tracker`; returns the records expected - of it, with the fit run's ``series`` (default: the search's).""" + of it, with the fit run's ``series`` (default: the search's) and + ``summaries``.""" recorder = _Recorder(tracker) session = TrackingSession(recorder) session.start(lambda: fit_spec(context)) @@ -136,7 +147,8 @@ def replay(tracker=None, events=EVENTS, result=RESULT, error=None, session.on_event(event) session.finish( None if result is None else lambda: fitted_spec( - session.spec, result, events[-1].metrics.get('seconds')), error) + session.spec, result, events[-1].metrics.get('seconds'), + summaries=summaries), error) records = { record.name: record for hook, record in recorder.calls if hook == 'on_trial' @@ -255,6 +267,7 @@ def history(run_id, key): a.path: mlflow.artifacts.load_text( f'{run.info.artifact_uri}/{a.path}') for a in client.list_artifacts(run.info.run_id) + if not a.is_dir }) for run in runs } @@ -394,6 +407,60 @@ def test_adapter_records_a_single_trial_fit_epoch_by_epoch(adapter): assert gaps(expected, adapter.observe(), adapter) == [] +SUMMARIES = ( + Scalar('test/brier', 0.25), + Series('search/trials', ((1, 3.0), (2, 4.0)), x_label='trial'), + Curve('test/roc', Series('test/roc/fpr', ((0, 1.0), (1, 0.5))), + Series('test/roc/tpr', ((0, 1.0), (1, 0.75)))), + Histogram('pred/True_PROB/hist', (0.5,), (3, 4)), + Table('graph/col_stats', ('table', 'rows'), (('ents', 400),)), + Confusion('test/confusion', ('a', 'b'), ((5, 1), (2, 7))), + Dataset('ents', INPUTS['ents'], {'ID': {'stype': 'id', 'dtype': 'int64'}}, + {'rows': 400}, 'training'), +) + + +def test_adapter_logs_each_summary_type(adapter, mlflow_adapter): + replay(adapter.make(), summaries=SUMMARIES) + + fit = adapter.observe()[FIT] + assert {name: fit.steps[name] for name in ( + 'test/brier', 'search/trials', 'test/roc/fpr', 'test/roc/tpr', + 'pred/True_PROB/hist/count', 'pred/True_PROB/hist/lower_edge')} == { + 'test/brier': ((0, 0.25),), 'search/trials': ((1, 3.0), (2, 4.0)), + 'test/roc/fpr': ((0, 1.0), (1, 0.5)), + 'test/roc/tpr': ((0, 1.0), (1, 0.75)), + 'pred/True_PROB/hist/count': ((0, 3.0), (1, 4.0)), + 'pred/True_PROB/hist/lower_edge': ((1, 0.5),)} + tables = {'col_stats.json': {'columns': ['table', 'rows'], + 'data': [['ents', 400]]}, + 'confusion.json': {'columns': ['actual', 'a', 'b'], + 'data': [['a', 5, 1], ['b', 2, 7]]}} + if adapter is not mlflow_adapter: + assert {name: json.loads(fit.artifacts[name]) + for name in tables} == tables + return + import mlflow + + client = mlflow.MlflowClient() + run, = (r for r in client.search_runs( + [e.experiment_id for e in client.search_experiments()]) + if r.data.tags['kumo.kind'] == 'fit') + assert {name: json.loads(mlflow.artifacts.load_text( + f'{run.info.artifact_uri}/{section}/{name}')) + for section, name in (('graph', 'col_stats.json'), + ('test', 'confusion.json'))} == tables + assert json.loads(run.data.tags['mlflow.loggedArtifacts']) == [ + {'path': 'graph/col_stats.json', 'type': 'table'}, + {'path': 'test/confusion.json', 'type': 'table'}] + dataset, = run.inputs.dataset_inputs + assert (dataset.dataset.name, dataset.dataset.digest, + dataset.tags[0].value) == ('ents', '1790000000', 'training') + assert json.loads(dataset.dataset.schema) == { + 'mlflow_colspec': [{'name': 'ID', 'type': 'long'}]} + assert json.loads(dataset.dataset.profile) == {'rows': 400} + + def test_a_raising_tracker_is_ended_and_not_called_again(caplog): class RaisesOnTrial(_Recorder): def on_trial(self, record, event): From e45527a01f5a6c6443c444baa0951097cb327e30 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 21:17:03 +0000 Subject: [PATCH 046/278] Log each table's and relation's data quality, and flag Kumo's data checks graph/
/{duplicate_pks,dropped_rows,null_time_dropped} and graph/edges//{count,src_match_frac,dst_match_frac, fk_lookup_failures} come from the snapshot's ingest and edge stats. The kumo.data_warnings tag names each table that breaks a Kumo platform check: 5% missing or invalid time, invalid or duplicate primary keys, 10% invalid values, a match below 20% in either direction, fewer than 10 rows. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-core/src/kumo_core/summaries.py | 9 +- .../kumo/engine/_data_summaries.py | 200 ++++++++++++++---- .../kumo/engine/integrations/mlflow.py | 5 + .../kumo/engine/integrations/snowflake.py | 5 + .../kumo/engine/trainer/trainer.py | 2 +- .../kumo-engine/test/test_data_summaries.py | 67 +++++- 6 files changed, 245 insertions(+), 43 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/summaries.py b/packages/kumo-core/src/kumo_core/summaries.py index 629a24bd7..2d7906f58 100644 --- a/packages/kumo-core/src/kumo_core/summaries.py +++ b/packages/kumo-core/src/kumo_core/summaries.py @@ -18,6 +18,12 @@ class Scalar: value: float +@dataclass(frozen=True) +class Tag: + name: str + value: str + + @dataclass(frozen=True) class Series: r"""``(step, value)`` points; ``x_label`` says what a step counts.""" @@ -78,4 +84,5 @@ class Dataset: context: str -Summary = Union[Scalar, Series, Curve, Histogram, Table, Confusion, Dataset] +Summary = Union[Scalar, Tag, Series, Curve, Histogram, Table, Confusion, + Dataset] diff --git a/packages/kumo-engine/kumo/engine/_data_summaries.py b/packages/kumo-engine/kumo/engine/_data_summaries.py index fd115e793..bc5fd7158 100644 --- a/packages/kumo-engine/kumo/engine/_data_summaries.py +++ b/packages/kumo-engine/kumo/engine/_data_summaries.py @@ -11,7 +11,7 @@ from typing import Any, Dict, List, Mapping, Optional from kumo_core.instrumentation import format_error -from kumo_core.summaries import Dataset, Scalar, Summary, Table +from kumo_core.summaries import Dataset, Scalar, Summary, Table, Tag logger = logging.getLogger(__name__) @@ -21,6 +21,14 @@ _COLUMN_STATS = ('count', 'null_rate', 'num_unique', 'min', 'max', 'mean', 'std', 'q25', 'q50', 'q75', 'avg_tokens', 'max_tokens') _TOP_CATEGORIES = 10 +_EDGE_STATS = {'FRAC_SRC_NONZERO_DEGREE': 'src_match_frac', + 'FRAC_DST_NONZERO_DEGREE': 'dst_match_frac'} +# The Kumo platform's data checks. +_BAD_TIME_FRACTION = 0.05 +_INVALID_FRACTION = 0.10 +_MIN_MATCH_FRACTION = 0.20 +_MIN_ROWS = 10 +_TAG_LENGTH = 5000 def split_summaries(staging_dir: Path) -> List[Summary]: @@ -51,51 +59,170 @@ def split_summaries(staging_dir: Path) -> List[Summary]: return summaries -def snapshot_summaries(index: Path, context: str) -> List[Summary]: - r"""``graph/rows/
``, ``graph/edges/``, ``graph/tables`` and - ``graph/edge_types``; ``graph/col_stats``, a row per column; and a +def snapshot_summaries(index: Path, graph_config: Mapping[str, Any], + context: str) -> List[Summary]: + r"""The snapshot's tables and relations (`_graph_scalars`), their data + quality (`_quality`), ``graph/col_stats`` with a row per column, and a `Dataset` per source table (its location, version, schema and profile) in ``context``, ``training`` or ``scoring``.""" ingest = _read(index / 'ingest_stats.json').get('stats') or {} columns: Dict[str, Dict[str, Any]] = _read(index / 'col_stats.json') - sources = (_read(index / '_provenance.json').get('parents') - or {}).get('source_tables') or {} - config = _read(index / 'diskgraph_config.json').get('tables') or [] - rows = dict(ingest.get('node_counts') or {}) - edges = {name: count - for name, count in (ingest.get('edge_counts') or {}).items() - if not name.endswith('_reverse')} - summaries: List[Summary] = [ - *(Scalar(f'graph/rows/{table}', count) - for table, count in sorted(rows.items())), - *(Scalar(f'graph/edges/{edge}', count) - for edge, count in sorted(edges.items())), - Scalar('graph/tables', len(rows)), - Scalar('graph/edge_types', len(edges)), - ] + relations = _relations(_read(index / 'edge_stats.json'), ingest) profiles = { table: {column: _profile(stats) for column, stats in table_stats.items()} for table, table_stats in sorted(columns.items()) } - summaries.append(Table('graph/col_stats', - ('table', 'column', 'stype', *_COLUMN_STATS, 'top'), - tuple((table, column, *profile.values()) - for table, table_profiles in profiles.items() - for column, profile in table_profiles.items()))) - for table in config: - name = table.get('kumo_source_table_name') or table.get('name') - if name not in sources: - continue - summaries.append(Dataset( - name=name, source=sources[name], context=context, - schema={column['name']: {'stype': column.get('semantic_type'), - 'dtype': column.get('source_dtype')} - for column in table.get('columns') or ()}, - profile={'rows': rows.get(table.get('name')), - 'columns': profiles.get(table.get('name'), {})})) + return [ + *_graph_scalars(ingest, relations), + *_quality(ingest, columns, relations, graph_config), + Table('graph/col_stats', + ('table', 'column', 'stype', *_COLUMN_STATS, 'top'), + tuple((table, column, *profile.values()) + for table, table_profiles in profiles.items() + for column, profile in table_profiles.items())), + *_datasets(index, ingest, profiles, context), + ] + + +def _relations(edge_stats: Mapping[str, Any], + ingest: Mapping[str, Any]) -> Dict[str, Dict[str, float]]: + r"""Each relation's ``count``, ``src_match_frac`` and ``dst_match_frac`` + (the share of source and destination rows with an edge) and + ``fk_lookup_failures``, under DiskGraph's ``__to__`` + name.""" + relations: Dict[str, Dict[str, float]] = { + name: {'count': count} + for name, count in (ingest.get('edge_counts') or {}).items() + if not name.endswith('_reverse') + } + for key, entry in edge_stats.items(): + parts = key.split('__') + if len(parts) == 7 and isinstance(entry, dict): + src, _, dst, _, fkey, pkey, stat = parts + name = _EDGE_STATS.get(stat) + relation = relations.get(f'{src}_{fkey}_to_{dst}_{pkey}') + if name and relation is not None: + relation[name] = entry.get('value') + failures = ingest.get('key_failures_by_context') or {} + for name, relation in relations.items(): + missed = failures.get(f'dst_lookup:{name}') or {} + relation['fk_lookup_failures'] = (missed.get('lookup_failures', 0) + + missed.get('membership_failures', 0)) + return relations + + +def _graph_scalars(ingest: Mapping[str, Any], + relations: Mapping[str, Mapping[str, float]] + ) -> List[Summary]: + rows = ingest.get('node_counts') or {} + return [ + *(Scalar(f'graph/rows/{table}', count) + for table, count in sorted(rows.items())), + *(Scalar(f'graph/edges/{relation}/{name}', value) + for relation, stats in sorted(relations.items()) + for name, value in stats.items() if _number(value) is not None), + Scalar('graph/tables', len(rows)), + Scalar('graph/edge_types', len(relations)), + ] + + +def _quality(ingest: Mapping[str, Any], columns: Mapping[str, Any], + relations: Mapping[str, Mapping[str, float]], + graph_config: Mapping[str, Any]) -> List[Summary]: + r"""``graph/
/{duplicate_pks,dropped_rows,null_time_dropped}`` and, + when any table breaks the Kumo platform's data checks, the tag + ``kumo.data_warnings`` naming each.""" + duplicates = _by_table(ingest.get('pk_duplicate_counts'), ':') + null_time = _by_table(ingest.get('pk_null_timestamp_drops'), ':shard-') + dropped = ingest.get('dropped_rows_counts') or {} + sources = ingest.get('source_rows_by_table') or {} + failures = ingest.get('key_failures_by_context') or {} + summaries: List[Summary] = [] + warnings: List[str] = [] + for table in graph_config.get('tables') or (): + name = table['table_name'] + counts = {'duplicate_pks': duplicates.get(name, 0), + 'dropped_rows': dropped.get(name, 0), + 'null_time_dropped': null_time.get(name, 0)} + summaries += [Scalar(f'graph/{name}/{key}', value) + for key, value in counts.items()] + stats = columns.get(name) or {} + nodes = (ingest.get('node_counts') or {}).get(name) + rows = sources.get(name) or nodes + pkey = (table.get('pkey') or {}).get('name') + bad_pks = (failures.get(f'pkey:{name}.{pkey}') + or {}).get('parse_failures', 0) + time = (table.get('time_col') or {}).get('name') + time_stats = stats.get(time) or {} + bad_times = (counts['null_time_dropped'] + + time_stats.get('null_count', 0) + + time_stats.get('parse_fail_count', 0)) + invalid = [column for column, column_stats in stats.items() + if column_stats.get('count') and _invalid(column_stats) >= + _INVALID_FRACTION * column_stats['count']] + warnings += [ + *([f'{name}: {bad_times / rows:.1%} missing or invalid time'] + if time and rows and bad_times >= _BAD_TIME_FRACTION * rows + else []), + *([f'{name}: {bad_pks} missing or invalid primary keys'] + if bad_pks else []), + *([f"{name}: {counts['duplicate_pks']} duplicate primary keys"] + if counts['duplicate_pks'] else []), + *(f'{name}.{column}: at least {_INVALID_FRACTION:.0%} invalid ' + f'values' for column in sorted(invalid)), + *([f'{name}: {nodes} rows'] if nodes is not None + and nodes < _MIN_ROWS else []), + ] + for fkey in table.get('fkeys') or (): + relation = relations.get( + f"{name}_{fkey['name']}_to_{fkey['dst_table']}_" + f"{fkey['dst_key']}") or {} + for side, table_name in (('src', name), + ('dst', fkey['dst_table'])): + match = _number(relation.get(f'{side}_match_frac')) + if match is not None and match < _MIN_MATCH_FRACTION: + warnings.append( + f"{name}.{fkey['name']} -> {fkey['dst_table']}: " + f"{match:.1%} of {table_name} rows match") + if warnings: + summaries.append(Tag('kumo.data_warnings', + '; '.join(warnings)[:_TAG_LENGTH])) return summaries +def _datasets(index: Path, ingest: Mapping[str, Any], + profiles: Mapping[str, Mapping[str, Any]], + context: str) -> List[Summary]: + sources = (_read(index / '_provenance.json').get('parents') + or {}).get('source_tables') or {} + rows = ingest.get('node_counts') or {} + datasets: List[Summary] = [] + for table in _read(index / 'diskgraph_config.json').get('tables') or (): + name = table.get('kumo_source_table_name') or table.get('name') + if name in sources: + datasets.append(Dataset( + name=name, source=sources[name], context=context, + schema={column['name']: {'stype': column.get('semantic_type'), + 'dtype': column.get('source_dtype')} + for column in table.get('columns') or ()}, + profile={'rows': rows.get(table.get('name')), + 'columns': profiles.get(table.get('name'), {})})) + return datasets + + +def _by_table(counts: Optional[Mapping[str, int]], + separator: str) -> Dict[str, int]: + tables: Dict[str, int] = {} + for key, count in (counts or {}).items(): + table = key.split(separator, 1)[0] + tables[table] = tables.get(table, 0) + count + return tables + + +def _invalid(stats: Mapping[str, Any]) -> int: + return stats.get('parse_fail_count', 0) + stats.get('invalid_count', 0) + + def table_source(manifest: Any) -> List[Summary]: r"""A `Dataset` for a training table read from its own source table, which a snapshot's sources do not cover.""" @@ -113,7 +240,8 @@ def _profile(stats: Mapping[str, Any]) -> Dict[str, Any]: value = (stats.get('null_count', 0) / count if count else None) \ if name == 'null_rate' else _number(stats.get(name)) profile[name] = value - top = stats.get('category_counts') or () + top = () if stats.get('semantic_type') == 'id' else ( + stats.get('category_counts') or ()) profile['top'] = ', '.join(f'{value}: {n}' for value, n in top[:_TOP_CATEGORIES]) or None return profile diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index f9300b970..5bcc71a0e 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -43,6 +43,7 @@ Series, Summary, Table, + Tag, ) from kumoapi.plugin import BUNDLE_ID_KEY, ModelRegistryLocation, PublishedModel @@ -461,6 +462,9 @@ def _log_scalar(self, summary: Scalar) -> None: self._log_metrics(self._parent, {summary.name: summary.value}, model_id=model_id) + def _log_tag(self, summary: Tag) -> None: + self._log_tags(self._parent, {summary.name: summary.value}) + def _log_series(self, summary: Series) -> None: self._log_points({summary.name: summary.points}) @@ -587,6 +591,7 @@ def _log_metrics(self, run_id: Optional[str], _SINKS: Dict[type, Callable[[MlflowTracking, Any], None]] = { Scalar: MlflowTracking._log_scalar, + Tag: MlflowTracking._log_tag, Series: MlflowTracking._log_series, Curve: MlflowTracking._log_curve, Histogram: MlflowTracking._log_histogram, diff --git a/packages/kumo-engine/kumo/engine/integrations/snowflake.py b/packages/kumo-engine/kumo/engine/integrations/snowflake.py index 5ffc17a6d..55a5e3192 100644 --- a/packages/kumo-engine/kumo/engine/integrations/snowflake.py +++ b/packages/kumo-engine/kumo/engine/integrations/snowflake.py @@ -35,6 +35,7 @@ Series, Summary, Table, + Tag, ) from kumo.engine._events import Event @@ -282,6 +283,9 @@ def log_summaries(self, summaries: Sequence[Summary]) -> None: def _log_scalar(self, summary: Scalar) -> None: self._log_metrics({summary.name: summary.value}, _MAX_METRICS) + def _log_tag(self, summary: Tag) -> None: + self._log_params({summary.name: summary.value}) + def _log_series(self, summary: Series) -> None: self._log_points({summary.name: summary.points}) @@ -363,6 +367,7 @@ def _log_metrics(self, metrics: Mapping[str, float], limit: int, _SINKS: Dict[type, Callable[[SnowflakeTracking, Any], None]] = { Scalar: SnowflakeTracking._log_scalar, + Tag: SnowflakeTracking._log_tag, Series: SnowflakeTracking._log_series, Curve: SnowflakeTracking._log_curve, Histogram: SnowflakeTracking._log_histogram, diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index b901b4bf9..7d305ff04 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -782,7 +782,7 @@ def fit( lambda: fitted_spec( cast(RunSpec, session.spec), result, seconds, alias, ( *snapshot_summaries(Path(snapshot.index_path), - 'training'), + snapshot.graph_config, 'training'), *table_source(train_table.manifest), *split_summaries(train_table.staging_dir))), error) raise_on_error(settings, session.error or error, diff --git a/packages/kumo-engine/test/test_data_summaries.py b/packages/kumo-engine/test/test_data_summaries.py index 71ac424c7..d6fcc8033 100644 --- a/packages/kumo-engine/test/test_data_summaries.py +++ b/packages/kumo-engine/test/test_data_summaries.py @@ -1,12 +1,13 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. # All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import json from pathlib import Path import numpy as np import pandas as pd import pytest -from kumo_core.summaries import Dataset, Scalar, Table +from kumo_core.summaries import Dataset, Scalar, Table, Tag import kumo.engine as kumo from kumo.engine._data_summaries import snapshot_summaries, split_summaries @@ -31,7 +32,7 @@ def test_summaries_match_the_snapshot_and_training_table(scratch): users.loc[0, 'AGE'] = None orders = pd.DataFrame({ 'ORDER_ID': np.arange(6000), - 'USER_ID': rng.integers(0, 300, 6000), + 'USER_ID': rng.integers(0, 360, 6000), 'ORDER_TIME': pd.Timestamp('2024-01-01') + pd.to_timedelta( rng.integers(0, 120, 6000), unit='D'), }) @@ -61,11 +62,22 @@ def test_summaries_match_the_snapshot_and_training_table(scratch): assert splits[f'data/{split}/positive_rate'] == pytest.approx( seeds['TARGET1'].mean()) - summaries = snapshot_summaries(Path(snapshot.index_path), 'training') - assert {s.name: s.value for s in summaries if isinstance(s, Scalar)} == { + summaries = snapshot_summaries(Path(snapshot.index_path), + snapshot.graph_config, 'training') + matched = orders['USER_ID'] < 300 + relation = 'graph/edges/orders_USER_ID_to_users_USER_ID' + scalars = {s.name: s.value for s in summaries if isinstance(s, Scalar)} + assert scalars == { 'graph/rows/users': 300, 'graph/rows/orders': 6000, - 'graph/edges/orders_USER_ID_to_users_USER_ID': 6000, + f'{relation}/count': matched.sum(), + f'{relation}/src_match_frac': pytest.approx(matched.mean()), + f'{relation}/dst_match_frac': pytest.approx( + orders.loc[matched, 'USER_ID'].nunique() / 300), + f'{relation}/fk_lookup_failures': (~matched).sum(), + **{f'graph/{table}/{name}': 0 for table in ('users', 'orders') + for name in ('duplicate_pks', 'dropped_rows', 'null_time_dropped')}, 'graph/tables': 2, 'graph/edge_types': 1} + assert not [s for s in summaries if isinstance(s, Tag)] stats, = (s for s in summaries if isinstance(s, Table)) age = next(dict(zip(stats.columns, row)) for row in stats.rows if row[:2] == ('users', 'AGE')) @@ -83,3 +95,48 @@ def test_summaries_match_the_snapshot_and_training_table(scratch): assert datasets['users'].source['location'] == str( (data / 'users.parquet').resolve()) assert datasets['users'].context == 'training' + + +def test_data_warnings_name_each_table_that_breaks_a_check(tmp_path): + def write(name, value): + (tmp_path / name).write_text(json.dumps(value)) + + write('ingest_stats.json', {'stats': { + 'node_counts': {'users': 9, 'orders': 100}, + 'source_rows_by_table': {'users': 12, 'orders': 100}, + 'edge_counts': {'orders_USER_ID_to_users_USER_ID': 10}, + 'pk_duplicate_counts': {'users:shard-0': 2}, + 'pk_null_timestamp_drops': {'orders:shard-0': 3}, + 'dropped_rows_counts': {'users': 1, 'orders': 3}, + 'key_failures_by_context': {'pkey:users.USER_ID': { + 'parse_failures': 1}}, + }}) + write('col_stats.json', {'orders': { + 'ORDER_TIME': {'semantic_type': 'timestamp', 'count': 97, + 'null_count': 2, 'parse_fail_count': 0}, + 'AMOUNT': {'semantic_type': 'numerical', 'count': 100, + 'null_count': 0, 'parse_fail_count': 4, + 'invalid_count': 6}, + }}) + write('edge_stats.json', { + f'orders__c__users__c__USER_ID__USER_ID__{stat}': {'value': value} + for stat, value in (('FRAC_SRC_NONZERO_DEGREE', 0.1), + ('FRAC_DST_NONZERO_DEGREE', 0.5))}) + config = {'tables': [ + {'table_name': 'users', 'pkey': {'name': 'USER_ID'}, 'fkeys': []}, + {'table_name': 'orders', 'pkey': {'name': 'ORDER_ID'}, + 'time_col': {'name': 'ORDER_TIME'}, + 'fkeys': [{'name': 'USER_ID', 'dst_table': 'users', + 'dst_key': 'USER_ID'}]}, + ]} + + tag, = (s for s in snapshot_summaries(tmp_path, config, 'training') + if isinstance(s, Tag)) + assert tag == Tag('kumo.data_warnings', '; '.join([ + 'users: 1 missing or invalid primary keys', + 'users: 2 duplicate primary keys', + 'users: 9 rows', + 'orders: 5.0% missing or invalid time', + 'orders.AMOUNT: at least 10% invalid values', + 'orders.USER_ID -> users: 10.0% of orders rows match', + ])) From 39d43e3c665b0b6bdfe81660b02fe5f6a94242cd Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 21:19:04 +0000 Subject: [PATCH 047/278] Name Snowflake fit runs as plain identifiers snowflake-ml puts a run's name unquoted into its artifact URIs, so a fit run named with a quoted identifier could not log its artifacts. A fit run is named fit___ unless run_name is given, and the contract compares names in each adapter's form. Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/integrations/snowflake.py | 26 ++++++++++++------- .../test/test_tracking_contract.py | 12 ++++++--- 2 files changed, 26 insertions(+), 12 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/integrations/snowflake.py b/packages/kumo-engine/kumo/engine/integrations/snowflake.py index 55a5e3192..6d915db55 100644 --- a/packages/kumo-engine/kumo/engine/integrations/snowflake.py +++ b/packages/kumo-engine/kumo/engine/integrations/snowflake.py @@ -75,7 +75,7 @@ class SnowflakeTracking: schema_name: `None` uses the session's current schema, or ``PUBLIC`` when `database` is given or the session has none. run_name: name of the fit run. `None` names it after its registry - model and UTC start, as ``"fit-USERS_CHURN-20261027-0618"``; a + model and UTC start, as ``fit_users_churn_20261027_0618``; a taken name gets a random suffix. tags: extra tags for the fit run. native_model_lineage: keep the run open as a context for the whole @@ -183,7 +183,9 @@ def start(self, record: RunRecord) -> None: database_name=database, schema_name=schema or 'PUBLIC') self._run = tracking.start_run( - self._new_run_name(tracking, experiment.name, record.name)) + self._new_run_name( + tracking, experiment.name, + self.run_name or _run_identifier(record.name or 'fit'))) self.run_id = str(self._run.name) if self.native_model_lineage: self._run_context.enter_context(self._run) @@ -246,9 +248,8 @@ def log_metrics(self, metrics: Mapping[str, float], step: int) -> None: self._log_metrics(metrics, _EPOCH_METRICS, step=step) def _new_run_name(self, tracking: Any, experiment: Any, - default: Optional[str]) -> Optional[str]: - wanted = self.run_name or (default and f'"{default}"') - if wanted is None: + run_name: Optional[str]) -> Optional[str]: + if run_name is None: return None # start_run(name) resumes, or raises on, any run matching name as a # SQL LIKE pattern and cannot create an existing run, so pick a name @@ -259,16 +260,16 @@ def _new_run_name(self, tracking: Any, experiment: Any, experiment_name=experiment) ] short = uuid.uuid4().hex - for name in (wanted, _suffixed(wanted, short[:8]), - _suffixed(wanted, short)): + for name in (run_name, _suffixed(run_name, short[:8]), + _suffixed(run_name, short)): if not _taken(name, existing): break else: name = None - if name != wanted: + if name != run_name: logger.warning( 'Snowflake run %s already exists in experiment %s; recording ' - 'this fit as %s', wanted, self.experiment_name, + 'this fit as %s', run_name, self.experiment_name, name or 'a generated run name') return name @@ -377,6 +378,13 @@ def _log_metrics(self, metrics: Mapping[str, float], limit: int, } +def _run_identifier(name: str) -> str: + r"""``name`` as an unquoted identifier: snowflake-ml puts a run's name + into its artifact URIs, where a quoted name is not found.""" + name = re.sub(r'[^A-Za-z0-9_]', '_', name) + return name if re.match(r'[A-Za-z_]', name) else f'_{name}' + + def _suffixed(name: str, suffix: str) -> str: if len(name) > 1 and name[0] == name[-1] == '"': return f'{name[:-1]}_{suffix}"' diff --git a/packages/kumo-engine/test/test_tracking_contract.py b/packages/kumo-engine/test/test_tracking_contract.py index 9fa6363e3..579092c53 100644 --- a/packages/kumo-engine/test/test_tracking_contract.py +++ b/packages/kumo-engine/test/test_tracking_contract.py @@ -185,6 +185,7 @@ class Adapter: nested_runs: bool = False step_metrics: bool = False artifacts: bool = False + run_name: Callable[[str], str] = str def gaps(expected: Dict[str, RunRecord], observed: Dict[str, Observed], @@ -207,7 +208,8 @@ def differ(what, have, want): found.append(f'{run_id} {what}: {have!r} != {want!r}') if record.name is not None: - differ('name', seen.name, record.name) + differ('name', adapter.run_name(seen.name), + adapter.run_name(record.name)) differ('params', dict(seen.params), dict(kept_params(record.params))) differ('metrics', {k: seen.metrics.get(k) for k in record.metrics}, dict(record.metrics)) @@ -289,7 +291,10 @@ def snowflake_adapter(monkeypatch): import _snowflake_experiment_fake fake = _snowflake_experiment_fake.install(monkeypatch) - from kumo.engine.integrations.snowflake import SnowflakeTracking + from kumo.engine.integrations.snowflake import ( + SnowflakeTracking, + _run_identifier, + ) def observe(): observed = {} @@ -326,7 +331,8 @@ def bucket(key): return Adapter( make=lambda: SnowflakeTracking(_snowflake_experiment_fake.FakeSession(), experiment_name='ents', tags=USER_TAGS), - observe=observe, step_metrics=True, artifacts=True) + observe=observe, step_metrics=True, artifacts=True, + run_name=lambda name: _run_identifier(name).upper()) @pytest.fixture(params=['mlflow', 'snowflake']) From 2dfbf93537596eedbdc3d3fbcb11bd80d8b5f256 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 21:52:39 +0000 Subject: [PATCH 048/278] Keep system metrics going past a failed log One failed log skips that reading instead of ending system metrics for the rest of the fit. The first failure is logged once at INFO with its cause; six failures in a row stop them. Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/integrations/mlflow.py | 17 +++++++--- .../kumo-engine/test/test_mlflow_tracking.py | 34 +++++++++++++++++++ 2 files changed, 47 insertions(+), 4 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index 5bcc71a0e..4a2f51761 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -600,6 +600,7 @@ def _log_metrics(self, run_id: Optional[str], Dataset: MlflowTracking._log_dataset, } _BATCH_METRICS = 1000 # MLflow's per-request limit +_SYSTEM_METRICS_FAILURES = 6 _COLUMN_TYPES = {'int64': 'long', 'int32': 'integer', 'float64': 'double', 'float32': 'float', 'bool': 'boolean', 'date': 'datetime', 'time': 'datetime'} @@ -617,6 +618,7 @@ def __init__(self, log: Callable[[Mapping[str, float], int], None], self._log = log self._interval = interval self._monitors = _metric_monitors() + self._failed = False self._stop = threading.Event() self._thread = threading.Thread(target=self._run, daemon=True, name='kumo-system-metrics') @@ -632,7 +634,7 @@ def __exit__(self, *exc: Any) -> None: self._thread.join(self._interval + 5) def _run(self) -> None: - step = 0 + step = failures = 0 while not self._stop.wait(self._interval): metrics: Dict[str, float] = {} for monitor in self._monitors: @@ -642,10 +644,17 @@ def _run(self) -> None: try: self._log({f'system/{k}': v for k, v in metrics.items()}, step) + failures = 0 except Exception as error: - logger.debug('Could not log system metrics: %s', - format_error(error)) - return + if not self._failed: + logger.info('Could not log system metrics, skipping that ' + 'reading: %s', format_error(error)) + self._failed = True + failures += 1 + if failures == _SYSTEM_METRICS_FAILURES: + logger.info('System metrics stopped after %d failed logs ' + 'in a row', failures) + return step += 1 diff --git a/packages/kumo-engine/test/test_mlflow_tracking.py b/packages/kumo-engine/test/test_mlflow_tracking.py index d2c260f75..b11f9db3f 100644 --- a/packages/kumo-engine/test/test_mlflow_tracking.py +++ b/packages/kumo-engine/test/test_mlflow_tracking.py @@ -448,3 +448,37 @@ def test_link_prediction_metric_names_are_logged(tracking_uri): 'test_map_at_10': 0.5, 'val_precision_at_10': 0.25 } + + +def test_system_metrics_skip_a_failed_log_and_stop_after_repeated_ones( + monkeypatch, caplog): + import logging + + from kumo.engine.integrations import mlflow as integration + + class Monitor: + def collect_metrics(self): + pass + + def aggregate_metrics(self): + return {'cpu': 1.0} + + def clear_metrics(self): + pass + + calls = [] + + def log(metrics, step): + calls.append(step) + if step == 0 or step >= 2: + raise ConnectionError('blip') + + monkeypatch.setattr(integration, '_metric_monitors', lambda: [Monitor()]) + caplog.set_level(logging.INFO, integration.logger.name) + with integration._SystemMetrics(log, interval=0.001) as metrics: + metrics._thread.join(10) + + assert calls == list(range(2 + integration._SYSTEM_METRICS_FAILURES)) + assert caplog.text.count('Could not log system metrics') == 1 + assert 'ConnectionError: blip' in caplog.text + assert 'stopped after 6 failed logs in a row' in caplog.text From 47f517298fbb1bdd018f358fecad3b555ff2b2e9 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 21:37:57 +0000 Subject: [PATCH 049/278] Record in each manifest where the artifact root stores it MLOps.locate writes the location into the local manifest before a snapshot, training table or model is copied to the root, so the stored copy and every child that embeds it carry it; assets record theirs too. Ids and the bundle id do not change. An asset dependency names its artifact. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_deps.py | 10 +++++---- packages/kumo-engine/kumo/engine/_manifest.py | 7 ++++++ packages/kumo-engine/kumo/engine/_mlops.py | 17 ++++++++++++-- .../kumo-engine/kumo/engine/_snapshot_io.py | 6 +++++ .../kumo/engine/pquery/training_table.py | 2 ++ .../kumo/engine/trainer/trainer.py | 5 +++++ packages/kumo-engine/test/test_artifacts.py | 22 +++++++++++++++++++ 7 files changed, 63 insertions(+), 6 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_deps.py b/packages/kumo-engine/kumo/engine/_deps.py index f8d6b5713..23b76117b 100644 --- a/packages/kumo-engine/kumo/engine/_deps.py +++ b/packages/kumo-engine/kumo/engine/_deps.py @@ -235,17 +235,19 @@ def _store(asset: Asset, sha256: str, settings: Any) -> Optional[str]: with open(path, 'rb') as reader, \ filesystem.open_output_stream(f'{root}/{name}') as writer: shutil.copyfileobj(reader, writer, 1 << 20) + location = settings.root.uri(relative) with filesystem.open_output_stream(marker) as writer: - writer.write(manifest_bytes(_manifest(asset, sha256))) - return settings.root.uri(relative) + writer.write(manifest_bytes(_manifest(asset, sha256, location))) + return location except Exception as error: raise RuntimeError( f"Could not store {asset.asset_kind} {asset.name!r} under the " f"artifact root: {error}") from error -def _manifest(asset: Asset, sha256: str) -> AssetManifest: - return AssetManifest(id=sha256, asset_kind=asset.asset_kind, +def _manifest(asset: Asset, sha256: str, location: str) -> AssetManifest: + return AssetManifest(id=sha256, location=location, + asset_kind=asset.asset_kind, name=asset.name, version=asset.version, format=entry_point(asset), source=str(getattr(asset, 'path', '') or '') or None, diff --git a/packages/kumo-engine/kumo/engine/_manifest.py b/packages/kumo-engine/kumo/engine/_manifest.py index df1fd4582..2f0a7e92e 100644 --- a/packages/kumo-engine/kumo/engine/_manifest.py +++ b/packages/kumo-engine/kumo/engine/_manifest.py @@ -340,6 +340,13 @@ class Dependency(Record): packing: Optional[str] = omit_none() location: Optional[str] = omit_none() + def as_artifact(self, key: str) -> Optional[Artifact]: + kind = _kinds().get(self.kind) + return None if kind is None else Artifact( + self.kind, str(self.sha256), self.location, kind.IS_DATASET, + f'kumo.{self.asset_kind}_uri', version=self.sha256, + version_kind='exact') + @dataclass(frozen=True, kw_only=True) class ModelManifest(Manifest): diff --git a/packages/kumo-engine/kumo/engine/_mlops.py b/packages/kumo-engine/kumo/engine/_mlops.py index fb6e501fe..f2977cacb 100644 --- a/packages/kumo-engine/kumo/engine/_mlops.py +++ b/packages/kumo-engine/kumo/engine/_mlops.py @@ -9,7 +9,8 @@ import posixpath import re import warnings -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace +from pathlib import Path from typing import ( Any, FrozenSet, @@ -18,12 +19,13 @@ Literal, Optional, Tuple, + Type, Union, ) from kumoapi.plugin import ModelRegistryLocation, check_model_alias -from kumo.engine._manifest import KINDS +from kumo.engine._manifest import KINDS, M, read_manifest, write_manifest @dataclass(frozen=True) @@ -157,6 +159,17 @@ def wants(self, kind: str) -> bool: r"""Whether artifacts of ``kind`` are stored under the root.""" return kind in self.kinds + def locate(self, path: Path, kind: Type[M], relative: str) -> Optional[M]: + r"""Records in the manifest at ``path`` that its artifact is stored + at ``relative`` under the root, before it is copied there; returns + the manifest, if there is one.""" + location = self.root.uri(relative) + manifest = read_manifest(path, kind) + if manifest is not None and manifest.location != location: + manifest = replace(manifest, location=location) + write_manifest(path, manifest) + return manifest + def _require_within_root(connector: Any, path: str) -> None: from kumo_core.sources.object_store_sessions import ObjectStoreSessions diff --git a/packages/kumo-engine/kumo/engine/_snapshot_io.py b/packages/kumo-engine/kumo/engine/_snapshot_io.py index cce308d7f..863b5dbbd 100644 --- a/packages/kumo-engine/kumo/engine/_snapshot_io.py +++ b/packages/kumo-engine/kumo/engine/_snapshot_io.py @@ -195,7 +195,13 @@ def export_location(snapshot: Any, path: str) -> str: def store(snapshot: Any, settings: Any) -> None: r"""Export ``snapshot`` under ``/snapshots/`` unless a complete export of it is there.""" + from kumo.engine._manifest import SnapshotManifest + from kumo.engine._mlops import snapshot_path + from kumo.engine.graph import SNAPSHOT_MANIFEST_FILE + try: + settings.locate(Path(snapshot.index_path) / SNAPSHOT_MANIFEST_FILE, + SnapshotManifest, snapshot_path(snapshot)) connector = settings.root.connector path = settings.root.location('snapshots') filesystem, root = _location(connector, diff --git a/packages/kumo-engine/kumo/engine/pquery/training_table.py b/packages/kumo-engine/kumo/engine/pquery/training_table.py index 82d616e11..098ca0ebe 100644 --- a/packages/kumo-engine/kumo/engine/pquery/training_table.py +++ b/packages/kumo-engine/kumo/engine/pquery/training_table.py @@ -142,6 +142,8 @@ def _store(self, settings: Any) -> None: path = settings.root.location(f'training_tables/{self.id}/') try: + settings.locate(self.staging_dir / STAGING_MANIFEST_FILE, + TrainingTableManifest, f'training_tables/{self.id}') filesystem, root = file_store(settings.root.connector, path) if filesystem.get_file_info( f'{root}/_SUCCESS').type is not FileType.File: diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 7d305ff04..c6b1d15b6 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -402,14 +402,19 @@ def _store_model(bundle_dir: Path, manifest: Optional[ModelManifest], settings: Any) -> None: r"""Copy the bundle to ``/models//`` unless a complete copy is there.""" + from kumo_core.artifacts import PROVENANCE_FILE, add_bundle_file from kumo_core.object_store_sink import holds_bundle, upload_bundle from pyarrow.fs import FileSelector from kumo.engine.connector import file_store + from kumo.engine.trainer._spec import PROVENANCE_ROLE if manifest is None: return try: + if settings.locate(bundle_dir / PROVENANCE_FILE, ModelManifest, + f'models/{manifest.id}') is not None: + add_bundle_file(bundle_dir, PROVENANCE_ROLE, PROVENANCE_FILE) filesystem, root = file_store( settings.root.connector, settings.root.location(f'models/{manifest.id}')) diff --git a/packages/kumo-engine/test/test_artifacts.py b/packages/kumo-engine/test/test_artifacts.py index 0d51cc262..54a2663a1 100644 --- a/packages/kumo-engine/test/test_artifacts.py +++ b/packages/kumo-engine/test/test_artifacts.py @@ -5,6 +5,9 @@ from dataclasses import dataclass from typing import List, Optional +from kumo_core.artifacts import bundle_id_from_manifest + +import kumo.engine as kumo from kumo.engine._manifest import ( Artifact, Manifest, @@ -16,7 +19,10 @@ SourceTableEntry, TrainingTableManifest, artifacts, + read_manifest, ) +from kumo.engine._snapshot_io import store +from kumo.engine.trainer.trainer import _store_and_register @dataclass(frozen=True, kw_only=True) @@ -60,3 +66,19 @@ def test_a_new_kind_names_its_artifacts_without_other_changes(): Artifact('snapshot', 'p', None, True, 'kumo.snapshot_uri', version='p', version_kind='exact'), ] + + +def test_each_store_records_its_location_and_keeps_the_ids(fitted, tmp_path): + trainer, _, snapshot = fitted + bundle = trainer.result.bundle_dir + ids = snapshot.id, trainer.result.id, bundle_id_from_manifest(bundle) + for root in (tmp_path / 'a', tmp_path / 'b'): + settings = kumo.MLOps(str(root)) + store(snapshot, settings) + _store_and_register(bundle, trainer.result.manifest, settings) + assert snapshot.manifest.location == settings.snapshot_uri(snapshot) + assert trainer.result.manifest.location == f'{root}/models/{ids[1]}' + assert read_manifest(root / 'models' / ids[1] / '_provenance.json', + ModelManifest).location == f'{root}/models/{ids[1]}' + assert (snapshot.id, trainer.result.id, + bundle_id_from_manifest(bundle)) == ids From 6278a75f19edd464737ccd92b18672d0e368ceb8 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 21:40:09 +0000 Subject: [PATCH 050/278] Give tracked runs their artifacts' locations as tags and datasets artifact_summaries turns a manifest's artifacts into kumo._uri tags (kumo.bundle_uri for a model, the nearest of each), kumo.artifact_root, and a Dataset per stored dataset with its real location, for the trackers' sinks to log. A new kind or asset needs no tracker change. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_manifest.py | 44 ++++++++++++++++++ packages/kumo-engine/test/test_artifacts.py | 46 +++++++++++++++++++ 2 files changed, 90 insertions(+) diff --git a/packages/kumo-engine/kumo/engine/_manifest.py b/packages/kumo-engine/kumo/engine/_manifest.py index 2f0a7e92e..5888e5404 100644 --- a/packages/kumo-engine/kumo/engine/_manifest.py +++ b/packages/kumo-engine/kumo/engine/_manifest.py @@ -18,6 +18,7 @@ List, Mapping, Optional, + Tuple, Type, TypeVar, Union, @@ -27,6 +28,7 @@ ) from kumo_core.io.filesystem import atomic_write_bytes +from kumo_core.summaries import Dataset, Summary, Tag SCHEMA_VERSION = 1 _OMIT_NONE = 'omit_none' @@ -214,6 +216,48 @@ def _walk(value: Any, key: str, root: Manifest) -> Iterator[Artifact]: yield from _walk(item, key, root) +#: Connector kinds of the artifact root's URI schemes. +_SCHEMES = {'s3': 'S3', 'gs': 'GCS', 'file': 'LOCAL_FILE', '': 'LOCAL_FILE'} + +Details = Mapping[str, Tuple[Mapping[str, Any], Mapping[str, Any]]] + + +def artifact_summaries(manifest: Manifest, context: str, + settings: Any = None, + details: Optional[Details] = None) -> List[Summary]: + r"""Where ``manifest``'s artifacts are, for its tracked run: each + artifact's tag (the nearest of each), ``kumo.artifact_root`` under MLOps + ``settings``, and a `Dataset` per stored dataset, read in ``context`` + or written, with the schema and profile ``details`` gives its id.""" + located = [item for item in artifacts(manifest) if item.location] + tags: Dict[str, str] = {} + for item in located: + if item.tag: + tags.setdefault(item.tag, str(item.location)) + if settings is not None: + tags['kumo.artifact_root'] = settings.root.uri('').rstrip('/') + summaries: List[Summary] = [Tag(name, value) + for name, value in tags.items()] + for item in located: + if not item.is_dataset: + continue + location = str(item.location) + source = {'connector': item.connector or _connector(location), + 'location': location, 'version': item.version, + 'version_kind': item.version_kind} + schema, profile = (details or {}).get(item.id, ({}, {})) + summaries.append(Dataset(item.id, source, schema, profile, + item.context or context)) + return summaries + + +def _connector(location: str) -> Optional[str]: + from urllib.parse import urlparse + + return None if location.startswith('@') else _SCHEMES.get( + urlparse(location).scheme) + + def stamp() -> Dict[str, Any]: """Envelope fields for a manifest written now.""" from kumo.engine._provenance import environment diff --git a/packages/kumo-engine/test/test_artifacts.py b/packages/kumo-engine/test/test_artifacts.py index 54a2663a1..1fc12bed2 100644 --- a/packages/kumo-engine/test/test_artifacts.py +++ b/packages/kumo-engine/test/test_artifacts.py @@ -6,10 +6,13 @@ from typing import List, Optional from kumo_core.artifacts import bundle_id_from_manifest +from kumo_core.summaries import Dataset, Tag import kumo.engine as kumo +from kumo.engine._deps import Dependencies, FileAsset from kumo.engine._manifest import ( Artifact, + Dependency, Manifest, ModelManifest, ModelParents, @@ -18,6 +21,7 @@ SnapshotParents, SourceTableEntry, TrainingTableManifest, + artifact_summaries, artifacts, read_manifest, ) @@ -36,6 +40,7 @@ class ReportManifest(Manifest): kind: str = 'report' parents: ReportParents pages: Optional[List[SnapshotManifest]] = None + dependencies: Optional[List[Dependency]] = None IS_DATASET = True @@ -82,3 +87,44 @@ def test_each_store_records_its_location_and_keeps_the_ids(fitted, tmp_path): ModelManifest).location == f'{root}/models/{ids[1]}' assert (snapshot.id, trainer.result.id, bundle_id_from_manifest(bundle)) == ids + + +def test_a_new_kind_and_asset_reach_tags_and_datasets(tmp_path): + settings = kumo.MLOps(str(tmp_path / 'root')) + (tmp_path / 'notes.txt').write_text('notes') + asset = FileAsset('dummy', 'notes', tmp_path / 'notes.txt') + dependency = Dependencies(settings=settings).require(asset) + root = f"{tmp_path / 'root'}" + orders = SourceTableEntry(connector='S3', location='s3://data/orders', + version='7', version_kind='exact') + report = ReportManifest( + id='r', location=f'{root}/reports/r', + parents=ReportParents( + snapshot=SnapshotManifest( + id='g@d', location=f'{root}/snapshots/g/d', + parents=SnapshotParents(source_tables={'orders': orders})), + model=ModelManifest(id='m', location=f'{root}/models/m')), + dependencies=[dependency]) + + assert dependency.location == f'{root}/assets/dummy/{dependency.sha256}' + assert ReportManifest.from_dict(report.to_dict()) == report + assert artifact_summaries(report, 'training', settings, { + 'orders': ({'ID': {'stype': 'ID', 'dtype': 'int64'}}, {'rows': 3}) + }) == [ + Tag('kumo.report_uri', f'{root}/reports/r'), + Tag('kumo.snapshot_uri', f'{root}/snapshots/g/d'), + Tag('kumo.bundle_uri', f'{root}/models/m'), + Tag('kumo.dummy_uri', dependency.location), + Tag('kumo.artifact_root', root), + Dataset('r', {'connector': 'LOCAL_FILE', + 'location': f'{root}/reports/r', 'version': 'r', + 'version_kind': 'exact'}, {}, {}, 'training'), + Dataset('g@d', {'connector': 'LOCAL_FILE', + 'location': f'{root}/snapshots/g/d', + 'version': 'g@d', 'version_kind': 'exact'}, {}, {}, + 'training'), + Dataset('orders', {'connector': 'S3', 'location': 's3://data/orders', + 'version': '7', 'version_kind': 'exact'}, + {'ID': {'stype': 'ID', 'dtype': 'int64'}}, {'rows': 3}, + 'training'), + ] From 73432ecc688527ff68e0c3ff0dcdaa856f5e299e Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 21:48:05 +0000 Subject: [PATCH 051/278] Log a fit run's datasets from its model manifest The fit run's tags and datasets come from artifact_summaries over the model it stored: the training table, snapshot and source tables at their real locations, with each source table's schema and profile from table_details. snapshot_summaries keeps the graph and quality summaries. Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/_data_summaries.py | 69 ++++++++----------- .../kumo/engine/trainer/trainer.py | 16 +++-- .../kumo-engine/test/test_data_summaries.py | 16 +++-- .../kumo-engine/test/test_mlflow_tracking.py | 5 +- 4 files changed, 58 insertions(+), 48 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_data_summaries.py b/packages/kumo-engine/kumo/engine/_data_summaries.py index bc5fd7158..2e2fe052e 100644 --- a/packages/kumo-engine/kumo/engine/_data_summaries.py +++ b/packages/kumo-engine/kumo/engine/_data_summaries.py @@ -11,7 +11,9 @@ from typing import Any, Dict, List, Mapping, Optional from kumo_core.instrumentation import format_error -from kumo_core.summaries import Dataset, Scalar, Summary, Table, Tag +from kumo_core.summaries import Scalar, Summary, Table, Tag + +from kumo.engine._manifest import Details logger = logging.getLogger(__name__) @@ -59,19 +61,14 @@ def split_summaries(staging_dir: Path) -> List[Summary]: return summaries -def snapshot_summaries(index: Path, graph_config: Mapping[str, Any], - context: str) -> List[Summary]: +def snapshot_summaries(index: Path, + graph_config: Mapping[str, Any]) -> List[Summary]: r"""The snapshot's tables and relations (`_graph_scalars`), their data - quality (`_quality`), ``graph/col_stats`` with a row per column, and a - `Dataset` per source table (its location, version, schema and profile) - in ``context``, ``training`` or ``scoring``.""" + quality (`_quality`) and ``graph/col_stats`` with a row per column.""" ingest = _read(index / 'ingest_stats.json').get('stats') or {} columns: Dict[str, Dict[str, Any]] = _read(index / 'col_stats.json') relations = _relations(_read(index / 'edge_stats.json'), ingest) - profiles = { - table: {column: _profile(stats) for column, stats in table_stats.items()} - for table, table_stats in sorted(columns.items()) - } + profiles = _profiles(columns) return [ *_graph_scalars(ingest, relations), *_quality(ingest, columns, relations, graph_config), @@ -80,10 +77,17 @@ def snapshot_summaries(index: Path, graph_config: Mapping[str, Any], tuple((table, column, *profile.values()) for table, table_profiles in profiles.items() for column, profile in table_profiles.items())), - *_datasets(index, ingest, profiles, context), ] +def _profiles(columns: Mapping[str, Mapping[str, Any]] + ) -> Dict[str, Dict[str, Any]]: + return { + table: {column: _profile(stats) for column, stats in table_stats.items()} + for table, table_stats in sorted(columns.items()) + } + + def _relations(edge_stats: Mapping[str, Any], ingest: Mapping[str, Any]) -> Dict[str, Dict[str, float]]: r"""Each relation's ``count``, ``src_match_frac`` and ``dst_match_frac`` @@ -190,24 +194,21 @@ def _quality(ingest: Mapping[str, Any], columns: Mapping[str, Any], return summaries -def _datasets(index: Path, ingest: Mapping[str, Any], - profiles: Mapping[str, Mapping[str, Any]], - context: str) -> List[Summary]: - sources = (_read(index / '_provenance.json').get('parents') - or {}).get('source_tables') or {} - rows = ingest.get('node_counts') or {} - datasets: List[Summary] = [] - for table in _read(index / 'diskgraph_config.json').get('tables') or (): - name = table.get('kumo_source_table_name') or table.get('name') - if name in sources: - datasets.append(Dataset( - name=name, source=sources[name], context=context, - schema={column['name']: {'stype': column.get('semantic_type'), - 'dtype': column.get('source_dtype')} - for column in table.get('columns') or ()}, - profile={'rows': rows.get(table.get('name')), - 'columns': profiles.get(table.get('name'), {})})) - return datasets +def table_details(index: Path) -> Details: + r"""Each source table's schema and profile, its rows and column + statistics, for its `Dataset`.""" + rows = (_read(index / 'ingest_stats.json').get('stats') + or {}).get('node_counts') or {} + profiles = _profiles(_read(index / 'col_stats.json')) + return { + table.get('kumo_source_table_name') or table.get('name'): ( + {column['name']: {'stype': column.get('semantic_type'), + 'dtype': column.get('source_dtype')} + for column in table.get('columns') or ()}, + {'rows': rows.get(table.get('name')), + 'columns': profiles.get(table.get('name'), {})}) + for table in _read(index / 'diskgraph_config.json').get('tables') or () + } def _by_table(counts: Optional[Mapping[str, int]], @@ -223,16 +224,6 @@ def _invalid(stats: Mapping[str, Any]) -> int: return stats.get('parse_fail_count', 0) + stats.get('invalid_count', 0) -def table_source(manifest: Any) -> List[Summary]: - r"""A `Dataset` for a training table read from its own source table, - which a snapshot's sources do not cover.""" - from kumo.engine._tracking import fit_inputs - - source = fit_inputs(manifest).get('training_table') - return [] if source is None else [ - Dataset('training_table', source, {}, {}, 'training')] - - def _profile(stats: Mapping[str, Any]) -> Dict[str, Any]: count = _number(stats.get('count')) profile: Dict[str, Any] = {'stype': stats.get('semantic_type')} diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index c6b1d15b6..c2eb8e3c3 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -60,7 +60,7 @@ from kumo.engine._data_summaries import ( snapshot_summaries, split_summaries, - table_source, + table_details, ) from kumo.engine._deps import Dependencies from kumo.engine._embeddings import ( @@ -71,7 +71,13 @@ with_text_encoders, ) from kumo.engine._events import Event, EventEmitter, query_sha256 -from kumo.engine._manifest import Dependency, ModelManifest, TrackingRef, stub +from kumo.engine._manifest import ( + Dependency, + ModelManifest, + TrackingRef, + artifact_summaries, + stub, +) from kumo.engine._mlops import full_error, warn_failure from kumo.engine._predictor import TRAINER_PREDICTOR from kumo.engine._scratch import new_run_id, trainer_scratch_dir @@ -787,8 +793,10 @@ def fit( lambda: fitted_spec( cast(RunSpec, session.spec), result, seconds, alias, ( *snapshot_summaries(Path(snapshot.index_path), - snapshot.graph_config, 'training'), - *table_source(train_table.manifest), + snapshot.graph_config), + *([] if result.manifest is None else artifact_summaries( + result.manifest, 'training', settings, + table_details(Path(snapshot.index_path)))), *split_summaries(train_table.staging_dir))), error) raise_on_error(settings, session.error or error, 'Trainer.fit (its model is in trainer.result)') diff --git a/packages/kumo-engine/test/test_data_summaries.py b/packages/kumo-engine/test/test_data_summaries.py index d6fcc8033..ce68db8f5 100644 --- a/packages/kumo-engine/test/test_data_summaries.py +++ b/packages/kumo-engine/test/test_data_summaries.py @@ -10,7 +10,12 @@ from kumo_core.summaries import Dataset, Scalar, Table, Tag import kumo.engine as kumo -from kumo.engine._data_summaries import snapshot_summaries, split_summaries +from kumo.engine._data_summaries import ( + snapshot_summaries, + split_summaries, + table_details, +) +from kumo.engine._manifest import artifact_summaries from kumo.engine._scratch import get_scratch_dir from kumo.engine.pquery import PredictiveQuery @@ -63,7 +68,7 @@ def test_summaries_match_the_snapshot_and_training_table(scratch): seeds['TARGET1'].mean()) summaries = snapshot_summaries(Path(snapshot.index_path), - snapshot.graph_config, 'training') + snapshot.graph_config) matched = orders['USER_ID'] < 300 relation = 'graph/edges/orders_USER_ID_to_users_USER_ID' scalars = {s.name: s.value for s in summaries if isinstance(s, Scalar)} @@ -85,7 +90,10 @@ def test_summaries_match_the_snapshot_and_training_table(scratch): age['max']) == ('numerical', 300, 1 / 300, users['AGE'].min(), users['AGE'].max()) assert age['mean'] == pytest.approx(users['AGE'].mean()) - datasets = {s.name: s for s in summaries if isinstance(s, Dataset)} + datasets = {s.name: s for s in artifact_summaries( + snapshot.manifest, 'training', + details=table_details(Path(snapshot.index_path))) + if isinstance(s, Dataset)} assert set(datasets) == {'users', 'orders'} assert datasets['orders'].schema == { 'ORDER_ID': {'stype': 'id', 'dtype': 'int64'}, @@ -130,7 +138,7 @@ def write(name, value): 'dst_key': 'USER_ID'}]}, ]} - tag, = (s for s in snapshot_summaries(tmp_path, config, 'training') + tag, = (s for s in snapshot_summaries(tmp_path, config) if isinstance(s, Tag)) assert tag == Tag('kumo.data_warnings', '; '.join([ 'users: 1 missing or invalid primary keys', diff --git a/packages/kumo-engine/test/test_mlflow_tracking.py b/packages/kumo-engine/test/test_mlflow_tracking.py index b11f9db3f..710356183 100644 --- a/packages/kumo-engine/test/test_mlflow_tracking.py +++ b/packages/kumo-engine/test/test_mlflow_tracking.py @@ -124,7 +124,10 @@ def test_fit_logs_a_run_per_fit_and_per_trial(table_and_snapshot, tmp_path, assert {k for k in parent.data.tags if k.startswith('kumo.')} == { 'kumo.kind', 'kumo.bundle_id', 'kumo.dropped_params', 'kumo.snapshot_id', 'kumo.task_type', 'kumo.tune_metric', - 'kumo.engine_version', 'kumo.best_trial'} + 'kumo.engine_version', 'kumo.best_trial', 'kumo.bundle_uri', + 'kumo.artifact_root'} + assert parent.data.tags['kumo.bundle_uri'] == str( + tmp_path / 'root' / 'models' / result.id) assert parent.data.tags['kumo.snapshot_id'] == snapshot.id assert parent.data.tags['kumo.task_type'] == 'regression' assert parent.data.tags['kumo.bundle_id'] == result.id From 2ceabc7539e93936da4fbe7d7b0160bcea6e1faa Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 23:53:26 +0000 Subject: [PATCH 052/278] Log a last system-metrics reading when a run ends A run shorter than the sampling interval logged none. Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/integrations/mlflow.py | 4 +- .../kumo-engine/test/test_mlflow_tracking.py | 38 +++++++++++++------ 2 files changed, 29 insertions(+), 13 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index 4a2f51761..4a6d7ddb1 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -635,7 +635,9 @@ def __exit__(self, *exc: Any) -> None: def _run(self) -> None: step = failures = 0 - while not self._stop.wait(self._interval): + stopped = False + while not stopped: # and a last reading when the run ends + stopped = self._stop.wait(self._interval) metrics: Dict[str, float] = {} for monitor in self._monitors: monitor.collect_metrics() diff --git a/packages/kumo-engine/test/test_mlflow_tracking.py b/packages/kumo-engine/test/test_mlflow_tracking.py index 710356183..526726d0e 100644 --- a/packages/kumo-engine/test/test_mlflow_tracking.py +++ b/packages/kumo-engine/test/test_mlflow_tracking.py @@ -436,7 +436,7 @@ def test_link_prediction_metric_names_are_logged(tracking_uri): from kumo.engine._tracking import RunRecord from kumo.engine.integrations.mlflow import MlflowTracking - tracker = MlflowTracking(experiment_name='lp') + tracker = MlflowTracking(experiment_name='lp', log_system_metrics=False) tags = {'kumo.kind': 'fit', 'kumo.run_id': 'r'} tracker.start(RunRecord(name='r', tags=tags)) tracker.finish( @@ -453,22 +453,23 @@ def test_link_prediction_metric_names_are_logged(tracking_uri): } +class _Monitor: + def collect_metrics(self): + pass + + def aggregate_metrics(self): + return {'cpu': 1.0} + + def clear_metrics(self): + pass + + def test_system_metrics_skip_a_failed_log_and_stop_after_repeated_ones( monkeypatch, caplog): import logging from kumo.engine.integrations import mlflow as integration - class Monitor: - def collect_metrics(self): - pass - - def aggregate_metrics(self): - return {'cpu': 1.0} - - def clear_metrics(self): - pass - calls = [] def log(metrics, step): @@ -476,7 +477,7 @@ def log(metrics, step): if step == 0 or step >= 2: raise ConnectionError('blip') - monkeypatch.setattr(integration, '_metric_monitors', lambda: [Monitor()]) + monkeypatch.setattr(integration, '_metric_monitors', lambda: [_Monitor()]) caplog.set_level(logging.INFO, integration.logger.name) with integration._SystemMetrics(log, interval=0.001) as metrics: metrics._thread.join(10) @@ -485,3 +486,16 @@ def log(metrics, step): assert caplog.text.count('Could not log system metrics') == 1 assert 'ConnectionError: blip' in caplog.text assert 'stopped after 6 failed logs in a row' in caplog.text + + +def test_system_metrics_log_a_reading_when_a_run_ends_within_the_interval( + monkeypatch): + from kumo.engine.integrations import mlflow as integration + + calls = [] + monkeypatch.setattr(integration, '_metric_monitors', lambda: [_Monitor()]) + with integration._SystemMetrics(lambda metrics, step: calls.append( + (metrics, step)), interval=60): + pass + + assert calls == [({'system/cpu': 1.0}, 0)] From 94784dc2968a41c983102b7798e32c33f90dadc4 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 18:52:04 +0000 Subject: [PATCH 053/278] Report an unfinished fit tracker through the error policy When the tracking thread is still running after the bounded wait, record a TimeoutError on the session so on_error='raise' raises and 'warn' warns, after the fitted result is kept. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_tracking.py | 6 +-- packages/kumo-engine/test/test_fit_events.py | 43 +++++++++++++++++++ .../test/test_snowflake_tracking.py | 8 ++-- .../test/test_tracking_contract.py | 9 ++-- 4 files changed, 55 insertions(+), 11 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_tracking.py b/packages/kumo-engine/kumo/engine/_tracking.py index 959886f7f..556fb4aa4 100644 --- a/packages/kumo-engine/kumo/engine/_tracking.py +++ b/packages/kumo-engine/kumo/engine/_tracking.py @@ -477,9 +477,9 @@ def finish_tracker() -> None: pass worker.join(max(0.0, deadline - time.monotonic())) if worker.is_alive(): - logger.warning( - 'Experiment tracking did not finish within %ss; it goes on ' - 'in the background and the fit returns', FINISH_SECONDS) + self._fail(TimeoutError( + f'experiment tracking did not finish within ' + f'{FINISH_SECONDS:g}s; it goes on in the background')) def _submit(self, job: Callable[[], Any]) -> bool: if self._behind: diff --git a/packages/kumo-engine/test/test_fit_events.py b/packages/kumo-engine/test/test_fit_events.py index 196b9e077..5e881c759 100644 --- a/packages/kumo-engine/test/test_fit_events.py +++ b/packages/kumo-engine/test/test_fit_events.py @@ -4,6 +4,7 @@ import json import logging +import threading import numpy as np import pandas as pd @@ -121,6 +122,48 @@ def fail(*args, **kwargs): } +@pytest.mark.parametrize('on_error', ['warn', 'raise']) +def test_an_unfinished_tracker_is_reported_by_the_error_policy( + table_and_snapshot, tmp_path, monkeypatch, on_error): + import kumo.engine._tracking as tracking + + monkeypatch.setattr(tracking, 'FINISH_SECONDS', 0.5) + release = threading.Event() + + class Blocked: + def start(self, record): + pass + + def on_trial(self, record, event): + pass + + def finish(self, record, error): + release.wait(30) + + table, snapshot = table_and_snapshot + plan = ModelPlan() + plan.training_job.num_experiments = 1 + plan.training_job.disable_explain = True + plan.optimization.max_epochs = 2 + plan.optimization.max_steps_per_epoch = 2 + trainer = kumo.Trainer(plan) + cfg = kumo.MLOps(str(tmp_path / 'root'), tracking=Blocked(), store=False, + on_error=on_error) + try: + if on_error == 'warn': + with pytest.warns(UserWarning, match='TimeoutError: experiment ' + 'tracking did not finish within 0.5s'): + assert trainer.fit(table, snapshot=snapshot, + mlops=cfg) is trainer.result + else: + with pytest.raises(RuntimeError) as raised: + trainer.fit(table, snapshot=snapshot, mlops=cfg) + assert isinstance(raised.value.__cause__, TimeoutError) + finally: + release.set() + assert trainer.result.bundle_dir is not None + + def test_event_metrics_keep_finite_numbers(): from kumo.engine._events import EventEmitter diff --git a/packages/kumo-engine/test/test_snowflake_tracking.py b/packages/kumo-engine/test/test_snowflake_tracking.py index bca68d51c..59d211abe 100644 --- a/packages/kumo-engine/test/test_snowflake_tracking.py +++ b/packages/kumo-engine/test/test_snowflake_tracking.py @@ -77,6 +77,7 @@ def _replay(tracking, result=RESULT, error=None, during=None): None if result is None else lambda: fitted_spec(session.spec, result, EVENTS[-1].metrics['seconds']), error) + return session def test_native_model_lineage_holds_the_run_open_during_the_fit(fake): @@ -200,8 +201,7 @@ def test_the_users_session_is_not_altered(fake): assert session.statements == [] -def test_a_hung_snowflake_call_does_not_stall_the_fit(fake, monkeypatch, - caplog): +def test_a_hung_snowflake_call_does_not_stall_the_fit(fake, monkeypatch): import kumo.engine._tracking as tracking monkeypatch.setattr(tracking, 'FINISH_SECONDS', 0.5) @@ -210,14 +210,14 @@ def test_a_hung_snowflake_call_does_not_stall_the_fit(fake, monkeypatch, lambda *args, **kwargs: release.wait(30)) began = time.monotonic() try: - _replay(_tracking()) + session = _replay(_tracking()) elapsed = time.monotonic() - began finally: release.set() join_tracking_threads() assert elapsed < 1.5 - assert 'did not finish within' in caplog.text + assert isinstance(session.error, TimeoutError) run, = fake._instance.runs.values() assert run['status'] == 'FINISHED' diff --git a/packages/kumo-engine/test/test_tracking_contract.py b/packages/kumo-engine/test/test_tracking_contract.py index 579092c53..22c72489a 100644 --- a/packages/kumo-engine/test/test_tracking_contract.py +++ b/packages/kumo-engine/test/test_tracking_contract.py @@ -503,7 +503,7 @@ def test_only_rank_zero_tracks(monkeypatch): assert recorder.calls == [] -def test_a_hung_tracker_does_not_stall_the_fit(monkeypatch, caplog): +def test_a_hung_tracker_does_not_stall_the_fit(monkeypatch): import kumo.engine._tracking as tracking monkeypatch.setattr(tracking, 'FINISH_SECONDS', 0.5) @@ -514,14 +514,15 @@ def start(self, record): super().start(record) release.wait(30) - hung = Hung() + session = TrackingSession(Hung()) began = time.monotonic() try: - replay(hung) + session.start(lambda: fit_spec(CONTEXT)) + session.finish(None, None) elapsed = time.monotonic() - began finally: release.set() join_tracking_threads() assert elapsed < 1.5 - assert 'did not finish within' in caplog.text + assert isinstance(session.error, TimeoutError) From e11ea474c2915a759fd173d49c301bce72582a81 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 20:46:35 +0000 Subject: [PATCH 054/278] Fail Snowflake tracking when ExperimentTracking is in use With an active ExperimentTracking run, or the singleton on another Snowpark session, start raises instead of returning untracked, so the error policy warns or raises; the user's run is left untouched. Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/integrations/snowflake.py | 9 ++++-- .../test/test_snowflake_tracking.py | 32 +++++++++++++++---- 2 files changed, 31 insertions(+), 10 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/integrations/snowflake.py b/packages/kumo-engine/kumo/engine/integrations/snowflake.py index 6d915db55..5bd9ebc0b 100644 --- a/packages/kumo-engine/kumo/engine/integrations/snowflake.py +++ b/packages/kumo-engine/kumo/engine/integrations/snowflake.py @@ -172,9 +172,12 @@ def start(self, record: RunRecord) -> None: tracking = cls(self.session) elif (not hasattr(tracking, '_run') or tracking._run is not None or getattr(tracking, '_session', None) is not self.session): - logger.warning('Snowflake ExperimentTracking is in use in this ' - 'process; this fit is not tracked in Snowflake') - return + raise RuntimeError( + f'Snowflake ExperimentTracking is in use in this process, by ' + f'an active run or another Snowpark session, so this ' + f'{record.tags.get("kumo.kind")} run is not tracked in ' + f'Snowflake; end that run, or pass SnowflakeTracking the ' + f'session it uses') self._tracking = tracking database = self.database or self.session.get_current_database() schema = self.schema or (None if self.database else diff --git a/packages/kumo-engine/test/test_snowflake_tracking.py b/packages/kumo-engine/test/test_snowflake_tracking.py index 59d211abe..37239caac 100644 --- a/packages/kumo-engine/test/test_snowflake_tracking.py +++ b/packages/kumo-engine/test/test_snowflake_tracking.py @@ -11,6 +11,8 @@ import _snowflake_experiment_fake import pytest from _snowflake_experiment_fake import FakeExperimentTracking, FakeSession +from kumoapi.model_plan import ModelPlan +from test_fit_events import table_and_snapshot # noqa: F401 from test_tracking_contract import ( CONTEXT, EVENTS, @@ -113,7 +115,12 @@ def test_failed_fit_ends_the_run_failed(fake): @pytest.mark.parametrize('busy', ['active run', 'other session']) -def test_an_instance_in_use_is_left_alone(fake, busy, caplog): +def test_an_instance_in_use_is_left_alone_and_fails_strict_mode( + fake, busy, table_and_snapshot, tmp_path, monkeypatch): + import kumo.engine as kumo + from kumo.engine.integrations.snowflake import SnowflakeTracking + + monkeypatch.setattr(SnowflakeTracking, 'check', lambda *args: None) session = FakeSession() user = fake(session, warnings_as_errors=True) user.set_experiment('mine') @@ -121,13 +128,24 @@ def test_an_instance_in_use_is_left_alone(fake, busy, caplog): user.start_run('mine_1') else: session = FakeSession() - calls = len(user.calls) - - _replay(_tracking(session)) - + calls, run = len(user.calls), user._run + table, snapshot = table_and_snapshot + plan = ModelPlan() + plan.training_job.num_experiments = 1 + plan.training_job.disable_explain = True + plan.optimization.max_epochs = 2 + plan.optimization.max_steps_per_epoch = 2 + trainer = kumo.Trainer(plan) + in_use = 'Snowflake ExperimentTracking is in use in this process' + + with pytest.raises(RuntimeError, match=in_use): + trainer.fit(table, snapshot=snapshot, mlops=kumo.MLOps( + str(tmp_path / 'root'), tracking=_tracking(session), + store=False, on_error='raise')) + + assert trainer.result is not None assert user.calls[calls:] == [] - assert user.warnings_as_errors is True - assert caplog.text.count('not tracked in Snowflake') == 1 + assert (user._run, user.warnings_as_errors) == (run, True) def test_location_defaults_to_the_session_not_the_last_fit(fake): From 6c2a7ffda44361a6a8d4be0796db287cf73e7588 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 21:23:08 +0000 Subject: [PATCH 055/278] Import MLflow's Databricks helper through import_provider cluster Co-Authored-By: Claude Opus 5.5 --- .../kumo-engine/kumo/engine/integrations/_mlflow_target.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py b/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py index 5ce74d53e..16d04af68 100644 --- a/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py +++ b/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py @@ -170,10 +170,8 @@ def _workspace_host() -> Optional[str]: # On a cluster the SDK's host is the regional one; the notebook context # names the workspace's own URL, which run links need. try: - from mlflow.utils.databricks_utils import ( - get_workspace_info_from_dbutils, - ) - host, _ = get_workspace_info_from_dbutils() + utils = import_provider('mlflow.utils.databricks_utils', 'mlflow') + host, _ = utils.get_workspace_info_from_dbutils() if host: return host except Exception: From b9d63fffcb448d480e6c7bfc039259f254bb6c55 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 23:39:26 +0000 Subject: [PATCH 056/278] Write a reconstructed snapshot manifest with manifest_bytes Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_snapshot_io.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_snapshot_io.py b/packages/kumo-engine/kumo/engine/_snapshot_io.py index 5b8060244..506176906 100644 --- a/packages/kumo-engine/kumo/engine/_snapshot_io.py +++ b/packages/kumo-engine/kumo/engine/_snapshot_io.py @@ -173,17 +173,17 @@ def _snapshot_manifest(snapshot: Any) -> bytes: RECONSTRUCTED, SnapshotManifest, SnapshotParents, + manifest_bytes, ) from kumo.engine.graph import SNAPSHOT_MANIFEST_FILE try: return (Path(snapshot.index_path) / SNAPSHOT_MANIFEST_FILE).read_bytes() except OSError: - manifest = SnapshotManifest( - id=snapshot.id, parents=SnapshotParents(source_tables={}), - **RECONSTRUCTED) - return (json.dumps(manifest.to_dict(), indent=2, sort_keys=True) + - '\n').encode() + return manifest_bytes( + SnapshotManifest(id=snapshot.id, + parents=SnapshotParents(source_tables={}), + **RECONSTRUCTED)) def export_location(snapshot: Any, path: str) -> str: From 23c3466ed7c14ac281ef49801871a1f90aa99196 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 23:39:40 +0000 Subject: [PATCH 057/278] List a remote directory's file sizes in one place Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_core/object_store_sink.py | 23 +++++++++++-------- .../kumo-engine/kumo/engine/_snapshot_io.py | 13 +++-------- 2 files changed, 17 insertions(+), 19 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/object_store_sink.py b/packages/kumo-core/src/kumo_core/object_store_sink.py index b6750b69a..4946877da 100644 --- a/packages/kumo-core/src/kumo_core/object_store_sink.py +++ b/packages/kumo-core/src/kumo_core/object_store_sink.py @@ -63,22 +63,15 @@ def upload_bundle(filesystem: FileSystem, destination: str, def holds_bundle(filesystem: FileSystem, destination: str, bundle_dir: str) -> bool: - """Whether ``destination`` holds a complete copy of ``bundle_dir``: the - same files and sizes, and the same provenance file, written last.""" + """Whether ``destination`` holds the same files, sizes and provenance.""" from kumo_core.artifacts import PROVENANCE_FILE - from pyarrow.fs import FileSelector, FileType root = Path(bundle_dir) local = { path.relative_to(root).as_posix(): path.stat().st_size for path in root.rglob('*') if path.is_file() } - remote = { - info.path[len(destination) + 1:]: info.size - for info in filesystem.get_file_info( - FileSelector(destination, recursive=True, allow_not_found=True)) - if info.type is FileType.File - } + remote = file_sizes(filesystem, destination) if remote != local or PROVENANCE_FILE not in local: return False with filesystem.open_input_stream( @@ -86,6 +79,18 @@ def holds_bundle(filesystem: FileSystem, destination: str, return reader.readall() == (root / PROVENANCE_FILE).read_bytes() +def file_sizes(filesystem: FileSystem, directory: str) -> dict[str, int]: + """The size of each file under ``directory``, by relative path.""" + from pyarrow.fs import FileSelector, FileType + + return { + info.path[len(directory) + 1:]: info.size + for info in filesystem.get_file_info( + FileSelector(directory, recursive=True, allow_not_found=True)) + if info.type is FileType.File + } + + def download_bundle(filesystem: FileSystem, source: str, local_dir: str) -> bool: """Replace ``local_dir`` with the files under ``source``; False if there diff --git a/packages/kumo-engine/kumo/engine/_snapshot_io.py b/packages/kumo-engine/kumo/engine/_snapshot_io.py index 506176906..e915a64d9 100644 --- a/packages/kumo-engine/kumo/engine/_snapshot_io.py +++ b/packages/kumo-engine/kumo/engine/_snapshot_io.py @@ -284,10 +284,8 @@ def upload(client: FileSystem, name: str, stop: threading.Event) -> None: def _holds_copy(connector: Any, location: str, snapshot_id: str, index: Path, snapshot_manifest: bytes) -> bool: - r"""Whether ``location`` holds a complete export of this snapshot: the - local index's bytes, every index file at its size, and the same - ``_provenance.json``.""" - from pyarrow.fs import FileSelector, FileType + r"""Whether ``location`` holds this index's bytes and manifest.""" + from kumo_core.object_store_sink import file_sizes from kumo.engine.graph import SNAPSHOT_MANIFEST_FILE @@ -297,12 +295,7 @@ def _holds_copy(connector: Any, location: str, snapshot_id: str, or not holds_export(index, manifest)): return False filesystem, _ = _location(connector, location) - remote = { - info.path[len(root) + 1:]: info.size - for info in filesystem.get_file_info( - FileSelector(root, recursive=True, allow_not_found=True)) - if info.type is FileType.File - } + remote = file_sizes(filesystem, root) listed = json.loads((index / _INDEX_MANIFEST).read_bytes())['files'] if (_DONE not in remote or f'index/{_INDEX_MANIFEST}' not in remote) or any( From a0f82e025869d7fc7e8a49fc80de8cfb9be08479 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 23:39:41 +0000 Subject: [PATCH 058/278] Build object-store URIs in one method that FilesystemSink overrides Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_core/object_store_sink.py | 48 ++++++------------- 1 file changed, 15 insertions(+), 33 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/object_store_sink.py b/packages/kumo-core/src/kumo_core/object_store_sink.py index 3951ea5b3..2d82c058e 100644 --- a/packages/kumo-core/src/kumo_core/object_store_sink.py +++ b/packages/kumo-core/src/kumo_core/object_store_sink.py @@ -6,7 +6,6 @@ import posixpath import shutil from abc import ABC, abstractmethod -from dataclasses import replace from pathlib import Path from typing import Any, Mapping, Optional, Tuple @@ -145,7 +144,7 @@ def write_files(self, connector, table_name, files, identity=export_id or str(files), file_type=kind, convert_csv=storage.convert_csv) - return f'{self.scheme}://{output}' + return self._uri(output) def publish_model( self, @@ -171,13 +170,12 @@ def publish_model( f"Publishing a model to {self.provider} needs a destination " f"path. Set output_path on the ModelExportConfig.") - filesystem, destination = self._destination(connector, path) - destination = destination.rstrip('/') + filesystem, destination = self.filesystem(connector, path) upload_bundle(filesystem, destination, bundle_dir) return PublishedModel( name=posixpath.basename(destination), version=destination, - uri=f'{self.scheme}://{destination}', + uri=self._uri(destination), metadata=dict(metadata), ) @@ -204,14 +202,16 @@ def fetch_model( f"was published to. Reuse the output_path from the export " f"config.") - filesystem, source = self._destination(connector, path) - source = source.rstrip('/') + filesystem, source = self.filesystem(connector, path) if not download_bundle(filesystem, source, local_dir): raise ValueError( - f"Nothing is published at {self.scheme}://{source}. Check the " + f"Nothing is published at {self._uri(source)}. Check the " f"path the model was exported to.") return PublishedModel(name=posixpath.basename(source), version=source, - uri=f'{self.scheme}://{source}') + uri=self._uri(source)) + + def _uri(self, path: str) -> str: + return f'{self.scheme}://{path}' def _write(self, filesystem: FileSystem, frame: pd.DataFrame, path: str, export_id: Optional[str] = None) -> str: @@ -233,12 +233,12 @@ def _write(self, filesystem: FileSystem, frame: pd.DataFrame, path: str, kind) with filesystem.open_output_stream(f'{directory}/_SUCCESS'): pass - return f'{self.scheme}://{directory}/' + return self._uri(f'{directory}/') if not path.casefold().endswith(f'.{kind}'): path += f'.{kind}' _prepared(filesystem, posixpath.dirname(path)) _write_one(filesystem, frame, path, kind) - return f'{self.scheme}://{path}' + return self._uri(path) def _occupied(self, filesystem: FileSystem, directory: str, marker: Optional[str]) -> bool: @@ -259,7 +259,7 @@ def _occupied(self, filesystem: FileSystem, directory: str, if reader.readall().decode('utf-8') == marker: return True raise ValueError( - f"Output directory '{self.scheme}://{directory}/' already " + f"Output directory '{self._uri(directory)}/' already " f"contains data from another export. Choose an empty destination " f"path.") @@ -392,11 +392,7 @@ def convert_csv(self, source, destination): class FilesystemSink(ObjectStoreSinkBackend): - """The object-store file writers over any sink's ``filesystem()``, for - the file paths of a sink that also writes tables: paths in, bare - resolved paths out.""" - scheme = 'filesystem' - + """The object-store file writers over any sink's ``filesystem()``.""" def __init__(self, sink: Any) -> None: self._sink = sink self.provider = type(sink).__name__ @@ -405,19 +401,5 @@ def _destination(self, connector, coordinate): filesystem, path = self._sink.filesystem(connector, coordinate) return filesystem, path + ('/' if coordinate.endswith('/') else '') - def _bare(self, uri: str) -> str: - return uri.removeprefix(f'{self.scheme}://') - - def write_table(self, *args, **kwargs): - return self._bare(super().write_table(*args, **kwargs)) - - def write_files(self, *args, **kwargs): - return self._bare(super().write_files(*args, **kwargs)) - - def publish_model(self, *args, **kwargs): - published = super().publish_model(*args, **kwargs) - return replace(published, uri=self._bare(published.uri)) - - def fetch_model(self, *args, **kwargs): - fetched = super().fetch_model(*args, **kwargs) - return replace(fetched, uri=self._bare(fetched.uri)) + def _uri(self, path: str) -> str: + return path From 1ed8e55fa24144722a14d0759a7c51ed81081af5 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 23:39:57 +0000 Subject: [PATCH 059/278] Let a missing sink extra raise instead of reading as a table destination Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/connector/__init__.py | 2 +- packages/kumo-engine/test/test_filestore_sink.py | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/connector/__init__.py b/packages/kumo-engine/kumo/engine/connector/__init__.py index 93265951d..d730d7a11 100644 --- a/packages/kumo-engine/kumo/engine/connector/__init__.py +++ b/packages/kumo-engine/kumo/engine/connector/__init__.py @@ -42,7 +42,7 @@ def _file_sink(connector: Any, path: str) -> Any: try: sink = connector.sink_backend - except Exception: + except NotImplementedError: return None return sink if isinstance(sink, FileStoreSink) and sink.holds_files( connector, path) else None diff --git a/packages/kumo-engine/test/test_filestore_sink.py b/packages/kumo-engine/test/test_filestore_sink.py index 8716310f8..40434672e 100644 --- a/packages/kumo-engine/test/test_filestore_sink.py +++ b/packages/kumo-engine/test/test_filestore_sink.py @@ -97,7 +97,8 @@ def test_sinks_with_a_filesystem_hold_files(tmp_path, connector): assert is_file_destination(connector, '@stage/p') assert not is_file_destination(connector, 'DB.SCHEMA.T') assert file_store(connector, 'DB.SCHEMA.T') is None - assert not is_file_destination(SimpleNamespace(), 'x') + assert not is_file_destination(SimpleNamespace(sink_backend=object()), + 'x') def test_output_configs_take_a_sinks_file_paths(connector): From 0df8cc0296f7fae73079e0cacdc8adfa37538f83 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 23:40:12 +0000 Subject: [PATCH 060/278] Trim the file-store docstrings and comments to one short line Co-Authored-By: Claude Opus 5.5 --- packages/kumo-api/kumoapi/plugin.py | 9 +++------ packages/kumo-core/src/kumo_core/artifacts.py | 2 -- .../kumo-core/src/kumo_core/object_store_sink.py | 16 +++++----------- .../kumo/engine/connector/__init__.py | 10 +++------- 4 files changed, 11 insertions(+), 26 deletions(-) diff --git a/packages/kumo-api/kumoapi/plugin.py b/packages/kumo-api/kumoapi/plugin.py index b649e302b..2e855462e 100644 --- a/packages/kumo-api/kumoapi/plugin.py +++ b/packages/kumo-api/kumoapi/plugin.py @@ -608,17 +608,14 @@ def write_files( @runtime_checkable class FileStoreSink(Protocol): - r"""A sink that also holds files: snapshot exports, model bundles and - export manifests go through its filesystem.""" + r"""A sink that can also hold files.""" def holds_files(self, connector: SourceConnector, path: str) -> bool: - r"""Whether ``path`` is a file location of this sink rather than a - table name; decided from the path alone, without I/O.""" + r"""Whether ``path`` is a file location, decided without I/O.""" ... def filesystem(self, connector: SourceConnector, path: str) -> Tuple[Any, str]: - r"""A ``pyarrow.fs.FileSystem`` and ``path`` resolved on it, for a - path :meth:`holds_files` accepts.""" + r"""A ``pyarrow.fs.FileSystem`` and ``path`` resolved on it.""" ... diff --git a/packages/kumo-core/src/kumo_core/artifacts.py b/packages/kumo-core/src/kumo_core/artifacts.py index 7a4226974..cfecddcbb 100644 --- a/packages/kumo-core/src/kumo_core/artifacts.py +++ b/packages/kumo-core/src/kumo_core/artifacts.py @@ -19,8 +19,6 @@ from .io.json import load_json, save_json from .manifest import RunManifest, RunPaths -#: Every engine artifact's manifest file; underscore-prefixed so Parquet -#: readers skip it. PROVENANCE_FILE = '_provenance.json' diff --git a/packages/kumo-core/src/kumo_core/object_store_sink.py b/packages/kumo-core/src/kumo_core/object_store_sink.py index 2d82c058e..0d37c6d90 100644 --- a/packages/kumo-core/src/kumo_core/object_store_sink.py +++ b/packages/kumo-core/src/kumo_core/object_store_sink.py @@ -31,8 +31,7 @@ def _chunked(filesystem: FileSystem) -> bool: def open_sized(filesystem: FileSystem, path: str, size: int) -> Any: - """``open_output_stream(path)``, telling a chunked filesystem the size - the file will have.""" + """``open_output_stream`` that tells a chunked filesystem the size.""" if _chunked(filesystem): from kumo_core.chunked_fs import SIZE_KEY @@ -43,8 +42,7 @@ def open_sized(filesystem: FileSystem, path: str, size: int) -> Any: def upload_bundle(filesystem: FileSystem, destination: str, bundle_dir: str) -> None: - """Copy every file of ``bundle_dir`` under ``destination``, the - provenance file last.""" + """Copy ``bundle_dir`` under ``destination``, provenance last.""" from kumo_core.artifacts import PROVENANCE_FILE root = Path(bundle_dir) @@ -60,8 +58,7 @@ def upload_bundle(filesystem: FileSystem, destination: str, def download_bundle(filesystem: FileSystem, source: str, local_dir: str) -> bool: - """Replace ``local_dir`` with the files under ``source``; False if there - are none.""" + """Replace ``local_dir`` with ``source``; False if it is empty.""" from kumo_core.sources.layout import confined from pyarrow.fs import FileSelector, FileType @@ -75,8 +72,7 @@ def download_bundle(filesystem: FileSystem, source: str, if info.type is not FileType.File: continue found = True - # Remote keys are untrusted; each relative path must stay inside - # the local bundle, even when it contains parent components. + # Remote keys are untrusted. target = destination / confined(info.path[len(source) + 1:]) target.parent.mkdir(parents=True, exist_ok=True) with filesystem.open_input_stream(info.path) as reader, \ @@ -96,12 +92,10 @@ def _destination(self, connector: SourceConnector, raise NotImplementedError def holds_files(self, connector: SourceConnector, path: str) -> bool: - """Every path of an object store is a file location.""" return True def filesystem(self, connector: SourceConnector, path: str) -> Tuple[FileSystem, str]: - """The authenticated filesystem and resolved ``path``.""" filesystem, resolved = self._destination(connector, path) return filesystem, resolved.rstrip('/') @@ -286,7 +280,7 @@ def _write_one(filesystem: FileSystem, frame: pd.DataFrame, path: str, import tempfile from kumo_core.chunked_fs import spool_root - # Spool so the size is known and a file up to one part stays plain. + # Spool to learn the size. with tempfile.TemporaryDirectory( dir=spool_root(filesystem.handler)) as spool: local = os.path.join(spool, 'frame.parquet') diff --git a/packages/kumo-engine/kumo/engine/connector/__init__.py b/packages/kumo-engine/kumo/engine/connector/__init__.py index d730d7a11..de9d897dd 100644 --- a/packages/kumo-engine/kumo/engine/connector/__init__.py +++ b/packages/kumo-engine/kumo/engine/connector/__init__.py @@ -49,22 +49,18 @@ def _file_sink(connector: Any, path: str) -> Any: def is_file_destination(connector: Any, path: str) -> bool: - r"""Whether ``path`` is a file location of ``connector``'s sink - (`kumoapi.plugin.FileStoreSink`) rather than a table; no I/O.""" + r"""Whether ``path`` is a file location rather than a table.""" return _file_sink(connector, path) is not None def file_store(connector: Any, path: str) -> Optional[Tuple[Any, str]]: - r"""``connector``'s ``(pyarrow filesystem, resolved path)`` for a file - destination; None for a table.""" + r"""The ``(filesystem, path)`` of a file destination; None for a table.""" sink = _file_sink(connector, path) return None if sink is None else sink.filesystem(connector, path) def model_sink(connector: Any, path: Optional[str]) -> Any: - r"""The sink that publishes or fetches a model at ``path``: the file - writers over ``connector``'s sink (`FilesystemSink`) for a file path of - a sink without its own, else the sink.""" + r"""The sink that writes files at ``path``.""" from kumo_core.object_store_sink import ( FilesystemSink, ObjectStoreSinkBackend, From 51b1e8b8e82e6ad3b8f864d0a0a6ad6e19e81e60 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 23:40:20 +0000 Subject: [PATCH 061/278] Stop shadowing stamp in _source_change and import at the top of source_changes Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/graph.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index ea5c66bed..1821d148f 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -321,9 +321,9 @@ def _source_change(table: Table, recorded: Optional[SourceTableEntry], source = table.definition.source_table if not isinstance(source, CustomSQLSourceTable): connector = table.connector - stamp = connector.source_backend.get_data_version( + version = connector.source_backend.get_data_version( connector, table.table_name) - now = source_table(connector, table.table_name, None, stamp, + now = source_table(connector, table.table_name, None, version, record_sources=record_sources) current = now.version comparable = comparable and now.version_kind in _COMPARABLE @@ -482,6 +482,8 @@ def load_from_connector(cls, connector: Connector, *, def source_changes(self, graph: 'Graph') -> Dict[str, Dict[str, Any]]: r"""Each of `graph`'s tables' source status since this snapshot: `unchanged`, `changed` or `unknown`. Never raises.""" + from kumo_core.sources.parallel import run_in_parallel + try: manifest = self.manifest recorded = ({} if manifest is None or manifest.reconstructed else @@ -490,8 +492,6 @@ def source_changes(self, graph: 'Graph') -> Dict[str, Dict[str, Any]]: or manifest.record_sources is not False) except Exception: recorded, record_sources = {}, True - from kumo_core.sources.parallel import run_in_parallel - tables = graph._named_tables changes = run_in_parallel( [ From 30429119daf49ce0bdba7956d6cac011aa6814ed Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 23:40:20 +0000 Subject: [PATCH 062/278] Move the local snapshot load out of GraphSnapshot.load_from_connector Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/graph.py | 87 ++++++++++++----------- 1 file changed, 47 insertions(+), 40 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index 1821d148f..dbfebd48e 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -302,6 +302,38 @@ def _loaded(store: Any, handle: Any) -> Optional[Any]: return None +def _load_local(connector: Any, path: str, root: str, + manifest: Dict[str, Any]) -> Any: + from kumo_core.diskgraph.snapshot import DiskgraphSnapshotHandle + + from kumo.engine._snapshot_io import holds_export, load_into, read + + store = _store() + handle = DiskgraphSnapshotHandle( + graph_key=manifest['graph_key'], + graph_data_version=manifest['data_version']) + target = Path(store.path(handle)) + with _lock(manifest['graph_key']): + result = _loaded(store, handle) + if result is not None and holds_export(Path(result.index_path), + manifest): + return result + try: + load_into(connector, path, root, manifest, target) + except FileNotFoundError: + # A newer export of this snapshot committed meanwhile. + newer_root, newer = read(connector, path) + if (newer_root == root + or newer['snapshot_id'] != manifest['snapshot_id']): + raise + load_into(connector, path, newer_root, newer, target) + result = _loaded(store, handle) + if result is None: + raise RuntimeError(f'The snapshot loaded from {path} did not pass ' + 'its integrity check.') + return result + + def _snapshot_id(graph_key: str, data_version: str) -> str: return f'{graph_key}@{data_version[:16]}' @@ -419,10 +451,8 @@ def load_from_connector(cls, connector: Connector, *, graph: if given, warn when its source tables changed since the snapshot was built (see `source_changes`). """ - from kumo_core.diskgraph.snapshot import DiskgraphSnapshotHandle - from kumo.engine._mlops import MLOps, on_root, snapshot_path - from kumo.engine._snapshot_io import holds_export, load_into, read + from kumo.engine._snapshot_io import read if (path is None) == (id is None): raise ValueError('Pass exactly one of path or id.') @@ -433,52 +463,29 @@ def load_from_connector(cls, connector: Connector, *, path = snapshot_path(id) connector, path = on_root(connector, path) root, manifest = read(connector, path) - graph_key = manifest['graph_key'] - if graph is not None and graph.id != graph_key: + if graph is not None and graph.id != manifest['graph_key']: report_other_graph( f"Snapshot at {path} was taken of another graph.", graph._graph_config, manifest['graph_config'], stacklevel=2) - store = _store() - handle = DiskgraphSnapshotHandle( - graph_key=graph_key, graph_data_version=manifest['data_version']) - with _lock(graph_key): - result = _loaded(store, handle) - if result is not None and not holds_export( - Path(result.index_path), manifest): - result = None - if result is None: - try: - load_into(connector, path, root, manifest, - Path(store.path(handle))) - except FileNotFoundError: - # A newer export of this snapshot committed meanwhile. - newer_root, newer = read(connector, path) - if (newer_root == root - or newer['snapshot_id'] != manifest['snapshot_id']): - raise - root, manifest = newer_root, newer - load_into(connector, path, root, manifest, - Path(store.path(handle))) - result = _loaded(store, handle) - if result is None: - raise RuntimeError(f'The snapshot loaded from {path} did not ' - 'pass its integrity check.') + result = _load_local(connector, path, root, manifest) snapshot = cls(id=manifest['snapshot_id'], index_path=Path(result.index_path), graph_config=manifest['graph_config']) if graph is not None: - changed = [ - alias - for alias, change in snapshot.source_changes(graph).items() - if change['status'] == 'changed' - ] - if changed: - warnings.warn( - f"Snapshot at {path} predates source table(s) " - f"{', '.join(changed)}; see `snapshot.source_changes`.", - UserWarning, stacklevel=2) + snapshot._warn_changed_sources(graph, path) return snapshot + def _warn_changed_sources(self, graph: 'Graph', path: str) -> None: + changed = [ + alias for alias, change in self.source_changes(graph).items() + if change['status'] == 'changed' + ] + if changed: + warnings.warn( + f"Snapshot at {path} predates source table(s) " + f"{', '.join(changed)}; see `snapshot.source_changes`.", + UserWarning, stacklevel=3) + def source_changes(self, graph: 'Graph') -> Dict[str, Dict[str, Any]]: r"""Each of `graph`'s tables' source status since this snapshot: `unchanged`, `changed` or `unknown`. Never raises.""" From 8374ae88661ad39d67c2c83d9dfea8be4dbf6353 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 23:40:29 +0000 Subject: [PATCH 063/278] Build the stub of a warm-start model without a manifest with stub() Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/trainer/_provenance.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/_provenance.py b/packages/kumo-engine/kumo/engine/trainer/_provenance.py index 67bb84d56..6bb54abb1 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_provenance.py +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -77,8 +77,7 @@ def parent_model(bundle_dir: Path, bundle_id: Optional[str] manifest = read_provenance(bundle_dir) if manifest is not None: return manifest - return None if bundle_id is None else ModelManifest(id=bundle_id, - reconstructed=True) + return None if bundle_id is None else stub('model', bundle_id) def read_provenance(bundle_dir: Path) -> Optional[ModelManifest]: From 754cf5e972c33cfd7e1fafdbd779c95411b37de7 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 23:40:29 +0000 Subject: [PATCH 064/278] Shorten a dev checkout's build commit to 12 characters, as wheels do Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_provenance.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/kumo-engine/kumo/engine/_provenance.py b/packages/kumo-engine/kumo/engine/_provenance.py index 7f0bef9eb..7b440d2bf 100644 --- a/packages/kumo-engine/kumo/engine/_provenance.py +++ b/packages/kumo-engine/kumo/engine/_provenance.py @@ -28,7 +28,7 @@ def _build_commit(version: str) -> Optional[str]: root = Path(__file__).resolve().parents[2] if not (root / 'pyproject.toml').is_file(): return None - command = ['git', '-C', str(root), 'rev-parse', '--short', 'HEAD'] + command = ['git', '-C', str(root), 'rev-parse', '--short=12', 'HEAD'] try: return subprocess.run(command, capture_output=True, text=True, check=True, timeout=10).stdout.strip() or None From f312483c522f5a0ce348ade15deabf8bc09b989a Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 23:40:30 +0000 Subject: [PATCH 065/278] Trim chunked_fs docstrings to one short line Co-Authored-By: Claude Opus 5.5 --- .../kumo-core/src/kumo_core/chunked_fs.py | 37 +++---------------- 1 file changed, 5 insertions(+), 32 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/chunked_fs.py b/packages/kumo-core/src/kumo_core/chunked_fs.py index 24a582f2b..57361d82b 100644 --- a/packages/kumo-core/src/kumo_core/chunked_fs.py +++ b/packages/kumo-core/src/kumo_core/chunked_fs.py @@ -1,21 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -r"""A pyarrow filesystem over stores that only put, list, read and remove. - -A file larger than :attr:`ChunkedFileHandler.part_bytes` is stored as -``.chunks//00000 ... NNNNN`` followed by an empty -``_SUCCESS``; anything smaller is one plain object. Readers see one logical -file either way. A write removes the other form and older generations only -once its own is committed, and a write that failed commits nothing; if a -crash leaves both forms, the plain object wins, and of several committed -generations the newest. A writer -spools one part at a time on local disk, under :func:`spool_root`. A handler -that can stream an upload implements ``_stream``: a writer told the size -(:data:`SIZE_KEY` in the stream metadata) streams a file of up to -``part_bytes`` as one plain object and a larger one as parts, without local -disk; without a size, up to ``plain_bytes`` is spooled and stays plain. -""" +r"""A pyarrow filesystem storing large files as ``.chunks/`` parts.""" import contextlib import io @@ -42,14 +28,12 @@ PART_BYTES = 256 << 20 CHUNKS = '.chunks' DONE = '_SUCCESS' -#: Stream metadata key carrying the size a written file will have. SIZE_KEY = 'kumo-size' _spool_dir: Union[None, str, Callable[[], str]] = None def set_spool_dir(path: Union[None, str, Callable[[], str]]) -> None: - r"""Where writers spool parts unless a handler sets ``spool_dir``: a - path, a callable read at each write, or ``None`` (tempfile's).""" + r"""Where writers spool parts unless a handler sets ``spool_dir``.""" global _spool_dir _spool_dir = path @@ -63,8 +47,6 @@ def close(self) -> None: def spool_root(handler: 'ChunkedFileHandler') -> Optional[str]: - r"""Where ``handler`` stages local copies: its ``spool_dir``, else the - :func:`set_spool_dir` directory (created), else None (tempfile's).""" root = handler.spool_dir or (_spool_dir() if callable(_spool_dir) else _spool_dir) if root is not None: @@ -73,35 +55,29 @@ def spool_root(handler: 'ChunkedFileHandler') -> Optional[str]: class ChunkedFileHandler(FileSystemHandler): - r"""Subclasses supply the five store primitives below, and optionally - ``_stream``.""" + r"""Subclasses supply the store primitives below.""" part_bytes = PART_BYTES spool_dir: Optional[str] = None - #: Largest plain file when ``_stream`` is implemented (<= part_bytes). plain_bytes = PART_BYTES def _stream(self, path: str) -> BinaryIO: - r"""Override to stream parts: a writable whose ``close`` commits the - object at ``path``.""" + r"""Optional: a writable whose ``close`` commits ``path``.""" raise NotImplementedError def _size(self, path: str) -> Optional[int]: - r"""Size of exactly the object at ``path``, or ``None``.""" raise NotImplementedError def _list(self, directory: str) -> Dict[str, int]: - r"""Every object below ``directory/``, by full path.""" raise NotImplementedError def _put(self, local: str, path: str) -> None: - r"""Uploads ``local``, whose basename matches ``path``'s.""" + r"""Upload ``local``; its basename matches ``path``'s.""" raise NotImplementedError def _part(self, path: str, size: int) -> Part: raise NotImplementedError def _remove(self, path: str, *, tree: bool) -> None: - r"""Removes the object at ``path``, or everything below it.""" raise NotImplementedError def normalize_path(self, path): @@ -120,8 +96,6 @@ def _layout(self, path: str) -> Optional[List[Tuple[str, int]]]: return layout def _forget(self, path: str) -> None: - r"""Drops cached layouts at or below ``path``; writes and deletes - through this handler call it.""" cache = self.__dict__.get('_layouts', {}) path = path.rstrip('/') for name in list(cache): @@ -229,7 +203,6 @@ def _logical_parts(objects: Dict[str, int], def _logical_files(objects: Dict[str, int]) -> Dict[str, int]: - r"""Collapses committed ``.chunks`` directories into their files.""" files: Dict[str, int] = {} chunked: Dict[str, Dict[str, int]] = {} for path, size in objects.items(): From 1a4cf51d997ab581040279edf1dbdb817a398ffd Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 23:40:36 +0000 Subject: [PATCH 066/278] Read the build in _table_id through the module's provenance import Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/pquery/training_table.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/pquery/training_table.py b/packages/kumo-engine/kumo/engine/pquery/training_table.py index db3888293..acb181d7a 100644 --- a/packages/kumo-engine/kumo/engine/pquery/training_table.py +++ b/packages/kumo-engine/kumo/engine/pquery/training_table.py @@ -321,12 +321,10 @@ def _table_id(table: TrainingTable, snapshot: GraphSnapshot, custom: bool) -> str: from kumo_core.io.json import to_json_dict - from kumo.engine._provenance import environment - if custom and source_version is None: from kumo.engine._scratch import new_run_id return f'tt-{new_run_id()}' - build = environment() + build = provenance.environment() key = json.dumps([ snapshot.id, table.query, to_json_dict(table.generation_plan), seed, From 1ab444bfaf6ab34f5fe2901f3010d3d13e31bcbb Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 23:40:38 +0000 Subject: [PATCH 067/278] Take only a callable as the chunked spool directory Co-Authored-By: Claude Opus 5.5 --- packages/kumo-core/src/kumo_core/chunked_fs.py | 10 ++++------ packages/kumo-core/tests/test_chunked_fs.py | 2 +- 2 files changed, 5 insertions(+), 7 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/chunked_fs.py b/packages/kumo-core/src/kumo_core/chunked_fs.py index 57361d82b..67c6df6ed 100644 --- a/packages/kumo-core/src/kumo_core/chunked_fs.py +++ b/packages/kumo-core/src/kumo_core/chunked_fs.py @@ -19,7 +19,6 @@ Optional, Protocol, Tuple, - Union, ) import pyarrow as pa @@ -29,13 +28,13 @@ CHUNKS = '.chunks' DONE = '_SUCCESS' SIZE_KEY = 'kumo-size' -_spool_dir: Union[None, str, Callable[[], str]] = None +_spool_dir: Optional[Callable[[], str]] = None -def set_spool_dir(path: Union[None, str, Callable[[], str]]) -> None: +def set_spool_dir(spool_dir: Optional[Callable[[], str]]) -> None: r"""Where writers spool parts unless a handler sets ``spool_dir``.""" global _spool_dir - _spool_dir = path + _spool_dir = spool_dir class Part(Protocol): @@ -47,8 +46,7 @@ def close(self) -> None: def spool_root(handler: 'ChunkedFileHandler') -> Optional[str]: - root = handler.spool_dir or (_spool_dir() - if callable(_spool_dir) else _spool_dir) + root = handler.spool_dir or (_spool_dir() if _spool_dir else None) if root is not None: os.makedirs(root, exist_ok=True) return root diff --git a/packages/kumo-core/tests/test_chunked_fs.py b/packages/kumo-core/tests/test_chunked_fs.py index 771014fd8..2d06b60d2 100644 --- a/packages/kumo-core/tests/test_chunked_fs.py +++ b/packages/kumo-core/tests/test_chunked_fs.py @@ -229,7 +229,7 @@ def test_parts_spool_under_the_spool_dir_and_are_removed( store, tmp_path, monkeypatch): handler, fs = store monkeypatch.setattr(chunked_fs, '_spool_dir', None) - chunked_fs.set_spool_dir(str(tmp_path / 'scratch')) + chunked_fs.set_spool_dir(lambda: str(tmp_path / 'scratch')) _write(fs, 'root/a.bin', b'a' * (2 * PART + 1)) handler.spool_dir = str(tmp_path / 'own') _write(fs, 'root/b.bin', b'b') From aef1273d58f986c646c0c6b215d93160dbf31180 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 23:40:44 +0000 Subject: [PATCH 068/278] Point the chunked spool at the scratch dir in one line Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_scratch.py | 14 +++----------- 1 file changed, 3 insertions(+), 11 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_scratch.py b/packages/kumo-engine/kumo/engine/_scratch.py index 7381e5cb5..7d43b4828 100644 --- a/packages/kumo-engine/kumo/engine/_scratch.py +++ b/packages/kumo-engine/kumo/engine/_scratch.py @@ -38,6 +38,8 @@ from pathlib import Path from typing import Union +from kumo_core.chunked_fs import set_spool_dir + _DEFAULT_ROOT = Path.home() / ".cache" / "kumo-engine" _root = _DEFAULT_ROOT @@ -68,17 +70,7 @@ def get_scratch_dir() -> Path: return _root -def _spool_dir() -> str: - return str(get_scratch_dir() / "spool") - - -def _spool_parts() -> None: - from kumo_core.chunked_fs import set_spool_dir - - set_spool_dir(_spool_dir) - - -_spool_parts() +set_spool_dir(lambda: str(get_scratch_dir() / "spool")) def new_run_id() -> str: From 2ab94a18ac6aa1de9446c7371bb32ed995fa8c1f Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 23:40:48 +0000 Subject: [PATCH 069/278] Name _provenance.json as what the bundle id leaves out Co-Authored-By: Claude Opus 5.5 --- packages/kumo-core/src/kumo_core/artifacts.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/artifacts.py b/packages/kumo-core/src/kumo_core/artifacts.py index 10b33ee11..e92e75871 100644 --- a/packages/kumo-core/src/kumo_core/artifacts.py +++ b/packages/kumo-core/src/kumo_core/artifacts.py @@ -258,8 +258,8 @@ def _sha256(path: Path) -> str: def verify_model_bundle(bundle: Path) -> str: """Verify fetched artifacts before parsing or staging a model bundle. - Returns the bundle id, a SHA-256 over every verified file except the - manifest (``_provenance.json``). Paths must stay inside the bundle after symlink + Returns the bundle id, a SHA-256 over every verified file except + ``_provenance.json``. Paths must stay inside the bundle after symlink resolution; declared sizes and SHA-256 digests must match. Older manifests may omit sizes or leave training-spec.json unlisted. """ From 3144052618cbf335d8943486b53a11ce821ff81c Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 23:40:48 +0000 Subject: [PATCH 070/278] Import warnings once at the top of artifact_export Co-Authored-By: Claude Opus 5.5 --- packages/kumo-core/src/kumo_core/artifact_export.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/kumo-core/src/kumo_core/artifact_export.py b/packages/kumo-core/src/kumo_core/artifact_export.py index 311b972bc..cc5371652 100644 --- a/packages/kumo-core/src/kumo_core/artifact_export.py +++ b/packages/kumo-core/src/kumo_core/artifact_export.py @@ -13,6 +13,7 @@ from __future__ import annotations import hashlib +import warnings from contextlib import contextmanager from typing import TYPE_CHECKING, Any @@ -239,7 +240,6 @@ def export_selected_files(files, output_path, *, fs, identity, try: fs.write_bytes(osp.join(output_path, name), payload) except Exception as error: - import warnings warnings.warn(f'Exporting without {name}: {error}', UserWarning, stacklevel=2) return publish_export_destination(output_path, fs=fs) From ca495f07acdf5be06a55c4afd4611cbe0b773d80 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 23:40:50 +0000 Subject: [PATCH 071/278] Resolve each prediction destination's writer once per export Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/trainer/trainer.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 3e87b954e..2f91d02bb 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -157,15 +157,17 @@ def export(self, output_config: PredictionOutputConfig) -> None: self._export(output_config) def _export(self, output_config: PredictionOutputConfig) -> None: - for destination in output_config.destinations().values(): - output_config.writer(destination) + writers = { + destination: output_config.writer(destination) + for destination in output_config.destinations().values() + } outputs = [] for name, destination in output_config.destinations().items(): kind = PredictionArtifactType(name.upper()) files = PredictionResult(self.selection, kind).files outputs.append((kind, destination, files)) for kind, destination, files in outputs: - output_config.writer(destination).write_files( + writers[destination].write_files( output_config.connector, destination, list(files), output_config.write_mode, export_id=f'{self.selection.job_id}-{kind.value}') From 62fcc1c265b95f8f2534b658613cd99ef9195c50 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 23:41:04 +0000 Subject: [PATCH 072/278] Drop the test stage sink's docstring Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/test/test_filestore_sink.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/packages/kumo-engine/test/test_filestore_sink.py b/packages/kumo-engine/test/test_filestore_sink.py index 40434672e..30de65fdd 100644 --- a/packages/kumo-engine/test/test_filestore_sink.py +++ b/packages/kumo-engine/test/test_filestore_sink.py @@ -68,8 +68,6 @@ def _remove(self, path, *, tree): class StageSink: - r"""Not an object store: holds '@stage/...' paths through - ``filesystem``; anything else is a table.""" def __init__(self): self.store = MemoryStore() From d24d4353017d6caa08a253a46b57c14a1f831c31 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 23:41:07 +0000 Subject: [PATCH 073/278] Shorten docstrings and comments in tracking and its integrations Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_tracking.py | 41 +++++----------- .../kumo/engine/integrations/_mlflow_model.py | 3 +- .../engine/integrations/_mlflow_target.py | 46 +++-------------- .../kumo/engine/integrations/mlflow.py | 49 +++++-------------- .../kumo/engine/integrations/snowflake.py | 11 +---- 5 files changed, 34 insertions(+), 116 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_tracking.py b/packages/kumo-engine/kumo/engine/_tracking.py index 556fb4aa4..4e489151b 100644 --- a/packages/kumo-engine/kumo/engine/_tracking.py +++ b/packages/kumo-engine/kumo/engine/_tracking.py @@ -42,8 +42,6 @@ DROPPED_PARAMS = 'kumo.dropped_params' DROPPED_METRICS = 'kumo.dropped_metrics' -# Tracker calls run on one background thread so a slow tracking server -# never stalls training; these bound how long the fit waits for it. QUEUE_SIZE = 1000 FINISH_SECONDS = 30.0 MODEL_SECONDS = 3600.0 @@ -51,11 +49,7 @@ @dataclass(frozen=True) class RunRecord: - r"""One tracked run; ``steps`` are per-epoch ``(epoch, value)`` - histories, ``artifacts`` map file names to text, ``inputs`` map - source table names to their ``connector``, ``location``, ``version`` and - ``version_kind``, and ``summaries`` are `kumo_core.summaries` results, - which each tracker logs in its own way.""" + r"""One tracked run; ``steps`` are ``(step, value)`` histories.""" name: Optional[str] = None params: Mapping[str, str] = field(default_factory=dict) @@ -71,9 +65,7 @@ class RunRecord: @dataclass(frozen=True) class RunSpec: - r"""What a run records, whatever trained or scored it: ``params`` are - flat, ``parents`` become ``kumo.`` pointer tags and ``tags`` are - the user's own.""" + r"""What a run records; ``parents`` become ``kumo.`` tags.""" kind: str name: Optional[str] = None @@ -95,8 +87,7 @@ class RunSpec: 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.""" try: table = manifest if getattr(manifest, 'kind', None) == 'model': @@ -129,9 +120,8 @@ def _text(value: Any) -> Optional[str]: def logged_name(name: str) -> str: - r"""A metric's name in a tracker: ``val_auroc`` is ``val/auroc`` and a - trial's ``best_val_auroc`` is ``val/best_auroc``, so charts group by - split.""" + r"""``val_auroc`` as ``val/auroc``, ``best_val_auroc`` as + ``val/best_auroc``.""" best = name.startswith('best_') split, _, metric = name.removeprefix('best_').partition('_') if split not in ('train', 'val', 'test') or not metric: @@ -145,8 +135,7 @@ def logged_names(metrics: Mapping[str, Any]) -> Dict[str, Any]: def by_step(steps: Mapping[str, Sequence[Tuple[int, float]]] ) -> List[Tuple[int, Dict[str, float]]]: - r"""Per-step histories regrouped as each step's metrics, in step - order.""" + r"""Each step's metrics, in step order.""" grouped: Dict[int, Dict[str, float]] = {} for name, history in steps.items(): for step, value in history: @@ -252,9 +241,7 @@ def finish(self, record: Optional[RunRecord], @dataclass class _Search: - r"""The fit run's own series: a single trial's epochs as they end, or, - for a search, the best validation ``metric`` after each trial; then - `best`, the selected trial, whose epochs the run records at the end.""" + r"""The fit run's series: one trial's epochs, or a search's best.""" single: bool metric: str @@ -286,8 +273,7 @@ def on(self, event: Event) -> Optional[Tuple[Dict[str, float], int]]: class RegistrationError(Exception): - r"""Raised by a tracker's ``log_model`` from the error that registering - the logged model raised.""" + r"""Registering, not logging, the model failed.""" class TrackingSession: @@ -337,8 +323,7 @@ def begin(tracker: Tracking) -> None: self._started.set() def pointer(self) -> Optional['TrackingRef']: - r"""The tracker's own run, once it started: its platform, run id - and URL.""" + r"""The tracker's own run, once it started.""" from kumo.engine._manifest import TrackingRef if self.spec is None or not self._started.wait(FINISH_SECONDS): @@ -381,8 +366,7 @@ def log_model( snapshot: Any, 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.""" + r"""The tracker's ``log_model``, run on the tracking thread.""" if self._worker is None or (self._tracker is not None and not callable( getattr(self._tracker, 'log_model', None))): return None, None @@ -438,8 +422,7 @@ def log() -> None: def finish(self, spec: Optional[Callable[[], RunSpec]], error: Optional[BaseException]) -> None: - r"""Ends the run with ``spec()``, its final spec (`None`: the started - one), failed if ``error``.""" + r"""Ends the run with ``spec()`` (`None`: the started spec).""" worker = self._worker if self.spec is None or worker is None: return @@ -524,8 +507,6 @@ def _fail(self, error: BaseException) -> None: def _debug() -> bool: - r"""Tracebacks of tracking failures only at DEBUG: the failure itself is - reported once, by the call that tracked it.""" return logger.isEnabledFor(logging.DEBUG) diff --git a/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py b/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py index 75d063c23..dd805bfa3 100644 --- a/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py +++ b/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py @@ -94,8 +94,7 @@ def predict(trainer: Any, snapshot: GraphSnapshot, rows: pd.DataFrame, def _in_input_order(rows: pd.DataFrame, scored: pd.DataFrame) -> pd.DataFrame: - r"""``scored`` in the order of ``rows``, which scoring does not keep; - rows scoring skipped stay out, and repeated rows pair up in order.""" + r"""``scored`` in the order of ``rows``.""" keys = ['ENTITY'] left = pd.DataFrame({'ENTITY': rows['ENTITY1'].to_numpy()}) right = scored[['ENTITY']].copy() diff --git a/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py b/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py index 16d04af68..7982f5ba5 100644 --- a/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py +++ b/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py @@ -1,18 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -r"""Where `MlflowTracking` records, and the credentials that reach it. - -MLflow clients take a tracking and a registry URI but no credentials: the -Databricks store reads them from ``DATABRICKS_*`` variables, a profile or the -notebook (``mlflow.utils.databricks_utils.get_databricks_host_creds``), and -the SageMaker plugin signs with the process-wide tracking URI and a boto3 -session (``sagemaker_mlflow.auth_provider.AuthProvider.get_auth``). With -explicit credentials, `Target.scope` sets exactly that state around each -call, under one lock, and restores it. Ambient credentials set no -credentials; a Databricks or SageMaker target still sets the process-wide -tracking URI. -""" +r"""Where `MlflowTracking` records, and the credentials that reach it.""" import logging import os @@ -33,8 +22,6 @@ def _no_header_error(record: logging.LogRecord) -> bool: - # A request-header provider that fails off its platform; the call - # proceeds without those headers. return not record.getMessage().startswith( 'Encountered unexpected error during resolving request headers') @@ -79,8 +66,7 @@ def __post_init__(self) -> None: class Target: - r"""A tracking URI, its registry URI and, for explicit credentials, the - state MLflow reads them from.""" + r"""A tracking URI, its registry URI and how to reach them.""" def __init__(self, tracking_uri: str, registry_uri: Optional[str], swap: Optional[Callable[[], ContextManager[Any]]] = None, @@ -100,8 +86,7 @@ def platform(self) -> str: return 'mlflow' def run_url(self, experiment_id: str, run_id: str) -> Optional[str]: - r"""The run's page in the MLflow UI; None where no link opens it - (SageMaker, local stores).""" + r"""The run's page in the MLflow UI, where there is one.""" if self.platform == 'databricks': host = self._host or _workspace_host() return None if host is None else ( @@ -133,8 +118,7 @@ def client(self) -> Any: return mlflow.MlflowClient(self.tracking_uri, self.registry_uri) def artifact_repository(self, uri: str) -> Any: - r"""The repository for artifacts at ``uri``, or `None` to go through - the client: MLflow's S3 store signs with the default AWS chain.""" + r"""An S3 repository signed as this target, or `None`.""" if self._aws_session is None or not uri.startswith('s3://'): return None s3 = import_provider('mlflow.store.artifact.s3_artifact_repo', @@ -167,8 +151,6 @@ def resolve(target: Any) -> Target: def _workspace_host() -> Optional[str]: - # On a cluster the SDK's host is the regional one; the notebook context - # names the workspace's own URL, which run links need. try: utils = import_provider('mlflow.utils.databricks_utils', 'mlflow') host, _ = utils.get_workspace_info_from_dbutils() @@ -184,8 +166,7 @@ def _workspace_host() -> Optional[str]: def registry_for(tracking_uri: str) -> str: - r"""The registry a tracking URI implies: Unity Catalog on Databricks, - the tracking server itself elsewhere.""" + r"""Unity Catalog on Databricks, else the tracking server itself.""" scheme, colon, rest = tracking_uri.partition(':') return f'databricks-uc{colon}{rest}' if scheme == 'databricks' else ( tracking_uri) @@ -203,8 +184,6 @@ def _databricks(connector: Any) -> Target: principal = _principal(connector.identity) if (principal is None or principal.config.authentication is DatabricksAuthenticationKind.AMBIENT): - # Unity Catalog registration downloads the logged model through the - # process-wide tracking URI, so it must point at the workspace. return Target('databricks', 'databricks-uc', swap=lambda: global_tracking_uri('databricks')) client = import_provider('kumo_databricks.session', @@ -216,10 +195,7 @@ def _databricks(connector: Any) -> Target: @contextmanager def _databricks_token(client: Any) -> Iterator[None]: - r"""``DATABRICKS_HOST`` and a short-lived ``DATABRICKS_TOKEN`` minted - from the identity's own client, with every other Databricks auth - variable unset, and the process-wide tracking URI set to - ``databricks``; the client secret never enters the environment.""" + r"""A short-lived ``DATABRICKS_TOKEN``; the secret stays out.""" config_type = import_provider('databricks.sdk.config', 'databricks').Config authorization = client.config.authenticate()['Authorization'] @@ -233,8 +209,6 @@ def _databricks_token(client: Any) -> Iterator[None]: for name in names}, 'DATABRICKS_HOST': client.config.host, 'DATABRICKS_TOKEN': authorization.partition(' ')[2] }), global_tracking_uri('databricks'): - # Unity Catalog registration downloads the logged model through the - # process-wide tracking URI, so it must point at the workspace too. yield @@ -254,9 +228,7 @@ def _sagemaker(target: SageMakerMLflow) -> Target: @contextmanager def _sagemaker_session(arn: str, session: Any) -> Iterator[None]: - r"""The plugin's signer and router read the process-wide tracking URI - and take a boto3 session from a per-thread context variable; - ``SAGEMAKER_MLFLOW_ASSUME_ROLE_ARN`` would make it assume a role.""" + r"""The SageMaker plugin's state, signing with ``session``.""" plugin = import_provider('sagemaker_mlflow.session', 'sagemaker') with ExitStack() as stack: stack.enter_context(environ({'SAGEMAKER_MLFLOW_ASSUME_ROLE_ARN': None})) @@ -267,9 +239,7 @@ def _sagemaker_session(arn: str, session: Any) -> Iterator[None]: @contextmanager def global_tracking_uri(uri: str) -> Iterator[None]: - r"""``mlflow.set_tracking_uri(uri)`` for the block. MLflow's own - ``_use_tracking_uri`` restores by unsetting ``MLFLOW_TRACKING_URI``, - so this restores the module value and the variable itself.""" + r"""``mlflow.set_tracking_uri(uri)`` for the block.""" mlflow = import_provider('mlflow', 'mlflow') utils = import_provider('mlflow.tracking._tracking_service.utils', 'mlflow') diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index 4a6d7ddb1..a7e37deef 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -77,14 +77,12 @@ _MAX_PARAM_KEY = 250 _MAX_PARAM_VALUE = 6000 -# MLflow's portable name charset (its Windows rule, which omits ':'). +# MLflow's Windows-safe charset. _PARAM_KEY = re.compile(r'[/\w.\- ]+') -# MLflow's defaults (7 retries, 120 s timeout) can hold a call for minutes. _HTTP_LIMITS = { 'MLFLOW_HTTP_REQUEST_MAX_RETRIES': '2', 'MLFLOW_HTTP_REQUEST_TIMEOUT': '30', } -# Model upload only, and only where the limits above apply. _UPLOAD_HTTP = {'MLFLOW_HTTP_REQUEST_TIMEOUT': '600'} _PATH_CONNECTORS = ('LOCAL_FILE', 'S3', 'GCS') _SAGEMAKER_MODEL = re.compile(r'[a-zA-Z0-9](-*[a-zA-Z0-9]){0,56}') @@ -190,15 +188,11 @@ def client(self) -> Iterator[Any]: yield self.target.client() def check_model_name(self, name: str) -> None: - r"""Raises :class:`ValueError` unless this tracker's registry - accepts ``name``.""" + r"""Raises :class:`ValueError` unless the registry accepts ``name``.""" check_model_name(name, self.target.registry_uri or '') def check(self, registry: Optional[ModelRegistryLocation] = None) -> None: - r"""Raises before any work unless the server is reachable with these - credentials and, given ``registry``, can register its model: one - experiment lookup, one registry listing and, on Unity Catalog, the - model's schema.""" + r"""Raises unless the server (and ``registry``) is usable.""" self._require_experiment() if registry is not None: self.check_model_name(registry.model_name) @@ -253,8 +247,7 @@ def _require_experiment(self) -> None: def load_model(self, location: ModelRegistryLocation, **kwargs: Any) -> 'Trainer': - r"""The `Trainer` of the version of ``location`` its alias or - version names, from this tracker's registry.""" + r"""The `Trainer` of ``location``'s alias or version.""" if (location.alias is None) == (location.version is None): raise ValueError('Loading from an MLflow registry takes a ' 'ModelRegistryLocation with an alias or a ' @@ -266,9 +259,8 @@ def load_model(self, location: ModelRegistryLocation, def register(self, result: Any, location: ModelRegistryLocation) -> PublishedModel: - r"""Registers a fitted or loaded model (a `TrainingResult`) at - ``location``: the version that already holds its bundle, else a new - version of the model logged in a new run of this experiment.""" + r"""Registers a `TrainingResult` at ``location``, reusing a version + that holds its bundle.""" from kumo.engine._manifest import stub from kumo.engine.graph import GraphSnapshot from kumo.engine.trainer._spec import read_training_spec @@ -318,7 +310,6 @@ def start(self, record: RunRecord) -> None: try: autologging = import_provider('mlflow.utils.autologging_utils', 'mlflow') - # Not reentrant: leaving it clears the flag. if not getattr(autologging, '_AUTOLOGGING_GLOBALLY_DISABLED', False): self._guard.enter_context(autologging.disable_autologging()) @@ -452,8 +443,6 @@ def _log_system_metrics(self, metrics: Mapping[str, float], self._log_metrics(self._parent, metrics, step) def log_summaries(self, summaries: Sequence[Summary]) -> None: - r"""Logs each summary on the current run through its type's call in - `_SINKS`; a test scalar also goes on the logged model.""" for summary in summaries: _SINKS[type(summary)](self, summary) @@ -607,11 +596,7 @@ def _log_metrics(self, run_id: Optional[str], class _SystemMetrics: - r"""Samples MLflow's CPU, disk, network and GPU monitors every - ``interval`` seconds on its own thread and hands each aggregate to - ``log(metrics, step)``, which logs it through the tracker's - credentials. (MLflow's own ``SystemMetricsMonitor`` logs with - process-wide ones.)""" + r"""MLflow's system monitors, logged with the tracker's credentials.""" def __init__(self, log: Callable[[Mapping[str, float], int], None], interval: float = 10.0) -> None: @@ -636,7 +621,7 @@ def __exit__(self, *exc: Any) -> None: def _run(self) -> None: step = failures = 0 stopped = False - while not stopped: # and a last reading when the run ends + while not stopped: stopped = self._stop.wait(self._interval) metrics: Dict[str, float] = {} for monitor in self._monitors: @@ -677,11 +662,7 @@ def _metric_monitors() -> List[Any]: def check_model_name(name: str, registry_uri: str) -> None: - r"""Raises :class:`ValueError` unless the MLflow registry at - ``registry_uri`` accepts ``name``: SageMaker managed MLflow takes letters, - digits and hyphens; Unity Catalog takes ``catalog.schema.model``, each - part at most 255 characters without spaces, ``/``, ``.`` or control - characters.""" + r"""Raises :class:`ValueError` unless ``registry_uri`` accepts ``name``.""" is_databricks_unity_catalog_uri = import_provider( 'mlflow.utils.uri', 'mlflow').is_databricks_unity_catalog_uri @@ -721,8 +702,7 @@ def _sagemaker_ui_hint(tracking_uri: str) -> Optional[str]: def metric_key(name: str) -> str: - r"""The MLflow name of a Kumo metric: MLflow rejects ``@``, so - ``map@10`` is logged as ``map_at_10``.""" + r"""``map@10`` as ``map_at_10``: MLflow rejects ``@``.""" return name.replace('@', '_at_') @@ -846,9 +826,7 @@ def _log_model(target: Any, client: Any, bundle_dir: Path, snapshot: Optional['GraphSnapshot'], *, name: str, run_id: str, experiment_id: Optional[str], serving_snapshot: Optional[str]) -> Any: - r"""What ``mlflow.pyfunc.log_model`` does, through ``client``: it only - logs to the process-wide tracking URI. A bundle without a training spec - is logged without a signature.""" + r"""``mlflow.pyfunc.log_model`` through ``client``.""" from kumo_core.artifacts import PROVENANCE_FILE from kumo.engine.integrations import _mlflow_model, _mlflow_model_code @@ -1030,9 +1008,7 @@ def _code(error: BaseException) -> str: def check_registry_schema(name: str, registry_uri: str) -> None: - r"""On Unity Catalog, raises :class:`ValueError` unless the catalog and - schema of a ``catalog.schema.model`` name exist and these credentials - can use them (one metadata call, no warehouse).""" + r"""On Unity Catalog, raises unless the model's schema is usable.""" uri = import_provider('mlflow.utils.uri', 'mlflow') parts = name.split('.') if not uri.is_databricks_unity_catalog_uri(registry_uri) or len( @@ -1060,7 +1036,6 @@ def check_registry_schema(name: str, registry_uri: str) -> None: def _no_serving_check(record: logging.LogRecord) -> bool: - # MLflow's input-example check predicts without a snapshot; no switch. return not record.getMessage().startswith( 'Failed to validate serving input example') diff --git a/packages/kumo-engine/kumo/engine/integrations/snowflake.py b/packages/kumo-engine/kumo/engine/integrations/snowflake.py index 5bd9ebc0b..cc2667786 100644 --- a/packages/kumo-engine/kumo/engine/integrations/snowflake.py +++ b/packages/kumo-engine/kumo/engine/integrations/snowflake.py @@ -51,7 +51,6 @@ _MAX_METRICS = 200 _MAX_PARAMS = 1000 -# Metric name budgets: per-epoch metrics give way first, then trial summaries. _EPOCH_METRICS = _MAX_METRICS - 60 _SUMMARY_METRICS = _MAX_METRICS - 30 _FINISH_PARAMS = 20 @@ -254,9 +253,7 @@ def _new_run_name(self, tracking: Any, experiment: Any, run_name: Optional[str]) -> Optional[str]: if run_name is None: return None - # start_run(name) resumes, or raises on, any run matching name as a - # SQL LIKE pattern and cannot create an existing run, so pick a name - # that is neither. + # start_run matches name as a LIKE pattern. existing = [ str(row['name']) for row in tracking._sql_client.show_runs_in_experiment( @@ -277,9 +274,6 @@ def _new_run_name(self, tracking: Any, experiment: Any, return name def log_summaries(self, summaries: Sequence[Summary]) -> None: - r"""Logs each summary on the run through its type's call in - `_SINKS`: values as metrics within the name budget, tables as JSON - artifacts; Snowflake has no datasets.""" if self._run is not None: for summary in summaries: _SINKS[type(summary)](self, summary) @@ -382,8 +376,7 @@ def _log_metrics(self, metrics: Mapping[str, float], limit: int, def _run_identifier(name: str) -> str: - r"""``name`` as an unquoted identifier: snowflake-ml puts a run's name - into its artifact URIs, where a quoted name is not found.""" + r"""``name`` as an unquoted identifier, which artifact URIs need.""" name = re.sub(r'[^A-Za-z0-9_]', '_', name) return name if re.match(r'[A-Za-z_]', name) else f'_{name}' From 83b472ef96f85a291aa2ebf379f77d46c1950305 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 23:41:08 +0000 Subject: [PATCH 074/278] Shorten docstrings in the predictor, dependencies, MLOps and fit helpers Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_deps.py | 22 +++++-------------- packages/kumo-engine/kumo/engine/_mlops.py | 18 +++++---------- .../kumo-engine/kumo/engine/_predictor.py | 10 +++------ .../kumo/engine/trainer/_run_spec.py | 5 +---- .../kumo/engine/trainer/trainer.py | 14 +++--------- 5 files changed, 17 insertions(+), 52 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_deps.py b/packages/kumo-engine/kumo/engine/_deps.py index 23b76117b..64d76a591 100644 --- a/packages/kumo-engine/kumo/engine/_deps.py +++ b/packages/kumo-engine/kumo/engine/_deps.py @@ -1,9 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -r"""What a model needs to predict besides its bundle and a snapshot: its -dependencies, recorded at fit, and the assets among them, resolved at -predict time.""" +r"""A model's dependencies, recorded at fit and resolved at predict.""" import hashlib import importlib @@ -33,15 +31,11 @@ stamp, ) -#: A place assets are stored: a URI or path, or a ``(connector, path)``. Location = Union[str, Tuple[Any, str]] class Asset(Protocol): - r"""A file or directory a model needs at predict time, addressed by its - content. ``files`` maps names to local files; the class's - ``load(directory, dependency)`` restores one from a directory holding - them, and must give the same ``fingerprint()``.""" + r"""A content-addressed file or directory a model predicts with.""" asset_kind: str name: str @@ -117,10 +111,7 @@ def imported(entry: str) -> Any: class Dependencies: - r"""Records a fit's dependencies (``require``) and resolves a fitted - model's assets (``resolve``). With MLOps ``settings`` whose export - includes models, each required asset is stored once under the - artifact root; ``locations`` are asset roots looked in first.""" + r"""Records a fit's dependencies and resolves a model's assets.""" def __init__(self, records: Sequence[Dependency] = (), settings: Any = None, locations: Sequence[Location] = ()) -> None: @@ -161,9 +152,7 @@ def assets(self, asset_kind: str) -> List[Dependency]: ] def resolve(self, asset_kind: str, name: str) -> Path: - r"""A local directory holding the asset, checked against its - sha256: from the given locations, then where it was stored, then - the current artifact root.""" + r"""A local directory holding the asset, checked by sha256.""" record = next((record for record in self.assets(asset_kind) if record.name == name), None) if record is None: @@ -216,8 +205,7 @@ def asset_cache() -> Path: def _store(asset: Asset, sha256: str, settings: Any) -> Optional[str]: - r"""Writes ``asset`` once to ``/assets///``, - its ``_provenance.json`` last; returns where.""" + r"""Writes ``asset`` once under the root; returns where.""" from kumo_core.object_store_sink import _prepared from kumo.engine.connector import file_store diff --git a/packages/kumo-engine/kumo/engine/_mlops.py b/packages/kumo-engine/kumo/engine/_mlops.py index f2977cacb..6a6814f98 100644 --- a/packages/kumo-engine/kumo/engine/_mlops.py +++ b/packages/kumo-engine/kumo/engine/_mlops.py @@ -151,8 +151,7 @@ def _check_registry(self) -> None: check_model_location(ModelOutputConfig(connector, location=registry)) def snapshot_uri(self, snapshot: Any) -> str: - r"""Where a `GraphSnapshot`, or a snapshot id, is stored under this - artifact root, e.g. for the pyfunc's ``snapshot`` param.""" + r"""Where a `GraphSnapshot`, or a snapshot id, is stored.""" return self.root.uri(snapshot_path(snapshot)) def wants(self, kind: str) -> bool: @@ -160,9 +159,7 @@ def wants(self, kind: str) -> bool: return kind in self.kinds def locate(self, path: Path, kind: Type[M], relative: str) -> Optional[M]: - r"""Records in the manifest at ``path`` that its artifact is stored - at ``relative`` under the root, before it is copied there; returns - the manifest, if there is one.""" + r"""Sets the manifest at ``path`` to say it is at ``relative``.""" location = self.root.uri(relative) manifest = read_manifest(path, kind) if manifest is not None and manifest.location != location: @@ -255,9 +252,7 @@ def resolve(settings: Optional[MLOps] = None) -> Optional[MLOps]: def raise_on_error(settings: Optional[MLOps], error: Optional[BaseException], call: str, stacklevel: int = 2) -> None: - r"""Reports ``error``, a storing, tracking or registration failure of a - finished ``call``: raises when ``settings.on_error`` is ``'raise'``, - else warns unless `warn_failure` already warned about it.""" + r"""Raises or warns about ``error``, as ``settings.on_error`` says.""" if error is None or settings is None: return message = (f'{call} finished, but MLOps storing, tracking or ' @@ -274,8 +269,7 @@ def full_error(error: BaseException) -> str: def warn_failure(message: str, error: BaseException, stacklevel: int = 2, also: Optional[BaseException] = None) -> None: - r"""Warns ``message`` (a :class:`UserWarning`) once per ``error``; it - also reports ``also``, an earlier error that caused ``error``.""" + r"""Warns ``message`` once per ``error`` (and ``also``).""" if getattr(error, '_kumo_warned', False): return warnings.warn(message, UserWarning, stacklevel=stacklevel + 1) @@ -305,9 +299,7 @@ def snapshot_path(snapshot: Any) -> str: def snapshot_uri(snapshot: Any, settings: Optional[MLOps] = None) -> Optional[str]: - r"""A snapshot to serve against, as a location: a URI or stage path as - given; a `GraphSnapshot` or snapshot id as stored under the artifact - root of ``settings`` (else of the enclosing block).""" + r"""A snapshot location; an object or id is read from the root.""" if snapshot is None or (isinstance(snapshot, str) and not _SNAPSHOT_ID.fullmatch(snapshot)): return snapshot diff --git a/packages/kumo-engine/kumo/engine/_predictor.py b/packages/kumo-engine/kumo/engine/_predictor.py index b65ebb906..8fd5f0977 100644 --- a/packages/kumo-engine/kumo/engine/_predictor.py +++ b/packages/kumo-engine/kumo/engine/_predictor.py @@ -1,8 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -r"""What scores a model bundle: the ``predictor`` dependency its manifest -names.""" +r"""The predictor a model bundle's manifest names.""" from pathlib import Path from typing import Any, Mapping, Protocol, Tuple, Union @@ -11,13 +10,11 @@ from kumo.engine._deps import imported -#: The predictor of bundles whose manifest names none. TRAINER_PREDICTOR = 'kumo.engine._predictor:TrainerPredictor' class Predictor(Protocol): - r"""A loaded bundle: scores ``rows`` (an ``ENTITY`` frame, or a - `PredictionTable` for files) against ``snapshot``.""" + r"""Scores ``rows`` (a frame or `PredictionTable`) on ``snapshot``.""" def predict(self, rows: Any, snapshot: Any, params: Mapping[str, Any]) -> Union[pd.DataFrame, Tuple[str, ...]]: @@ -25,8 +22,7 @@ def predict(self, rows: Any, snapshot: Any, def load_predictor(bundle_dir: Union[str, Path]) -> Predictor: - r"""The predictor the bundle at ``bundle_dir`` names, loaded with - ``.load(bundle_dir)``.""" + r"""The predictor the bundle at ``bundle_dir`` names, loaded.""" from kumo.engine.trainer._provenance import read_provenance manifest = read_provenance(Path(bundle_dir)) diff --git a/packages/kumo-engine/kumo/engine/trainer/_run_spec.py b/packages/kumo-engine/kumo/engine/trainer/_run_spec.py index f93bed7c2..167b0877d 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_run_spec.py +++ b/packages/kumo-engine/kumo/engine/trainer/_run_spec.py @@ -56,10 +56,7 @@ def fit_spec(context: FitContext, *, def fitted_spec(spec: RunSpec, result: Any, seconds: Optional[float] = None, alias: Optional[str] = None, summaries: Sequence[Summary] = ()) -> RunSpec: - r"""``spec`` once the fit returned ``result``, a `TrainingResult`, - after ``seconds``: its test metrics, since the run's train and - validation series show the selected trial's epochs, the version it - registered with ``alias``, and ``summaries``.""" + r"""``spec`` once the fit returned ``result``.""" artifacts = dict(spec.artifacts) provenance = Path(result.bundle_dir) / PROVENANCE_FILE if provenance.is_file(): diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index c2eb8e3c3..7d0198fd5 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -371,11 +371,7 @@ def _store_and_register( warn: bool = True, stored: Optional[BaseException] = None, ) -> Tuple[Optional[PublishedModel], Optional[BaseException]]: - r"""Store a fitted bundle under the artifact root, log it in the - session's run and register it, as MLOps ``settings`` say; returns the - registered version and the first failure. A model whose bundle or - assets (``stored``, the asset failure) could not be stored is not - registered.""" + r"""Stores, logs and registers a bundle; returns it and any error.""" if stored is None and settings is not None and settings.wants('model'): try: _store_model(bundle_dir, manifest, settings) @@ -496,8 +492,7 @@ def _register_result(result: TrainingResult, settings: Any) -> PublishedModel: def _with_metrics(model: Optional[PublishedModel], metrics: Mapping[str, float]) -> Optional[PublishedModel]: - r"""``model`` with the fit's metrics, where its registry returned none - (an MLflow version, unlike a Snowflake one, carries none).""" + r"""``model`` with ``metrics``, where its registry returned none.""" if model is None or model.metrics: return model return replace(model, metrics=dict(metrics)) @@ -572,10 +567,7 @@ def _stage_model_bundle(artifacts: Mapping[str, str], dest: Path) -> Path: def _without_torch_scatter_hint(method: _F) -> _F: r"""Runs ``method`` without PyTorch Geometric's GPU hint that min/max aggregation "can be accelerated via the 'torch-scatter' package", - Lightning's "Tip" log lines advertising its cloud products, Lightning's - note that the loader's IterableDataset defines ``__len__``, and Python's - note that forking a multi-threaded process may deadlock (the loader - workers fork before any thread that matters holds a lock). + Lightning's tips and ``__len__`` note, and Python's fork warning. kumo-engine uses PyTorch's own scatter instead of that optional extension (see ``test_wheel_uses_torch_fallbacks_without_native_extensions``), so From 16a36aebbcfe1b38c130faa5b9f00dbec3109882 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 23:41:14 +0000 Subject: [PATCH 075/278] Log MLflow system metrics through log_metrics, its exact copy Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/integrations/mlflow.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index a7e37deef..657e63f1a 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -340,7 +340,7 @@ def start(self, record: RunRecord) -> None: self._experiment_id = run.info.experiment_id if self.log_system_metrics and record.tags.get('kumo.kind') == 'fit': self._guard.enter_context( - _SystemMetrics(self._log_system_metrics)) + _SystemMetrics(self.log_metrics)) self._log_params(self._parent, record.params) self._log_artifacts(record.artifacts) if record.tags.get('kumo.kind') != 'fit': # a fit logs Datasets @@ -438,10 +438,6 @@ def finish(self, record: Optional[RunRecord], def log_metrics(self, metrics: Mapping[str, float], step: int) -> None: self._log_metrics(self._parent, metrics, step) - def _log_system_metrics(self, metrics: Mapping[str, float], - step: int) -> None: - self._log_metrics(self._parent, metrics, step) - def log_summaries(self, summaries: Sequence[Summary]) -> None: for summary in summaries: _SINKS[type(summary)](self, summary) From 1b53960155eaf08b0091c6c69d77341383ec15e7 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 23:41:24 +0000 Subject: [PATCH 076/278] Cut provenance and snapshot-export docstrings to one line Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_manifest.py | 23 ++++++------------- .../kumo-engine/kumo/engine/_snapshot_io.py | 14 ++++------- packages/kumo-engine/kumo/engine/graph.py | 6 ++--- .../kumo/engine/trainer/_provenance.py | 3 +-- 4 files changed, 14 insertions(+), 32 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_manifest.py b/packages/kumo-engine/kumo/engine/_manifest.py index 64d73bddc..1e10f6e31 100644 --- a/packages/kumo-engine/kumo/engine/_manifest.py +++ b/packages/kumo-engine/kumo/engine/_manifest.py @@ -36,8 +36,7 @@ def omit_none() -> Any: - """A field written only when it is not None, or read as an explicit - null (kept in ``extra``).""" + """A field left out of the JSON while it is None.""" return field(default=None, metadata={_OMIT_NONE: True}) @@ -66,8 +65,7 @@ class Record: compare=False) def as_artifact(self, key: str) -> Optional[Artifact]: - r"""The artifact this record names, if any, as ``key`` in its - parent.""" + r"""The artifact this record names as ``key``, if any.""" return None def to_dict(self) -> Dict[str, Any]: @@ -161,8 +159,7 @@ def as_artifact(self, key: str) -> Artifact: @dataclass(frozen=True, kw_only=True) class Manifest(Record): - """An artifact's record. A stub (a parent with no record of its own) - sets only ``kind``, ``id`` and ``reconstructed``.""" + """An artifact's record; a stub sets only kind, id and reconstructed.""" kind: str = '' schema_version: int = SCHEMA_VERSION @@ -186,10 +183,7 @@ def as_artifact(self, key: str) -> Artifact: def artifacts(manifest: Manifest) -> List[Artifact]: - r"""``manifest``'s artifact, then each one it names (parents, their - source tables, dependencies, outputs), depth first, each once. A named - manifest that is not a dataset, such as a warm-start model, is listed - but not opened: its own manifest names what it was made from.""" + r"""``manifest``'s artifact and those it names, depth first, once each.""" found: Dict[Any, Artifact] = {} for artifact in _walk(manifest, manifest.kind, manifest): found.setdefault((artifact.kind, artifact.id), artifact) @@ -245,8 +239,7 @@ def stub(kind: str, id: str) -> 'Manifest': def embedded(manifest: M) -> M: - """``manifest`` as a child embeds it: without its statistics and the - fields that change after it is written.""" + """``manifest`` without the fields a child's copy leaves out.""" return replace(manifest, **{ name: None for name in manifest.EMBED_OMITS @@ -270,8 +263,7 @@ class SnapshotParents(Record): @dataclass(frozen=True, kw_only=True) class SnapshotManifest(Manifest): - """What a graph snapshot was built from: each source table's connector, - location and version, and the engine build that wrote it.""" + """The source tables and engine build a snapshot was made from.""" kind: str = 'snapshot' IS_DATASET: ClassVar[bool] = True @@ -287,8 +279,7 @@ class TrainingTableParents(Record): @dataclass(frozen=True, kw_only=True) class TrainingTableManifest(Manifest): - """How a training table was generated: from which snapshot (embedded) - and, for caller-supplied rows, which source table.""" + """The snapshot, query and rows a training table was made from.""" kind: str = 'training_table' IS_DATASET: ClassVar[bool] = True diff --git a/packages/kumo-engine/kumo/engine/_snapshot_io.py b/packages/kumo-engine/kumo/engine/_snapshot_io.py index e915a64d9..6c331811e 100644 --- a/packages/kumo-engine/kumo/engine/_snapshot_io.py +++ b/packages/kumo-engine/kumo/engine/_snapshot_io.py @@ -311,8 +311,7 @@ def _holds_copy(connector: Any, location: str, snapshot_id: str, def _commit(filesystem: FileSystem, root: str, generation: str, marker: bytes) -> None: - r"""Writes the marker in one step: a rename on a local filesystem (an - in-place write there truncates first), a single object put elsewhere.""" + r"""Writes the marker atomically: a local rename, else one object put.""" if filesystem.type_name == 'local': staged = f'{root}/{generation}/{_MANIFEST}' _write(filesystem, staged, marker) @@ -324,9 +323,7 @@ def _commit(filesystem: FileSystem, root: str, generation: str, def _clean(filesystem: FileSystem, root: str) -> None: - r"""Best effort: remove generations that a later commit superseded, and - abandoned ones. A generation is marked done only after its commit, so - one listed done before the marker is re-read lost to a later one.""" + r"""Removes superseded and abandoned generations, best effort.""" from pyarrow.fs import FileSelector, FileType try: @@ -390,17 +387,14 @@ def read(connector: Any, path: str) -> Tuple[str, dict]: def export_identity(manifest: dict) -> str: - r"""What makes one export's bytes differ from another's of the same - snapshot: its range checksums and index manifest.""" + r"""A hash of an export's range checksums and index manifest.""" return hashlib.sha256(json.dumps( [manifest['range_bytes'], manifest['crc32'], manifest['index_manifest']], sort_keys=True).encode()).hexdigest() def holds_export(index: Path, manifest: dict) -> bool: - r"""Whether the local index at ``index`` has the export's bytes. An - index without a recorded export identity (built here, or loaded by an - older engine) is checksummed once.""" + r"""Whether the local ``index`` has the export's bytes.""" from diskgraph import snapshot_range_checksums identity = export_identity(manifest) diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index dbfebd48e..8c21e3cf9 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -415,14 +415,12 @@ def _graph_key(self) -> str: @property def _data_version(self) -> str: - """The full fingerprint of the source versions the index was built - from; `id` carries its first 16 characters.""" + """The full source-version fingerprint; `id` holds its first 16.""" return Path(self.index_path).name @cached_property def _native_graph(self) -> NativeGraph: - """`graph_config` in `kumo_core`'s normalized form, as - `Graph.native_graph`.""" + """`graph_config` normalized, as `Graph.native_graph`.""" return normalize_graph(self.graph_config) def export(self, output_config: 'SnapshotOutputConfig') -> str: diff --git a/packages/kumo-engine/kumo/engine/trainer/_provenance.py b/packages/kumo-engine/kumo/engine/trainer/_provenance.py index 6bb54abb1..72bb3f859 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_provenance.py +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -56,8 +56,7 @@ def write_provenance( def embedded_model(manifest: ModelManifest) -> ModelManifest: - r"""``manifest`` as a child model embeds it: its own model parents cut - to stubs, so a chain of retrains does not grow each manifest.""" + r"""``manifest`` for embedding, its own model parents cut to stubs.""" copy = embedded(manifest) parents = copy.parents if parents is None: From 98a66a5653326ee347d8092b79a6256c232a6d8b Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 23:41:28 +0000 Subject: [PATCH 077/278] Trim snapshot-export comments Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_snapshot_io.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_snapshot_io.py b/packages/kumo-engine/kumo/engine/_snapshot_io.py index 6c331811e..fdbf59b0a 100644 --- a/packages/kumo-engine/kumo/engine/_snapshot_io.py +++ b/packages/kumo-engine/kumo/engine/_snapshot_io.py @@ -235,8 +235,7 @@ def export(snapshot: Any, connector: Any, path: str) -> str: entry['rel_path']: int(entry['size_bytes']) for entry in json.loads(index_manifest)['files'] } - # Each export writes its own generation; the marker, written last, - # names the one that loads. Concurrent exports never share a file. + # One generation per export; the marker, written last, names it. generation = uuid.uuid4().hex filesystem, root = _location(connector, location) base = f'{root}/{generation}' @@ -412,7 +411,6 @@ def holds_export(index: Path, manifest: dict) -> bool: def _download(root: str, staging: Path, client: FileSystem, item: Tuple[str, int, int], stop: threading.Event) -> None: name, offset, length = item - # Each open is a HEAD: pyarrow cannot open a file from a known size. with client.open_input_file(f'{root}/index/{name}') as reader: fd = os.open(staging / name, os.O_WRONLY) try: From f606845055d485b5c992509a9301f25ef7e2aa5c Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 23:41:51 +0000 Subject: [PATCH 078/278] Fix the over-indented test signature flake8 flags Co-Authored-By: Claude Opus 5.5 --- packages/kumo-core/tests/test_chunked_fs.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/kumo-core/tests/test_chunked_fs.py b/packages/kumo-core/tests/test_chunked_fs.py index 2d06b60d2..8b4b444dc 100644 --- a/packages/kumo-core/tests/test_chunked_fs.py +++ b/packages/kumo-core/tests/test_chunked_fs.py @@ -137,8 +137,8 @@ def test_a_write_removes_the_other_form(store, first, second): @pytest.mark.parametrize('sized', [False, True]) @pytest.mark.parametrize('streams', [False, True]) @pytest.mark.parametrize('old', [3, 2 * PART]) -def test_a_failed_chunked_write_leaves_the_old_file(monkeypatch, streams, - sized, old): +def test_a_failed_chunked_write_leaves_the_old_file( + monkeypatch, streams, sized, old): from kumo_core.object_store_sink import open_sized handler = StreamingHandler() if streams else MemoryHandler() From 68875f1ee2231598f3441e0b759a8b6c51b5df53 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 23:41:55 +0000 Subject: [PATCH 079/278] Record a source's location wherever its backend can name it Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_provenance.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_provenance.py b/packages/kumo-engine/kumo/engine/_provenance.py index 7b440d2bf..2aea72b08 100644 --- a/packages/kumo-engine/kumo/engine/_provenance.py +++ b/packages/kumo-engine/kumo/engine/_provenance.py @@ -13,9 +13,6 @@ from kumo.engine._manifest import Environment, SourceTableEntry -_LOCATED_CONNECTORS = frozenset( - {'LOCAL_FILE', 'S3', 'GCS', 'SNOWFLAKE', 'DATABRICKS', 'BIGQUERY'}) - @cache def _build_commit(version: str) -> Optional[str]: @@ -68,9 +65,8 @@ def source_location(connector: Any, table_name: str, sql: Optional[str] = None) -> str: if sql is not None: return 'sql:' + hashlib.sha256(sql.encode()).hexdigest() - if _connector_kind(connector) not in _LOCATED_CONNECTORS: - return table_name - return connector.source_backend.table_location(connector, table_name) + locate = getattr(connector.source_backend, 'table_location', None) + return table_name if locate is None else locate(connector, table_name) def _connector_kind(connector: Any) -> str: From fdd6dc9935a9f58a4c2df71c42d925ab70e8d19c Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 23:42:25 +0000 Subject: [PATCH 080/278] Expect volume chunks under a generation directory Co-Authored-By: Claude Opus 5.5 --- packages/kumo-databricks/README.md | 6 +++--- packages/kumo-databricks/tests/test_volume.py | 3 ++- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/packages/kumo-databricks/README.md b/packages/kumo-databricks/README.md index 865ae85a1..03f99568d 100644 --- a/packages/kumo-databricks/README.md +++ b/packages/kumo-databricks/README.md @@ -89,9 +89,9 @@ For an MLOps artifact root on a volume, see `packages/kumo-engine/docs/mlops.md` A file whose size is known when it is written (bundle files, exported Parquet parts) is stored as one object up to 4 GiB, so Spark and other tools read it directly. A larger one streams into -`.chunks/00000`, `00001`, ... (4 GiB each, under the Files API's 5 GiB -upload cap) with an empty `_SUCCESS` written last, and Kumo reads the parts -back as one file. A file written without a known size stays one object up to +`.chunks//00000`, `00001`, ... (4 GiB each, under the Files +API's 5 GiB upload cap) with an empty `_SUCCESS` written last, and Kumo reads +the newest committed generation back as one file. A file written without a known size stays one object up to 256 MiB, staged on local disk under the engine scratch directory, and streams into parts above that. The identity needs `USE CATALOG`, `USE SCHEMA`, and `READ VOLUME` (loads) and `WRITE VOLUME` (exports) on the volume; no diff --git a/packages/kumo-databricks/tests/test_volume.py b/packages/kumo-databricks/tests/test_volume.py index db60e65fd..25cb8fc05 100644 --- a/packages/kumo-databricks/tests/test_volume.py +++ b/packages/kumo-databricks/tests/test_volume.py @@ -51,7 +51,8 @@ def test_sized_writes_stay_plain_up_to_a_part(): assert reader.readall() == body assert volume.objects[f'{ROOT}/plain'] == BODY[:1000] assert f'{ROOT}/parts' not in volume.objects - assert f'{ROOT}/parts.chunks/_SUCCESS' in volume.objects + assert any(path.startswith(f'{ROOT}/parts.chunks/') + and path.endswith('/_SUCCESS') for path in volume.objects) def test_a_failed_upload_raises_without_the_signed_url(): From 04f02f1c204881d59081cadcd6ad2144adb6c1f9 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 23:45:14 +0000 Subject: [PATCH 081/278] Abort, rather than leave open, a stream whose write failed Co-Authored-By: Claude Opus 5.5 --- .../kumo-core/src/kumo_core/chunked_fs.py | 16 ++++++++++---- packages/kumo-core/tests/test_chunked_fs.py | 21 +++++++++++++++++++ 2 files changed, 33 insertions(+), 4 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/chunked_fs.py b/packages/kumo-core/src/kumo_core/chunked_fs.py index 67c6df6ed..2124096ca 100644 --- a/packages/kumo-core/src/kumo_core/chunked_fs.py +++ b/packages/kumo-core/src/kumo_core/chunked_fs.py @@ -59,7 +59,7 @@ class ChunkedFileHandler(FileSystemHandler): plain_bytes = PART_BYTES def _stream(self, path: str) -> BinaryIO: - r"""Optional: a writable whose ``close`` commits ``path``.""" + r"""Optional: a writable committed by ``close``, dropped by ``abort``.""" raise NotImplementedError def _size(self, path: str) -> Optional[int]: @@ -301,6 +301,15 @@ def close(self): super().close() +def _abandon(stream: Optional[BinaryIO], close: bool = False) -> None: + if stream is None: + return + abort = getattr(stream, 'abort', stream.close if close else None) + if abort is not None: + with contextlib.suppress(Exception): + abort() + + class _Writer(io.RawIOBase): def __init__(self, handler: ChunkedFileHandler, path: str, size: Optional[int] = None) -> None: @@ -413,9 +422,8 @@ def close(self): return try: if self._failed: - with contextlib.suppress(Exception): - if self._part is not None: - self._part.close() + _abandon(self._plain) + _abandon(self._part, close=True) elif self._plain is not None: self._plain.close() self._plain = None diff --git a/packages/kumo-core/tests/test_chunked_fs.py b/packages/kumo-core/tests/test_chunked_fs.py index 8b4b444dc..4fa2ee3f0 100644 --- a/packages/kumo-core/tests/test_chunked_fs.py +++ b/packages/kumo-core/tests/test_chunked_fs.py @@ -167,6 +167,27 @@ def second_part_fails(*args): ] +def test_a_failed_plain_stream_is_aborted_not_committed(monkeypatch): + from kumo_core.object_store_sink import open_sized + + def fail(self, data): + raise OSError('upload failed') + + def abort(self): + aborted.append(self.path) + io.BytesIO.close(self) + + handler = StreamingHandler() + aborted = [] + monkeypatch.setattr(_Upload, 'write', fail) + monkeypatch.setattr(_Upload, 'abort', abort, raising=False) + with pytest.raises(OSError, match='upload failed'): + with open_sized(PyFileSystem(handler), 'root/a.bin', 3) as writer: + writer.write(b'abc') + assert aborted == ['root/a.bin'] + assert handler.objects == {} + + def test_the_newest_committed_generation_is_read(store): handler, fs = store _write(fs, 'root/a.bin', b'1' * (2 * PART)) From 1ef521a3e932f13b9bd92c4da4630a3dc237b2ba Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 23:41:37 +0000 Subject: [PATCH 082/278] Resolve the MLflow experiment in one place, so log_model creates a missing one log_model(experiment_name=...) logged into the Default experiment when the named one did not exist. The three copies of the HTTP-limit set share a helper. Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/integrations/mlflow.py | 46 ++++++++----------- 1 file changed, 18 insertions(+), 28 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index 657e63f1a..cedad020d 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -197,8 +197,7 @@ def check(self, registry: Optional[ModelRegistryLocation] = None) -> None: if registry is not None: self.check_model_name(registry.model_name) uri = self.target.tracking_uri - owned = frozenset(name for name in _HTTP_LIMITS - if name not in os.environ) + owned = _unset_limits() missing = False try: with _limited_http(owned), self.target.scope(): @@ -240,6 +239,13 @@ def prepare(self) -> None: self._outer = (active.info.run_id if active is not None and same else None) + def _experiment_for(self, client: Any) -> str: + if self.experiment_name is None: + return self._default_experiment or '0' + experiment = client.get_experiment_by_name(self.experiment_name) + return (experiment.experiment_id if experiment else + client.create_experiment(self.experiment_name)) + def _require_experiment(self) -> None: if self._named and self.experiment_name is None: raise ValueError('MlflowTracking(target=...) records runs in an ' @@ -269,8 +275,7 @@ def register(self, result: Any, self._require_experiment() client = self.target.client() pointers = registry_pointers(result.id) - with _limited_http(frozenset(name for name in _HTTP_LIMITS - if name not in os.environ)): + with _limited_http(_unset_limits()): with self.target.scope(): held = _registered_version(client, location.model_name, pointers) @@ -280,14 +285,7 @@ def register(self, result: Any, snapshot = GraphSnapshot(spec.snapshot_id, Path(), spec.graph_config or {}) with self.target.scope(): - experiment_id = self._default_experiment or '0' - if self.experiment_name is not None: - experiment = client.get_experiment_by_name( - self.experiment_name) - experiment_id = (experiment.experiment_id - if experiment else - client.create_experiment( - self.experiment_name)) + experiment_id = self._experiment_for(client) run_id = client.create_run( experiment_id, tags={'kumo.bundle_id': str(result.id)}).info.run_id @@ -304,8 +302,7 @@ def start(self, record: RunRecord) -> None: self._require_experiment() self._reset() self.run_id = None - self._http_owned = frozenset(name for name in _HTTP_LIMITS - if name not in os.environ) + self._http_owned = _unset_limits() self._guard.enter_context(_limited_http(self._http_owned)) try: autologging = import_provider('mlflow.utils.autologging_utils', @@ -320,18 +317,11 @@ def start(self, record: RunRecord) -> None: context = import_provider('mlflow.tracking.context.registry', 'mlflow') with self.target.scope(): - experiment_id = self._default_experiment - if self.experiment_name is not None: - experiment = client.get_experiment_by_name( - self.experiment_name) - experiment_id = (experiment.experiment_id if experiment else - client.create_experiment( - self.experiment_name)) parent = {} if self._outer is None else { 'mlflow.parentRunId': self._outer } run = client.create_run( - experiment_id or '0', + self._experiment_for(client), run_name=(self.run_name if self.run_name and record.tags.get('kumo.kind') == 'fit' else record.name), @@ -732,7 +722,7 @@ def log_model( raise ValueError('log_model requires a fitted or loaded Trainer') tracking = MlflowTracking(experiment_name, target=target) target = tracking.target - if run_id is None and tracking._named and experiment_name is None: + if run_id is None: tracking._require_experiment() client = target.client() active = mlflow.active_run() @@ -742,12 +732,8 @@ def log_model( created = run_id is None with target.scope(): if created: - experiment = (client.get_experiment_by_name( - tracking.experiment_name) if tracking.experiment_name else - None) run_id = client.create_run( - experiment.experiment_id if experiment else - tracking._default_experiment or '0').info.run_id + tracking._experiment_for(client)).info.run_id experiment_id = client.get_run(run_id).info.experiment_id _log_inputs(mlflow, client, run_id, fit_inputs(result.manifest)) info = _log_model(target, client, result.bundle_dir, result.manifest or @@ -1085,6 +1071,10 @@ def _quietly(function: Callable[..., Any], *args: Any) -> None: logger.debug('MLflow cleanup call failed', exc_info=True) +def _unset_limits() -> FrozenSet[str]: + return frozenset(name for name in _HTTP_LIMITS if name not in os.environ) + + @contextmanager def _limited_http(names: FrozenSet[str]) -> Iterator[None]: os.environ.update({name: _HTTP_LIMITS[name] for name in names}) From 73b3d2423f81550be1c6ab520874df6b681bb14d Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 23:41:44 +0000 Subject: [PATCH 083/278] Check a temporal query through _temporal instead of a copy of it Co-Authored-By: Claude Opus 5.5 --- .../kumo-engine/kumo/engine/integrations/_mlflow_model.py | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py b/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py index dd805bfa3..a2c2e4ba1 100644 --- a/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py +++ b/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py @@ -138,15 +138,9 @@ def _seconds(values: pd.Series) -> pd.Series: def _require_inputs(trainer: Any, snapshot: GraphSnapshot, rows: pd.DataFrame, time: Optional[pd.Timestamp]) -> None: - from kumo_core.pquery import query_is_temporal - - from kumo.engine.pquery import PredictiveQuery - if len(rows) == 0 or time is not None or 'TIMESTAMP' in rows: return - parsed = PredictiveQuery._from_snapshot(snapshot, - trainer._training_query)._parsed - if query_is_temporal(parsed, snapshot._native_graph): + if _temporal(trainer, snapshot): raise ValueError("temporal query: pass params={'prediction_time': " "...} or a TIMESTAMP column") From 7293b0f4a43c57094dde9f4cea286335ad844e57 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 23:41:53 +0000 Subject: [PATCH 084/278] Build a snapshot's URI through MLOps.snapshot_uri Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_mlops.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/kumo-engine/kumo/engine/_mlops.py b/packages/kumo-engine/kumo/engine/_mlops.py index 6a6814f98..ff0ba7224 100644 --- a/packages/kumo-engine/kumo/engine/_mlops.py +++ b/packages/kumo-engine/kumo/engine/_mlops.py @@ -309,4 +309,4 @@ def snapshot_uri(snapshot: Any, 'A snapshot given by object or id is read from the artifact ' 'root: call this in a kumo.mlops block, or pass the location ' 'snapshot.export(...) returned.') - return settings.root.uri(snapshot_path(snapshot)) + return settings.snapshot_uri(snapshot) From f9e79deac651b6a176f73e2375116a23c110c592 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 23:41:53 +0000 Subject: [PATCH 085/278] Read ModelRegistryLocation.alias directly in the Databricks registry Co-Authored-By: Claude Opus 5.5 --- packages/kumo-databricks/src/kumo_databricks/registry.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/kumo-databricks/src/kumo_databricks/registry.py b/packages/kumo-databricks/src/kumo_databricks/registry.py index d51e45474..9c7f9e3ab 100644 --- a/packages/kumo-databricks/src/kumo_databricks/registry.py +++ b/packages/kumo-databricks/src/kumo_databricks/registry.py @@ -386,7 +386,7 @@ def fetch_model( _version_number(location) return fetch_bundle(self._client(connector), name, Path(local_dir).resolve(), version=location.version, - alias=getattr(location, 'alias', None)) + alias=location.alias) def set_model_alias(self, connector: SourceConnector, location: ModelRegistryLocation, alias: str) -> None: From deb6eb86ee1fbadc5e521956b55446a55b3fafe0 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 23:42:15 +0000 Subject: [PATCH 086/278] Shorten docstrings in the summaries modules Co-Authored-By: Claude Opus 5.5 --- packages/kumo-core/src/kumo_core/summaries.py | 7 +----- .../kumo/engine/_data_summaries.py | 22 +++++-------------- 2 files changed, 7 insertions(+), 22 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/summaries.py b/packages/kumo-core/src/kumo_core/summaries.py index 2d7906f58..bbb9e05dd 100644 --- a/packages/kumo-core/src/kumo_core/summaries.py +++ b/packages/kumo-core/src/kumo_core/summaries.py @@ -1,12 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -r"""Small typed results of an evaluation or a data summary, named -``
/`` (``test/roc``, ``pred/True_PROB/hist``, ``graph/rows``). - -They know nothing about experiment trackers: each tracker maps these types to -its own calls, so another rendering (a figure, a UI) is one more mapping. -""" +r"""Typed, tracker-neutral results, named ``
/``.""" from dataclasses import dataclass from typing import Any, Mapping, Optional, Tuple, Union diff --git a/packages/kumo-engine/kumo/engine/_data_summaries.py b/packages/kumo-engine/kumo/engine/_data_summaries.py index 2e2fe052e..8a589f003 100644 --- a/packages/kumo-engine/kumo/engine/_data_summaries.py +++ b/packages/kumo-engine/kumo/engine/_data_summaries.py @@ -1,9 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -r"""Summaries of a training table's splits and of a graph snapshot, read -from the statistics their builds wrote (``split_stats.json``, -``col_stats.json``, ``ingest_stats.json``), with no pass over the data.""" +r"""Split and snapshot summaries, from the stats files builds wrote.""" import json import logging @@ -34,8 +32,7 @@ def split_summaries(staging_dir: Path) -> List[Summary]: - r"""``data//rows`` and ``entities`` per split, and its label - balance: ``positive_rate``, or the label's mean, median, min and max.""" + r"""Rows, entities and label balance per split.""" stats = _read(staging_dir / 'split_stats.json') summaries: List[Summary] = [] for split in _SPLITS: @@ -63,8 +60,7 @@ def split_summaries(staging_dir: Path) -> List[Summary]: def snapshot_summaries(index: Path, graph_config: Mapping[str, Any]) -> List[Summary]: - r"""The snapshot's tables and relations (`_graph_scalars`), their data - quality (`_quality`) and ``graph/col_stats`` with a row per column.""" + r"""The snapshot's tables, relations, data quality and column stats.""" ingest = _read(index / 'ingest_stats.json').get('stats') or {} columns: Dict[str, Dict[str, Any]] = _read(index / 'col_stats.json') relations = _relations(_read(index / 'edge_stats.json'), ingest) @@ -90,10 +86,7 @@ def _profiles(columns: Mapping[str, Mapping[str, Any]] def _relations(edge_stats: Mapping[str, Any], ingest: Mapping[str, Any]) -> Dict[str, Dict[str, float]]: - r"""Each relation's ``count``, ``src_match_frac`` and ``dst_match_frac`` - (the share of source and destination rows with an edge) and - ``fk_lookup_failures``, under DiskGraph's ``__to__`` - name.""" + r"""Each relation's counts and match fractions, by DiskGraph name.""" relations: Dict[str, Dict[str, float]] = { name: {'count': count} for name, count in (ingest.get('edge_counts') or {}).items() @@ -133,9 +126,7 @@ def _graph_scalars(ingest: Mapping[str, Any], def _quality(ingest: Mapping[str, Any], columns: Mapping[str, Any], relations: Mapping[str, Mapping[str, float]], graph_config: Mapping[str, Any]) -> List[Summary]: - r"""``graph/
/{duplicate_pks,dropped_rows,null_time_dropped}`` and, - when any table breaks the Kumo platform's data checks, the tag - ``kumo.data_warnings`` naming each.""" + r"""Per-table data quality, and ``kumo.data_warnings`` if any fail.""" duplicates = _by_table(ingest.get('pk_duplicate_counts'), ':') null_time = _by_table(ingest.get('pk_null_timestamp_drops'), ':shard-') dropped = ingest.get('dropped_rows_counts') or {} @@ -195,8 +186,7 @@ def _quality(ingest: Mapping[str, Any], columns: Mapping[str, Any], def table_details(index: Path) -> Details: - r"""Each source table's schema and profile, its rows and column - statistics, for its `Dataset`.""" + r"""Each source table's schema and profile, for its `Dataset`.""" rows = (_read(index / 'ingest_stats.json').get('stats') or {}).get('node_counts') or {} profiles = _profiles(_read(index / 'col_stats.json')) From 9771572994c1ee651a004122377594126ae2cd13 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 23:53:07 +0000 Subject: [PATCH 087/278] Test the S3 sink's delete with a fake session instead of moto Co-Authored-By: Claude Opus 5.5 --- packages/kumo-s3/tests/test_sink_delete.py | 24 +++++++++++----------- 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/packages/kumo-s3/tests/test_sink_delete.py b/packages/kumo-s3/tests/test_sink_delete.py index b3f97e114..cc727db78 100644 --- a/packages/kumo-s3/tests/test_sink_delete.py +++ b/packages/kumo-s3/tests/test_sink_delete.py @@ -1,22 +1,22 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -import boto3 +from types import SimpleNamespace + +from fakes import FakeSession from kumo_s3.sink import S3SinkBackend -from moto import mock_s3 -@mock_s3 -def test_deleting_the_last_object_under_a_prefix_leaves_no_marker( - monkeypatch): - client = boto3.client('s3', region_name='us-east-1') - client.create_bucket(Bucket='bucket') - client.put_object(Bucket='bucket', Key='root/.probe', Body=b'probe') +def test_deleting_an_object_deletes_only_that_key(monkeypatch): + calls = [] + session = FakeSession(region_name='us-east-1') + session.client = lambda service, endpoint_url=None: SimpleNamespace( + delete_object=lambda **kwargs: calls.append((service, endpoint_url, + kwargs))) sink = S3SinkBackend() - monkeypatch.setattr( - sink, '_resolved_session', lambda connector, bucket: - (boto3.Session(region_name='us-east-1'), True, None)) + monkeypatch.setattr(sink, '_resolved_session', + lambda connector, bucket: (session, True, None)) sink.delete_file(None, 'bucket/root/.probe') - assert client.list_objects_v2(Bucket='bucket')['KeyCount'] == 0 + assert calls == [('s3', None, {'Bucket': 'bucket', 'Key': 'root/.probe'})] From be32bb469432052489cf70378257846dbbea4bfd Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 23:52:35 +0000 Subject: [PATCH 088/278] Expect a single S3 object's provider version, not its modification time Co-Authored-By: Claude Opus 5.5 --- test/connector/test_file_connector.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/test/connector/test_file_connector.py b/test/connector/test_file_connector.py index dae3817e3..0410b7322 100644 --- a/test/connector/test_file_connector.py +++ b/test/connector/test_file_connector.py @@ -35,5 +35,7 @@ def test_get_data_version_s3_single_file(): version = file_connector.get_data_version('table') - # Epoch timestamp corresponding to original creation time of the file in s3 - assert version == '1663190838' + # A single object is versioned by its S3 VersionId, else its ETag. + assert version and version != '1663190838' + assert file_connector.get_data_version( + 'table', skip_cache=True, save_to_cache=False) == version From 6763e745eaffa9e52871f52f418173bd9938b150 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 00:00:47 +0000 Subject: [PATCH 089/278] Refuse an occupied output directory before scoring, in FilesystemSink A write into a directory that already holds files fails after scoring whatever the write mode, so the preflight refuses it up front. Co-Authored-By: Claude Opus 5.5 --- .../kumo-core/src/kumo_core/object_store_sink.py | 14 ++++++++++++++ packages/kumo-engine/test/test_filestore_sink.py | 11 +++++++++++ 2 files changed, 25 insertions(+) diff --git a/packages/kumo-core/src/kumo_core/object_store_sink.py b/packages/kumo-core/src/kumo_core/object_store_sink.py index 0d37c6d90..6be0a233c 100644 --- a/packages/kumo-core/src/kumo_core/object_store_sink.py +++ b/packages/kumo-core/src/kumo_core/object_store_sink.py @@ -397,3 +397,17 @@ def _destination(self, connector, coordinate): def _uri(self, path: str) -> str: return path + + def preflight_write_files(self, connector, table_name, write_mode=None): + """Refuse an occupied output directory before scoring.""" + del write_mode + from kumo_core.artifact_export import ArtifactDestinationConflict + from pyarrow.fs import FileSelector + + if not table_name.endswith('/'): + return + filesystem, path = self._sink.filesystem(connector, table_name) + if filesystem.get_file_info(FileSelector(path, allow_not_found=True)): + raise ArtifactDestinationConflict( + f"Output directory '{path}/' already contains data from " + f"another export. Choose an empty destination path.") diff --git a/packages/kumo-engine/test/test_filestore_sink.py b/packages/kumo-engine/test/test_filestore_sink.py index 30de65fdd..170a2d056 100644 --- a/packages/kumo-engine/test/test_filestore_sink.py +++ b/packages/kumo-engine/test/test_filestore_sink.py @@ -190,3 +190,14 @@ def test_parquet_writes_stay_plain_up_to_one_part(tmp_path, connector): assert files.write_files(connector, '@stage/out/', [str(part)]) assert pd.read_parquet('@stage/out/part-0.parquet', filesystem=PyFileSystem(store)).equals(frame) + + +def test_an_occupied_directory_is_refused_before_scoring(tmp_path, connector): + from kumo_core.artifact_export import ArtifactDestinationConflict + + files = FilesystemSink(connector.sink_backend) + files.preflight_write_files(connector, '@stage/out/') + connector.sink_backend.store.objects['@stage/out/old.parquet'] = b'x' + for mode in (None, 'OVERWRITE', 'APPEND'): + with pytest.raises(ArtifactDestinationConflict, match='already'): + files.preflight_write_files(connector, '@stage/out/', mode) From 85b8b06a633e8d05a01a72916e8d9d046231f355 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 23:40:03 +0000 Subject: [PATCH 090/278] Refuse an occupied volume directory whatever the write mode, through FilesystemSink Co-Authored-By: Claude Opus 5.5 --- .../kumo-databricks/src/kumo_databricks/sink.py | 17 +++-------------- packages/kumo-databricks/tests/test_volume.py | 8 ++++---- 2 files changed, 7 insertions(+), 18 deletions(-) diff --git a/packages/kumo-databricks/src/kumo_databricks/sink.py b/packages/kumo-databricks/src/kumo_databricks/sink.py index 56d7717fb..bdbdf7554 100644 --- a/packages/kumo-databricks/src/kumo_databricks/sink.py +++ b/packages/kumo-databricks/src/kumo_databricks/sink.py @@ -240,20 +240,9 @@ def fetch_model( def preflight_write_files(self, connector: SourceConnector, table_name: str, write_mode: Optional[WriteMode] = None) -> None: - """Refuse an occupied volume directory before scoring, unless - ``write_mode`` is ``OVERWRITE``; tables are not checked.""" - from kumo_core.artifact_export import ArtifactDestinationConflict - from pyarrow.fs import FileSelector - - if (write_mode is WriteMode.OVERWRITE or not table_name.endswith('/') - or not self.holds_files(connector, table_name)): - return - filesystem, path = self.filesystem(connector, table_name) - if filesystem.get_file_info(FileSelector(path, allow_not_found=True)): - raise ArtifactDestinationConflict( - f"Output directory '{path}/' already contains data from " - f"another export. Choose an empty destination path, or " - f"replace it with write_mode=WriteMode.OVERWRITE.") + """Refuse an occupied volume directory; tables are not checked.""" + if self.holds_files(connector, table_name): + FilesystemSink(self).preflight_write_files(connector, table_name) def holds_files(self, connector: SourceConnector, path: str) -> bool: """``/Volumes////...``; anything else is a diff --git a/packages/kumo-databricks/tests/test_volume.py b/packages/kumo-databricks/tests/test_volume.py index 25cb8fc05..15cef7d32 100644 --- a/packages/kumo-databricks/tests/test_volume.py +++ b/packages/kumo-databricks/tests/test_volume.py @@ -173,9 +173,9 @@ def test_an_occupied_volume_directory_is_refused_before_scoring(tmp_path): pq.write_table(pa.table({'users': [1]}), part) sink.write_files(SimpleNamespace(), target, [str(part)], export_id='r1') - with pytest.raises(ArtifactDestinationConflict, match='/predictions/'): - sink.preflight_write_files(SimpleNamespace(), target) - sink.preflight_write_files(SimpleNamespace(), target, - write_mode=WriteMode.OVERWRITE) + for mode in (None, WriteMode.OVERWRITE): + with pytest.raises(ArtifactDestinationConflict, match='/predictions/'): + sink.preflight_write_files(SimpleNamespace(), target, + write_mode=mode) sink.preflight_write_files(SimpleNamespace(), f'{ROOT}/other/') sink.preflight_write_files(SimpleNamespace(), 'main.ml.scores') From cceb78e4f80328d340e2d07f540f8492eda21b4d Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 23:40:03 +0000 Subject: [PATCH 091/278] Drop the unused write_files options passthrough Co-Authored-By: Claude Opus 5.5 --- packages/kumo-databricks/src/kumo_databricks/sink.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/packages/kumo-databricks/src/kumo_databricks/sink.py b/packages/kumo-databricks/src/kumo_databricks/sink.py index bdbdf7554..8cfd36711 100644 --- a/packages/kumo-databricks/src/kumo_databricks/sink.py +++ b/packages/kumo-databricks/src/kumo_databricks/sink.py @@ -162,13 +162,12 @@ def write_table( def write_files(self, connector, table_name, files, write_mode=WriteMode.OVERWRITE, project_id=None, *, - export_id=None, extra_columns=None, **options): + export_id=None, extra_columns=None): """Copy Parquet to a volume path, or COPY it into a table.""" if self.holds_files(connector, table_name): return FilesystemSink(self).write_files( connector, table_name, files, write_mode, - export_id=export_id, extra_columns=extra_columns, - **options) + export_id=export_id, extra_columns=extra_columns) del project_id, export_id from functools import partial From dc1405fa585e77e80322b1ef47fb37b5618c9dba Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 23:40:17 +0000 Subject: [PATCH 092/278] Shorten the volume sink docstrings Co-Authored-By: Claude Opus 5.5 --- packages/kumo-databricks/src/kumo_databricks/sink.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/packages/kumo-databricks/src/kumo_databricks/sink.py b/packages/kumo-databricks/src/kumo_databricks/sink.py index 8cfd36711..4b21d4d55 100644 --- a/packages/kumo-databricks/src/kumo_databricks/sink.py +++ b/packages/kumo-databricks/src/kumo_databricks/sink.py @@ -244,14 +244,13 @@ def preflight_write_files(self, connector: SourceConnector, FilesystemSink(self).preflight_write_files(connector, table_name) def holds_files(self, connector: SourceConnector, path: str) -> bool: - """``/Volumes////...``; anything else is a - table.""" + """Volume paths hold files; anything else is a table.""" return is_volume_path(path) def filesystem(self, connector: SourceConnector, path: str) -> tuple[FileSystem, str]: - """A volume path as a filesystem, as the connector's identity.""" - if not is_volume_path(path): + """A volume path's filesystem, as the connector's identity.""" + if not self.holds_files(connector, path): raise ValueError( f"{path!r} is not a Unity Catalog volume path; Databricks " f"artifacts go under /Volumes////.") From a7fdac952409191aba21fdccf9a0b48251827684 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 23:40:24 +0000 Subject: [PATCH 093/278] Tidy the volume handler: one-line docs, top-level imports, no custom equality Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_databricks/volume.py | 32 ++++++------------- 1 file changed, 9 insertions(+), 23 deletions(-) diff --git a/packages/kumo-databricks/src/kumo_databricks/volume.py b/packages/kumo-databricks/src/kumo_databricks/volume.py index 3a7b39e61..8d65c0b30 100644 --- a/packages/kumo-databricks/src/kumo_databricks/volume.py +++ b/packages/kumo-databricks/src/kumo_databricks/volume.py @@ -16,10 +16,12 @@ from databricks.sdk import WorkspaceClient from databricks.sdk.errors import DatabricksError, NotFound, PermissionDenied from kumo_core.chunked_fs import ChunkedFileHandler +from kumo_core.sources.diagnostics import safe_summary +from kumo_databricks.staging import clear_volume from pyarrow.fs import FileSystem, PyFileSystem _FILES = '/api/2.0/fs/files' -_READ = 8 << 20 +_BUFFER = 8 << 20 def is_volume_path(path: str) -> bool: @@ -34,10 +36,7 @@ def volume_filesystem(client: WorkspaceClient) -> FileSystem: @contextmanager def _translated(path: str) -> Iterator[None]: - """Nothing is chained, and only SDK errors keep their text: transport - errors from presigned uploads carry the signed URL.""" - from kumo_core.sources.diagnostics import safe_summary - + """OS errors, unchained; non-SDK text may hold a signed URL.""" try: yield except NotFound as error: @@ -51,23 +50,12 @@ def _translated(path: str) -> Iterator[None]: class VolumeHandler(ChunkedFileHandler): - """``/Volumes////...`` through ``client.files``. - - Parts stay under the Files API's 5 GiB single-upload cap. - """ - part_bytes = 4 << 30 + part_bytes = 4 << 30 # under the Files API's 5 GiB upload cap plain_bytes = 256 << 20 def __init__(self, client: WorkspaceClient) -> None: self._client = client - def __eq__(self, other: object) -> bool: - return (isinstance(other, VolumeHandler) - and other._client is self._client) - - def __ne__(self, other: object) -> bool: - return not self == other - def get_type_name(self) -> str: return 'databricks-volume' @@ -107,8 +95,6 @@ def _part(self, path: str, size: int) -> _Part: return _Part(self._client, path) def _remove(self, path: str, *, tree: bool) -> None: - from kumo_databricks.staging import clear_volume - path = path.rstrip('/') with _translated(path): if tree: @@ -135,13 +121,13 @@ def read_at(self, offset: int, length: int) -> bytes: with _translated(self._path): pieces, got = [], 0 while got < length: - piece = self._stream.read(min(_READ, length - got)) + piece = self._stream.read(min(_BUFFER, length - got)) if not piece: break pieces.append(piece) got += len(piece) self._at += got - return pieces[0] if len(pieces) == 1 else b''.join(pieces) + return b''.join(pieces) def _open(self, offset: int) -> None: self.close() @@ -158,7 +144,7 @@ def _open(self, offset: int) -> None: self.close() raise OSError(f'{self._path}: asked for bytes from {offset}, ' f'got {served or "the whole file"}') - self._stream.set_chunk_size(_READ) + self._stream.set_chunk_size(_BUFFER) self._at = offset def close(self) -> None: @@ -195,7 +181,7 @@ def __init__(self, client: WorkspaceClient, path: str) -> None: self._path = path self._feed = _Feed() self._error: Optional[BaseException] = None - source = io.BufferedReader(self._feed, 8 << 20) + source = io.BufferedReader(self._feed, _BUFFER) def upload() -> None: try: From 973dc8a1abc2b05d7bda43475c1f2776bffb6900 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 23:40:29 +0000 Subject: [PATCH 094/278] Show the chunk generation directory in the volume README Co-Authored-By: Claude Opus 5.5 --- packages/kumo-databricks/README.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/packages/kumo-databricks/README.md b/packages/kumo-databricks/README.md index 03f99568d..2c0a1a53f 100644 --- a/packages/kumo-databricks/README.md +++ b/packages/kumo-databricks/README.md @@ -91,7 +91,8 @@ Parquet parts) is stored as one object up to 4 GiB, so Spark and other tools read it directly. A larger one streams into `.chunks//00000`, `00001`, ... (4 GiB each, under the Files API's 5 GiB upload cap) with an empty `_SUCCESS` written last, and Kumo reads -the newest committed generation back as one file. A file written without a known size stays one object up to +the newest committed generation back as one file. A file written without a +known size stays one object up to 256 MiB, staged on local disk under the engine scratch directory, and streams into parts above that. The identity needs `USE CATALOG`, `USE SCHEMA`, and `READ VOLUME` (loads) and `WRITE VOLUME` (exports) on the volume; no From 53d23fc25c73c98b14193ff626a271e19f898d7b Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 00:03:23 +0000 Subject: [PATCH 095/278] Cancel a volume upload whose write failed instead of leaving it open Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_databricks/volume.py | 21 +++++++++++++++++-- packages/kumo-databricks/tests/test_volume.py | 1 + 2 files changed, 20 insertions(+), 2 deletions(-) diff --git a/packages/kumo-databricks/src/kumo_databricks/volume.py b/packages/kumo-databricks/src/kumo_databricks/volume.py index 8d65c0b30..ae7def3fe 100644 --- a/packages/kumo-databricks/src/kumo_databricks/volume.py +++ b/packages/kumo-databricks/src/kumo_databricks/volume.py @@ -10,7 +10,7 @@ import queue import threading from contextlib import contextmanager -from typing import BinaryIO, Dict, Iterator, Optional +from typing import Any, BinaryIO, Dict, Iterator, Optional from urllib.parse import quote from databricks.sdk import WorkspaceClient @@ -153,9 +153,12 @@ def close(self) -> None: self._stream = None +_CANCEL = object() + + class _Feed(io.RawIOBase): def __init__(self) -> None: - self.pieces: queue.Queue[Optional[bytes]] = queue.Queue(maxsize=4) + self.pieces: queue.Queue[Any] = queue.Queue(maxsize=4) self._pending = b'' self._done = False @@ -165,6 +168,8 @@ def readable(self) -> bool: def readinto(self, buffer) -> int: while not self._pending and not self._done: piece = self.pieces.get() + if piece is _CANCEL: + raise OSError('upload cancelled') if piece is None: self._done = True else: @@ -214,6 +219,18 @@ def _put(self, piece: Optional[bytes]) -> None: if piece is not None: raise OSError(f'{self._path}: the upload stopped reading') + def abort(self) -> None: + if self.closed: + return + while self._thread.is_alive(): + try: + self._feed.pieces.put(_CANCEL, timeout=1) + break + except queue.Full: + continue + self._thread.join() + super().close() + def close(self) -> None: if self.closed: return diff --git a/packages/kumo-databricks/tests/test_volume.py b/packages/kumo-databricks/tests/test_volume.py index 15cef7d32..c13650ffc 100644 --- a/packages/kumo-databricks/tests/test_volume.py +++ b/packages/kumo-databricks/tests/test_volume.py @@ -179,3 +179,4 @@ def test_an_occupied_volume_directory_is_refused_before_scoring(tmp_path): write_mode=mode) sink.preflight_write_files(SimpleNamespace(), f'{ROOT}/other/') sink.preflight_write_files(SimpleNamespace(), 'main.ml.scores') + From ea03c6d7fee45e42153fa5f81fed0415b5ab5645 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 00:03:35 +0000 Subject: [PATCH 096/278] Drop a trailing blank line flake8 flags in the volume tests Co-Authored-By: Claude Opus 5.5 --- packages/kumo-databricks/tests/test_volume.py | 1 - 1 file changed, 1 deletion(-) diff --git a/packages/kumo-databricks/tests/test_volume.py b/packages/kumo-databricks/tests/test_volume.py index c13650ffc..15cef7d32 100644 --- a/packages/kumo-databricks/tests/test_volume.py +++ b/packages/kumo-databricks/tests/test_volume.py @@ -179,4 +179,3 @@ def test_an_occupied_volume_directory_is_refused_before_scoring(tmp_path): write_mode=mode) sink.preflight_write_files(SimpleNamespace(), f'{ROOT}/other/') sink.preflight_write_files(SimpleNamespace(), 'main.ml.scores') - From 619fe197a38e598e1b7aec33dc6c3de0724dfbcb Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 01:08:20 +0000 Subject: [PATCH 097/278] Pass the write mode to the shared preflight for volume directories Co-Authored-By: Claude Opus 5.5 --- packages/kumo-databricks/src/kumo_databricks/sink.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/packages/kumo-databricks/src/kumo_databricks/sink.py b/packages/kumo-databricks/src/kumo_databricks/sink.py index 4b21d4d55..47734fddf 100644 --- a/packages/kumo-databricks/src/kumo_databricks/sink.py +++ b/packages/kumo-databricks/src/kumo_databricks/sink.py @@ -241,7 +241,8 @@ def preflight_write_files(self, connector: SourceConnector, write_mode: Optional[WriteMode] = None) -> None: """Refuse an occupied volume directory; tables are not checked.""" if self.holds_files(connector, table_name): - FilesystemSink(self).preflight_write_files(connector, table_name) + FilesystemSink(self).preflight_write_files( + connector, table_name, write_mode) def holds_files(self, connector: SourceConnector, path: str) -> bool: """Volume paths hold files; anything else is a table.""" From 7b3fb84173a146d50fa0d33160cdd0b808a24b4c Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 03:02:26 +0000 Subject: [PATCH 098/278] Give an ADLS source's version its kind: exact for one object, listing for a directory main's ADLS reader returned a plain string, so a snapshot read an ADLS table's source changes as unknown. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-adls/src/kumo_adls/reader.py | 9 +++++++-- packages/kumo-adls/tests/test_identity_listing.py | 11 +++++++++++ 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/packages/kumo-adls/src/kumo_adls/reader.py b/packages/kumo-adls/src/kumo_adls/reader.py index 88cf549dd..1b258d8a1 100644 --- a/packages/kumo-adls/src/kumo_adls/reader.py +++ b/packages/kumo-adls/src/kumo_adls/reader.py @@ -8,6 +8,7 @@ from typing import Literal from kumo_core.sources.csv_format import CsvOptions +from kumo_core.sources.identity import SourceVersion from kumo_core.sources.object_store import ObjectStoreSourceReader from kumo_core.types.enums import ConnectorType from kumoapi.adls import ( @@ -48,8 +49,12 @@ def _data_version(self, path, entries): zip((entry.path for entry in ordered), self._identities(ordered, None))) # Discovery and preparation spell the same directory with/without '/'. - return hashlib.sha256( - json.dumps((path.rstrip('/'), metadata)).encode()).hexdigest() + single = (len(ordered) == 1 + and ordered[0].path.rstrip('/') == path.rstrip('/')) + return SourceVersion( + hashlib.sha256( + json.dumps((path.rstrip('/'), metadata)).encode()).hexdigest(), + 'exact' if single else 'listing') def source_reader(client, account: str, suffix: str = 'dfs.core.windows.net', *, diff --git a/packages/kumo-adls/tests/test_identity_listing.py b/packages/kumo-adls/tests/test_identity_listing.py index 86d5fa248..a164fab46 100644 --- a/packages/kumo-adls/tests/test_identity_listing.py +++ b/packages/kumo-adls/tests/test_identity_listing.py @@ -304,3 +304,14 @@ def failed(*args, **kwargs): expected = ReconnectRequiredError if status == 401 else TransientSourceError with pytest.raises(expected): filesystem.object_identities(PATHS) + + +def test_versions_say_whether_they_name_one_object_or_a_listing( + listing_service): + from kumo_core.sources.identity import version_kind + + service, _ = listing_service + source = reader(ADLSFileSystem(service, 'companylake', + 'dfs.core.windows.net')) + assert version_kind(source.data_version('')) == 'listing' + assert version_kind(source.data_version('a.parquet')) == 'exact' From a3d88da2a7d520aa632497f2c6a95d355cb6eb37 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 03:49:57 +0000 Subject: [PATCH 099/278] Ship chunked_fs in the public test image, which object_store_sink now imports Co-Authored-By: Claude Opus 5.5 --- tests/public/Dockerfile.dockerignore | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/public/Dockerfile.dockerignore b/tests/public/Dockerfile.dockerignore index 54d2bca52..c6d21c61a 100644 --- a/tests/public/Dockerfile.dockerignore +++ b/tests/public/Dockerfile.dockerignore @@ -19,6 +19,7 @@ !packages/kumo-core/src/kumo_core/io/__init__.py !packages/kumo-core/src/kumo_core/io/filesystem.py !packages/kumo-core/src/kumo_core/io/json.py +!packages/kumo-core/src/kumo_core/chunked_fs.py !packages/kumo-core/src/kumo_core/object_store_sink.py !packages/kumo-core/src/kumo_core/artifact_export.py !packages/kumo-core/src/kumo_core/stype_inference.py From bd8fd5c092f8ff456a64f6390e5ff1aec20a82af Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 21:46:08 +0000 Subject: [PATCH 100/278] Import FitMode at module level in the engine trainer Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/trainer/trainer.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 28c3156a3..ea8adb8e4 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -34,6 +34,7 @@ resolve_query_forecast, ) from kumo_core.train.artifact_manifest import PredictionSelectionRef +from kumo_core.training.config import FitMode from kumo_core.training.explain_aggregation import aggregate_subgraph_features from kumo_core.training.workflow import TrainingWorkflowRequest from kumo_core.types import ExplainMethod, PredictionArtifactType @@ -189,8 +190,6 @@ def _find_holdout(log_dir: str) -> tuple[Optional[Path], Optional[bool]]: def _published_holdout(trainer: Any) -> tuple[Optional[Path], Optional[bool]]: - from kumo_core.training.config import FitMode - published = trainer.published_fit_mode() if published is FitMode.refit_full: logger.info('The published model is refit on the full data, including ' From 64d9dce768669faf44b5da59b197fe7e2b10dd2b Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 21:46:25 +0000 Subject: [PATCH 101/278] Share the refit suffix step in Trainer.get_log_dir Co-Authored-By: Claude Opus 5.5 --- .../kumo-core/src/kumo_core/train/trainer.py | 25 +++++++------------ 1 file changed, 9 insertions(+), 16 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/train/trainer.py b/packages/kumo-core/src/kumo_core/train/trainer.py index efe6968f2..c26afd9b7 100644 --- a/packages/kumo-core/src/kumo_core/train/trainer.py +++ b/packages/kumo-core/src/kumo_core/train/trainer.py @@ -1521,26 +1521,19 @@ def get_log_dir( fit_mode = self.fit_mode if fit_mode is None else fit_mode path = self.root_dir if not local else self.local_root_dir + refit = not original and fit_mode.is_refit # In case the `root_dir` contains a trainer config, it is a log dir # for a previous model version. if fsutil.exists(osp.join(path, 'config.json')): - if not original and fit_mode.is_refit: - if path[-1] == '/': - path = path[:-1] - path = f"{path}_{fit_mode.value}" - fsutil.makedirs(path, exist_ok=True) + if not refit: return path - else: - return path - - if self.name is not None: - path = osp.join(path, self.name) - if self.track_version: - path = osp.join(path, self.version_dir_name) - if not original and fit_mode.is_refit: - if path[-1] == '/': - path = path[:-1] - path = f"{path}_{fit_mode.value}" + else: + if self.name is not None: + path = osp.join(path, self.name) + if self.track_version: + path = osp.join(path, self.version_dir_name) + if refit: + path = f"{path.removesuffix('/')}_{fit_mode.value}" fsutil.makedirs(path, exist_ok=True) return path From bfb44972d277339089c48121b044e7a6d92b2208 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 21:47:33 +0000 Subject: [PATCH 102/278] Name the provenance file PROVENANCE_FILE everywhere instead of three aliases Co-Authored-By: Claude Opus 5.5 --- .../kumo-engine/kumo/engine/_snapshot_io.py | 21 ++++++------------- packages/kumo-engine/kumo/engine/graph.py | 6 ++---- .../kumo/engine/pquery/training_table.py | 12 +++++------ .../test/test_snapshot_manifest.py | 13 ++++++------ 4 files changed, 20 insertions(+), 32 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_snapshot_io.py b/packages/kumo-engine/kumo/engine/_snapshot_io.py index fdbf59b0a..45bc4dd98 100644 --- a/packages/kumo-engine/kumo/engine/_snapshot_io.py +++ b/packages/kumo-engine/kumo/engine/_snapshot_io.py @@ -15,6 +15,7 @@ from pathlib import Path from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple +from kumo_core.artifacts import PROVENANCE_FILE from pyarrow.fs import FileSystem _FORMAT = 1 @@ -175,10 +176,8 @@ def _snapshot_manifest(snapshot: Any) -> bytes: SnapshotParents, manifest_bytes, ) - from kumo.engine.graph import SNAPSHOT_MANIFEST_FILE - try: - return (Path(snapshot.index_path) / SNAPSHOT_MANIFEST_FILE).read_bytes() + return (Path(snapshot.index_path) / PROVENANCE_FILE).read_bytes() except OSError: return manifest_bytes( SnapshotManifest(id=snapshot.id, @@ -219,8 +218,6 @@ def export(snapshot: Any, connector: Any, path: str) -> str: from diskgraph import snapshot_range_checksums from kumo.engine._provenance import environment - from kumo.engine.graph import SNAPSHOT_MANIFEST_FILE - source = Path(snapshot.index_path) graph_key, data_version = snapshot._graph_key, snapshot._data_version location = export_location(snapshot, path) @@ -260,7 +257,7 @@ def upload(client: FileSystem, name: str, stop: threading.Event) -> None: [sizes[name] for name in names], upload, 'export') crc32 = checksums.result() _write(filesystem, f'{base}/index/{_INDEX_MANIFEST}', index_manifest) - _write(filesystem, f'{base}/{SNAPSHOT_MANIFEST_FILE}', snapshot_manifest) + _write(filesystem, f'{base}/{PROVENANCE_FILE}', snapshot_manifest) marker = { 'format': _FORMAT, 'snapshot_id': snapshot.id, @@ -286,8 +283,6 @@ def _holds_copy(connector: Any, location: str, snapshot_id: str, r"""Whether ``location`` holds this index's bytes and manifest.""" from kumo_core.object_store_sink import file_sizes - from kumo.engine.graph import SNAPSHOT_MANIFEST_FILE - try: root, manifest = read(connector, location) if (manifest.get('snapshot_id') != snapshot_id @@ -302,7 +297,7 @@ def _holds_copy(connector: Any, location: str, snapshot_id: str, entry['size_bytes']) for entry in listed): return False with filesystem.open_input_stream( - f'{root}/{SNAPSHOT_MANIFEST_FILE}') as reader: + f'{root}/{PROVENANCE_FILE}') as reader: return reader.readall() == snapshot_manifest except Exception: return False @@ -431,8 +426,6 @@ def _index_files( ) -> Tuple[bytes, List[Tuple[str, int]], Optional[bytes]]: from kumo_core.sources.layout import confined - from kumo.engine.graph import SNAPSHOT_MANIFEST_FILE - filesystem, _ = _location(connector, path) with filesystem.open_input_stream( f'{root}/index/{_INDEX_MANIFEST}') as reader: @@ -448,7 +441,7 @@ def _index_files( raise _corrupt(path, f'{_INDEX_MANIFEST} lists a file twice') try: with filesystem.open_input_stream( - f'{root}/{SNAPSHOT_MANIFEST_FILE}') as reader: + f'{root}/{PROVENANCE_FILE}') as reader: snapshot_manifest = reader.readall() except FileNotFoundError: snapshot_manifest = None @@ -469,8 +462,6 @@ def load_into(connector: Any, path: str, root: str, manifest: dict, from kumo_core.io.filesystem import fsync_directory from kumo_core.sources.transfer import discard_destination - from kumo.engine.graph import SNAPSHOT_MANIFEST_FILE - span = manifest['range_bytes'] index_manifest, files, snapshot_manifest = _index_files( path, root, manifest, connector) @@ -520,7 +511,7 @@ def load_into(connector: Any, path: str, root: str, manifest: dict, _check_readable(staging, manifest) (staging / _EXPORT_IDENTITY).write_text(export_identity(manifest)) if snapshot_manifest is not None: - (staging / SNAPSHOT_MANIFEST_FILE).write_bytes(snapshot_manifest) + (staging / PROVENANCE_FILE).write_bytes(snapshot_manifest) for directory in sorted({(staging / name).parent for name, _ in files} | {staging}): fsync_directory(directory) diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index 8c21e3cf9..34bd20707 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -93,8 +93,6 @@ from kumo.engine.trainer.output_config import SnapshotOutputConfig from kumo.engine.table import Table -SNAPSHOT_MANIFEST_FILE = PROVENANCE_FILE - _warned_scratch: set[Path] = set() _warned_backfill: set[Path] = set() @@ -509,14 +507,14 @@ def source_changes(self, graph: 'Graph') -> Dict[str, Dict[str, Any]]: @property def manifest(self) -> SnapshotManifest | None: """The snapshot manifest stored next to the index, if any.""" - return read_manifest(Path(self.index_path) / SNAPSHOT_MANIFEST_FILE, + return read_manifest(Path(self.index_path) / PROVENANCE_FILE, SnapshotManifest) def _write_snapshot_manifest(result: Any, snapshot_id: str, source_tables: Dict[str, SourceTableEntry], record_sources: bool, *, fresh: bool) -> None: - path = Path(result.index_path) / SNAPSHOT_MANIFEST_FILE + path = Path(result.index_path) / PROVENANCE_FILE try: if not fresh and path.exists(): existing = read_manifest(path, SnapshotManifest) diff --git a/packages/kumo-engine/kumo/engine/pquery/training_table.py b/packages/kumo-engine/kumo/engine/pquery/training_table.py index acb181d7a..1ce3cc828 100644 --- a/packages/kumo-engine/kumo/engine/pquery/training_table.py +++ b/packages/kumo-engine/kumo/engine/pquery/training_table.py @@ -30,8 +30,6 @@ if TYPE_CHECKING: from kumo.engine.trainer.output_config import TrainingTableOutputConfig -STAGING_MANIFEST_FILE = PROVENANCE_FILE -EXPORT_MANIFEST_FILE = PROVENANCE_FILE _RANDOM_SPLITS = ('RandomSplit', 'RandomLinkSplit') @@ -68,7 +66,7 @@ def tt_id(self) -> str: @property def manifest(self) -> Optional[TrainingTableManifest]: """Lineage recorded at generation; None for older tables.""" - return read_manifest(self.staging_dir / STAGING_MANIFEST_FILE, + return read_manifest(self.staging_dir / PROVENANCE_FILE, TrainingTableManifest) def export(self, output_config: "TrainingTableOutputConfig") -> str: @@ -105,17 +103,17 @@ def _export_sidecars(self, connector, output_config) -> Dict[str, Any]: from kumo.engine.connector import is_file_destination - manifest = self.staging_dir / STAGING_MANIFEST_FILE + manifest = self.staging_dir / PROVENANCE_FILE name = output_config.table_name.rstrip('/') try: if (not is_file_destination(connector, output_config.table_name) or layout.file_type(name) == layout.CSV or not manifest.is_file()): return {} - return {'sidecars': {EXPORT_MANIFEST_FILE: manifest.read_bytes()}} + return {'sidecars': {PROVENANCE_FILE: manifest.read_bytes()}} except Exception as error: warnings.warn( - f"Exporting without {EXPORT_MANIFEST_FILE}: {error}", + f"Exporting without {PROVENANCE_FILE}: {error}", UserWarning, stacklevel=3) return {} @@ -340,7 +338,7 @@ def _record_manifest(table: TrainingTable, snapshot: GraphSnapshot, source_table: Optional[SourceTableEntry] = None) -> None: from kumo.engine._manifest import embedded, stub - path = table.staging_dir / STAGING_MANIFEST_FILE + path = table.staging_dir / PROVENANCE_FILE try: from kumo_core.io.json import to_json_dict diff --git a/packages/kumo-engine/test/test_snapshot_manifest.py b/packages/kumo-engine/test/test_snapshot_manifest.py index 58ad7eab0..192128c50 100644 --- a/packages/kumo-engine/test/test_snapshot_manifest.py +++ b/packages/kumo-engine/test/test_snapshot_manifest.py @@ -12,6 +12,7 @@ import numpy as np import pandas as pd import pytest +from kumo_core.artifacts import PROVENANCE_FILE from kumo_core.custom_sql import custom_sql_synthetic_name from kumo_core.sources.identity import SourceVersion from kumo_core.types import ConnectorType @@ -39,7 +40,7 @@ DatabricksIdentity, SnowflakeIdentity, ) -from kumo.engine.graph import SNAPSHOT_MANIFEST_FILE, GraphSnapshot +from kumo.engine.graph import GraphSnapshot @pytest.fixture @@ -80,7 +81,7 @@ def test_snapshot_manifest_records_sources_and_diffs(make_snapshot, manifest = before.manifest assert manifest.kind == 'snapshot' - assert SNAPSHOT_MANIFEST_FILE == '_provenance.json' + assert PROVENANCE_FILE == '_provenance.json' assert manifest.id == before.id data = (tmp_path / 'data').resolve() assert {a: e.location @@ -111,7 +112,7 @@ def test_index_without_manifest_is_reused_and_backfilled( snapshot = make_snapshot(30) original = snapshot.manifest - (snapshot.index_path / SNAPSHOT_MANIFEST_FILE).unlink() + (snapshot.index_path / PROVENANCE_FILE).unlink() monkeypatch.setattr(ingest, 'execute_diskgraph_ingest', lambda *args: pytest.fail('cached index re-ingested')) @@ -141,7 +142,7 @@ def fail(path, manifest): ]) def test_unsupported_manifest_is_ignored(tmp_path, content): if content is not None: - (tmp_path / SNAPSHOT_MANIFEST_FILE).write_text(content) + (tmp_path / PROVENANCE_FILE).write_text(content) with pytest.warns(UserWarning): assert GraphSnapshot('s', tmp_path, {}).manifest is None else: @@ -229,7 +230,7 @@ def test_record_sources_false_keeps_table_names_only(make_snapshot, def test_toggling_record_sources_keeps_unknown_parent_fields(make_snapshot): - written = make_snapshot(30).index_path / SNAPSHOT_MANIFEST_FILE + written = make_snapshot(30).index_path / PROVENANCE_FILE data = json.loads(written.read_text()) data['parents']['future'] = {'kept': True} written.write_text(json.dumps(data)) @@ -240,7 +241,7 @@ def test_toggling_record_sources_keeps_unknown_parent_fields(make_snapshot): def test_snapshot_manifest_round_trips_unchanged(make_snapshot, tmp_path): - written = make_snapshot(30).index_path / SNAPSHOT_MANIFEST_FILE + written = make_snapshot(30).index_path / PROVENANCE_FILE newer = json.loads(written.read_text()) newer['future'] = {'kept': [1, 2]} newer['environment']['cuda'] = '13.0' From 6afedc7b451d7093d25674fa0d656c44f541cbd1 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 21:47:49 +0000 Subject: [PATCH 103/278] Read ingest stats inline where they are used Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/graph.py | 15 +++++---------- 1 file changed, 5 insertions(+), 10 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index 34bd20707..01b38600f 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -139,14 +139,6 @@ def _warn_low_fk_match_rates(stats: Dict[str, Any], ) -def _ingest_stats(path: Any) -> Dict[str, Any]: - try: - stats = json.loads(Path(path).read_text())['stats'] - except (OSError, ValueError, KeyError, TypeError): - return {} - return stats if isinstance(stats, dict) else {} - - def _warn_ingest_data_loss(result: Any, graph_config: Dict[str, Any]) -> None: from kumo_core.diskgraph.validation import parse_failure_warning @@ -170,8 +162,11 @@ def _warn_ingest_data_loss(result: Any, graph_config: Dict[str, Any]) -> None: warnings.warn(parse_failure_warning(name, column, count), UserWarning, stacklevel=3) - stats = _ingest_stats(getattr(result, 'ingest_stats_path', None)) - if not stats: + try: + stats = json.loads(Path(result.ingest_stats_path).read_text())['stats'] + except (AttributeError, OSError, ValueError, KeyError, TypeError): + return + if not isinstance(stats, dict): return tables = {table['table_name']: table for table in graph_config['tables']} duplicates: Dict[str, int] = {} From 39eae02ef85a544c5780233936a56c0f41036342 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 21:47:54 +0000 Subject: [PATCH 104/278] Warn once in report_other_graph with the detail that applies Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/graph.py | 25 ++++++++++------------- 1 file changed, 11 insertions(+), 14 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index 01b38600f..7582f6bb3 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -244,21 +244,18 @@ def schema_changes(before: Mapping[str, Any], def report_other_graph(message: str, before: Optional[Mapping[str, Any]], after: Mapping[str, Any], stacklevel: int) -> None: - if before is None: - warnings.warn( - f"{message} The model does not record its training graph, so " - "what differs is unknown; a different connector or table " - "location alone also changes the graph id.", UserWarning, - stacklevel=stacklevel + 1) - return - changes = schema_changes(before, after) - if changes: - warnings.warn(f"{message} Schema changes: {'; '.join(changes)}.", - UserWarning, stacklevel=stacklevel + 1) + changes = None if before is None else schema_changes(before, after) + if changes is None: + detail = ("The model does not record its training graph, so what " + "differs is unknown; a different connector or table " + "location alone also changes the graph id.") + elif changes: + detail = f"Schema changes: {'; '.join(changes)}." else: - warnings.warn(f"{message} The schema is the same; only table " - "connectors or locations differ.", UserWarning, - stacklevel=stacklevel + 1) + detail = ("The schema is the same; only table connectors or " + "locations differ.") + warnings.warn(f"{message} {detail}", UserWarning, + stacklevel=stacklevel + 1) def _store() -> Any: From ae1717cfd1eb5333769c50450d75a9703e32eebc Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 21:48:07 +0000 Subject: [PATCH 105/278] Share one artifact-root store warning across snapshots, tables and models Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_mlops.py | 6 ++++++ packages/kumo-engine/kumo/engine/_snapshot_io.py | 6 ++---- packages/kumo-engine/kumo/engine/pquery/training_table.py | 5 ++--- packages/kumo-engine/kumo/engine/trainer/trainer.py | 5 ++--- 4 files changed, 12 insertions(+), 10 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_mlops.py b/packages/kumo-engine/kumo/engine/_mlops.py index e38d35329..fc4b31d97 100644 --- a/packages/kumo-engine/kumo/engine/_mlops.py +++ b/packages/kumo-engine/kumo/engine/_mlops.py @@ -7,6 +7,7 @@ import contextvars import os import posixpath +import warnings from dataclasses import dataclass, field from typing import Any, FrozenSet, Iterable, Iterator, Optional, Tuple, Union @@ -172,6 +173,11 @@ def resolve(settings: Optional[MLOps] = None) -> Optional[MLOps]: return settings if settings is not None else _CURRENT.get() +def warn_not_stored(what: str, error: Exception, stacklevel: int) -> None: + warnings.warn(f'Could not store {what} under the artifact root: {error}', + UserWarning, stacklevel=stacklevel + 1) + + def on_root(source: Any, path: str) -> Tuple[Any, str]: r"""The connector and path a ``load_from_connector(source, path=...)`` reads: an `MLOps` ``source`` reads ``path`` under its artifact root.""" diff --git a/packages/kumo-engine/kumo/engine/_snapshot_io.py b/packages/kumo-engine/kumo/engine/_snapshot_io.py index 45bc4dd98..36b6e605e 100644 --- a/packages/kumo-engine/kumo/engine/_snapshot_io.py +++ b/packages/kumo-engine/kumo/engine/_snapshot_io.py @@ -194,7 +194,7 @@ def export_location(snapshot: Any, path: str) -> str: def store(snapshot: Any, settings: Any) -> None: r"""Export ``snapshot`` under ``/snapshots/`` unless a complete export of it is there; failures warn.""" - import warnings + from kumo.engine._mlops import warn_not_stored try: connector = settings.root.connector @@ -209,9 +209,7 @@ def store(snapshot: Any, settings: Any) -> None: and marker.get('snapshot_id') == snapshot.id): export(snapshot, connector, path) except Exception as error: - warnings.warn( - f"Could not store snapshot {snapshot.id} under the artifact " - f"root: {error}", UserWarning, stacklevel=3) + warn_not_stored(f'snapshot {snapshot.id}', error, stacklevel=3) def export(snapshot: Any, connector: Any, path: str) -> str: diff --git a/packages/kumo-engine/kumo/engine/pquery/training_table.py b/packages/kumo-engine/kumo/engine/pquery/training_table.py index 1ce3cc828..30bfb343d 100644 --- a/packages/kumo-engine/kumo/engine/pquery/training_table.py +++ b/packages/kumo-engine/kumo/engine/pquery/training_table.py @@ -135,6 +135,7 @@ def load_from_connector(cls, connector: Any, *, path: str, def _store(self, settings: Any) -> None: from pyarrow.fs import FileType + from kumo.engine._mlops import warn_not_stored from kumo.engine.connector import file_store from kumo.engine.trainer.output_config import TrainingTableOutputConfig @@ -146,9 +147,7 @@ def _store(self, settings: Any) -> None: self.export( TrainingTableOutputConfig(settings.root.connector, path)) except Exception as error: - warnings.warn( - f"Could not store training table {self.id} under the " - f"artifact root: {error}", UserWarning, stacklevel=3) + warn_not_stored(f'training table {self.id}', error, stacklevel=3) @classmethod def from_source_table( diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 32b76cf4f..af692a88a 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -283,6 +283,7 @@ def _store_model(result: TrainingResult, settings: Any) -> None: from kumo_core.object_store_sink import holds_bundle, upload_bundle from pyarrow.fs import FileSelector + from kumo.engine._mlops import warn_not_stored from kumo.engine.connector import file_store if result.bundle_id is None: @@ -298,9 +299,7 @@ def _store_model(result: TrainingResult, settings: Any) -> None: filesystem.delete_dir_contents(root, missing_dir_ok=True) upload_bundle(filesystem, root, str(result.bundle_dir)) except Exception as error: - warnings.warn( - f"Could not store model {result.bundle_id} under the artifact " - f"root: {error}", UserWarning, stacklevel=4) + warn_not_stored(f'model {result.bundle_id}', error, stacklevel=4) def _published_holdout(trainer: Any) -> tuple[Optional[Path], Optional[bool]]: From 95e3ce05e8767cc65a7c3f445866ba4875ba3932 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 21:48:18 +0000 Subject: [PATCH 106/278] Drop the never-overridden _upload_file hook and its export-file indirection Co-Authored-By: Claude Opus 5.5 --- packages/kumo-adls/tests/test_native_output.py | 8 +++----- .../src/kumo_core/object_store_sink.py | 17 +++-------------- 2 files changed, 6 insertions(+), 19 deletions(-) diff --git a/packages/kumo-adls/tests/test_native_output.py b/packages/kumo-adls/tests/test_native_output.py index a20060c28..c42670006 100644 --- a/packages/kumo-adls/tests/test_native_output.py +++ b/packages/kumo-adls/tests/test_native_output.py @@ -7,7 +7,7 @@ import pytest from kumo_adls.native_output import NativeOutputHandler from kumo_adls.sdk_filesystem import ADLSFileSystem -from kumo_core.object_store_sink import ObjectStoreSinkBackend +from kumo_core.object_store_sink import upload_file from pyarrow.fs import ( FileSelector, FileType, @@ -84,8 +84,7 @@ def test_shared_sink_copies_prepared_file_through_native_stream( payload = bytes(range(256)) * ((8 << 20) // 256) + b'last chunk' source = tmp_path / 'prepared.bin' source.write_bytes(payload) - ObjectStoreSinkBackend._upload_file(None, arrow, str(source), - AUTHORITY + '/orders/copied.bin') + upload_file(arrow, str(source), AUTHORITY + '/orders/copied.bin') with native.open_input_file('curated/orders/copied.bin') as reader: assert reader.read() == payload @@ -204,8 +203,7 @@ def test_sink_integrates_real_local_native_stream_and_shared_copy( filesystem, destination = backend._destination(connector(), 'native.bin') source = tmp_path / 'prepared.bin' source.write_bytes(b'native output') - assert ADLSSinkBackend._upload_file is ObjectStoreSinkBackend._upload_file - backend._upload_file(filesystem, str(source), destination) + upload_file(filesystem, str(source), destination) with native.open_input_file('curated/orders/native.bin') as stream: assert stream.read() == b'native output' assert 'curated/orders/native.bin' not in service.files diff --git a/packages/kumo-core/src/kumo_core/object_store_sink.py b/packages/kumo-core/src/kumo_core/object_store_sink.py index 346766cd1..d49153b95 100644 --- a/packages/kumo-core/src/kumo_core/object_store_sink.py +++ b/packages/kumo-core/src/kumo_core/object_store_sink.py @@ -159,9 +159,7 @@ def write_files(self, connector, table_name, files, from kumo_core.sources import layout filesystem, path = self._destination(connector, table_name) kind = layout.file_type(path.rstrip('/')) or layout.PARQUET - storage = _ArrowExportFiles( - filesystem, extra_columns, upload=partial(self._upload_file, - filesystem)) + storage = _ArrowExportFiles(filesystem, extra_columns) _, output = export_selected_files(files, path, fs=storage, identity=export_id or str(files), file_type=kind, @@ -193,8 +191,7 @@ def publish_model( f"path. Set output_path on the ModelExportConfig.") filesystem, destination = self.filesystem(connector, path) - upload_bundle(filesystem, destination, bundle_dir, - partial(self._upload_file, filesystem)) + upload_bundle(filesystem, destination, bundle_dir) return PublishedModel( name=posixpath.basename(destination), version=destination, @@ -202,10 +199,6 @@ def publish_model( metadata=dict(metadata), ) - def _upload_file(self, filesystem, source, destination): - """Providers may transfer an already prepared local file directly.""" - upload_file(filesystem, source, destination) - def _download_file(self, filesystem, source, destination): """Providers may transfer a model file without the streaming fallback.""" download_file(filesystem, source, destination) @@ -346,10 +339,9 @@ def _marker(export_id: Optional[str]) -> Optional[str]: class _ArrowExportFiles: """Adapt authenticated Arrow output to the shared hosted export operation.""" - def __init__(self, filesystem, extra_columns, upload=None): + def __init__(self, filesystem, extra_columns): self.filesystem = filesystem self.extra_columns = extra_columns - self.upload = upload def exists(self, path): from pyarrow.fs import FileType @@ -397,9 +389,6 @@ def copy_file(self, source, destination, overwrite=True): # CSV conversion already applied this destination's metadata. columns = self.extra_columns if source.endswith('.parquet') else None with prepared_artifact_file(source, columns) as prepared: - if self.upload is not None: - self.upload(prepared, destination) - return upload_file(self.filesystem, prepared, destination) def copy_files_threaded(self, pairs, overwrite, max_workers): From 72819088ca917e1de282ae43136631740c1f0ba9 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 21:49:01 +0000 Subject: [PATCH 107/278] Pass fit provenance to _finish_fit as a FitRecord instead of a dict Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/trainer/_provenance.py | 64 ++++++++++--------- .../kumo/engine/trainer/trainer.py | 19 +++--- .../test/test_bundle_provenance.py | 6 +- 3 files changed, 45 insertions(+), 44 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/_provenance.py b/packages/kumo-engine/kumo/engine/trainer/_provenance.py index 72bb3f859..965ddb05d 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_provenance.py +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -3,7 +3,7 @@ # SPDX-License-Identifier: Apache-2.0 """A model bundle's ``_provenance.json``; excluded from the bundle id.""" -from dataclasses import replace +from dataclasses import dataclass, replace from datetime import datetime, timezone from pathlib import Path from typing import Optional @@ -22,37 +22,39 @@ from kumo.engine.trainer._spec import PROVENANCE_ROLE, listed_file -def write_provenance( - bundle_dir: Path, - *, - bundle_id: str, - train_table: TrainingTable, - warm_start: Optional[ModelManifest], - started_at: datetime, - fit_mode: Optional[str], -) -> None: - from kumo_core.artifacts import PROVENANCE_FILE, add_bundle_file +@dataclass(frozen=True) +class FitRecord: + train_table: TrainingTable + warm_start: Optional[ModelManifest] + started_at: datetime + fit_mode: Optional[str] - finished_at = datetime.now(timezone.utc) - table = train_table.manifest - write_manifest( - bundle_dir / PROVENANCE_FILE, - ModelManifest( - id=bundle_id, - parents=ModelParents( - training_table=embedded(table) if table is not None else - stub('training_table', train_table.id or train_table.tt_id), - warm_start=None - if warm_start is None else embedded_model(warm_start)), - fit=FitTiming( - started_at=started_at.isoformat(timespec='seconds'), - finished_at=finished_at.isoformat(timespec='seconds'), - seconds=round((finished_at - started_at).total_seconds(), 3), - ), - fit_mode=fit_mode, - **stamp(), - )) - add_bundle_file(bundle_dir, PROVENANCE_ROLE, PROVENANCE_FILE) + def write(self, bundle_dir: Path, bundle_id: str) -> None: + from kumo_core.artifacts import PROVENANCE_FILE, add_bundle_file + + finished_at = datetime.now(timezone.utc) + started_at, train_table = self.started_at, self.train_table + table = train_table.manifest + write_manifest( + bundle_dir / PROVENANCE_FILE, + ModelManifest( + id=bundle_id, + parents=ModelParents( + training_table=embedded(table) if table is not None else + stub('training_table', train_table.id + or train_table.tt_id), + warm_start=None if self.warm_start is None else + embedded_model(self.warm_start)), + fit=FitTiming( + started_at=started_at.isoformat(timespec='seconds'), + finished_at=finished_at.isoformat(timespec='seconds'), + seconds=round( + (finished_at - started_at).total_seconds(), 3), + ), + fit_mode=self.fit_mode, + **stamp(), + )) + add_bundle_file(bundle_dir, PROVENANCE_ROLE, PROVENANCE_FILE) def embedded_model(manifest: ModelManifest) -> ModelManifest: diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index af692a88a..bd420c82d 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -62,6 +62,11 @@ from kumo.engine.graph import GraphSnapshot, report_other_graph from kumo.engine.pquery.prediction_table import PredictionTable from kumo.engine.pquery.training_table import TrainingTable +from kumo.engine.trainer._provenance import ( + FitRecord, + parent_model, + read_provenance, +) from kumo.engine.trainer._spec import ( HOLDOUT_META_ROLE, HOLDOUT_ROLE, @@ -113,8 +118,6 @@ def id(self) -> Optional[str]: @property def manifest(self) -> Optional[ModelManifest]: """The bundle's ``_provenance.json``; None for older bundles.""" - from kumo.engine.trainer._provenance import read_provenance - return read_provenance(self.bundle_dir) @property @@ -491,7 +494,6 @@ def fit( f'{type(export_config.connector).__name__} cannot publish ' f'models') from kumo.engine._mlops import resolve - from kumo.engine.trainer._provenance import parent_model settings = resolve(mlops) started_at = datetime.now(timezone.utc) @@ -665,7 +667,7 @@ def fit( training_table_generation=train_table.generation_plan, export_config=export_config, model_plan=effective_plan, - provenance=dict( + provenance=FitRecord( train_table=train_table, warm_start=None if warm_start is None else parent_model( warm_start._training_result.bundle_dir, @@ -684,14 +686,12 @@ def _finish_fit( export_config: Optional[ModelOutputConfig], graph_config: Optional[Dict[str, Any]] = None, model_plan: Optional[ModelPlan] = None, - provenance: Optional[Mapping[str, Any]] = None, + provenance: Optional[FitRecord] = None, settings: Optional[Any] = None, ) -> TrainingResult: """Commit durable fit outputs before making this trainer usable.""" from kumo_core.artifacts import bundle_id_from_manifest - from kumo.engine.trainer._provenance import write_provenance - effective_plan = model_plan if model_plan is not None else self._model_plan write_training_spec( training_result.bundle_dir, @@ -711,9 +711,8 @@ def _finish_fit( UserWarning, stacklevel=3) if provenance is not None and training_result.bundle_id is not None: try: - write_provenance(training_result.bundle_dir, - bundle_id=training_result.bundle_id, - **provenance) + provenance.write(training_result.bundle_dir, + training_result.bundle_id) except Exception as error: warnings.warn( f"Publishing the model bundle without _provenance.json: " diff --git a/packages/kumo-engine/test/test_bundle_provenance.py b/packages/kumo-engine/test/test_bundle_provenance.py index 854e6e0dc..5210581c8 100644 --- a/packages/kumo-engine/test/test_bundle_provenance.py +++ b/packages/kumo-engine/test/test_bundle_provenance.py @@ -20,7 +20,7 @@ read_manifest, write_manifest, ) -from kumo.engine.trainer._provenance import embedded_model +from kumo.engine.trainer._provenance import FitRecord, embedded_model from kumo.engine.trainer._spec import MANIFEST_FILE from kumo.engine.trainer.trainer import TrainingResult @@ -107,8 +107,8 @@ def _finish_stub_fit(bundle, trainer, warm_start=None, settings=None): TrainingResult({}, {}, bundle), predictive_query='PREDICT x', snapshot_id='snap', export_config=None, training_table_generation=TrainingTableGenerationPlan(), - provenance=dict(train_table=table, warm_start=warm_start, - started_at=datetime.now(timezone.utc), fit_mode=None), + provenance=FitRecord(table, warm_start, datetime.now(timezone.utc), + None), settings=settings) From e48d607b35515243640362b99e2c8fea64ea98ef Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 21:49:24 +0000 Subject: [PATCH 108/278] Rename the shared training-table finish step to finish_generated and type its request Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/pquery/predictive_query.py | 8 ++++++-- .../kumo/engine/pquery/training_table.py | 15 ++++++++------- 2 files changed, 14 insertions(+), 9 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/pquery/predictive_query.py b/packages/kumo-engine/kumo/engine/pquery/predictive_query.py index dab929730..e1ed2c1f8 100644 --- a/packages/kumo-engine/kumo/engine/pquery/predictive_query.py +++ b/packages/kumo-engine/kumo/engine/pquery/predictive_query.py @@ -94,7 +94,10 @@ PredictionTable, _compile_prediction_table, ) -from kumo.engine.pquery.training_table import TrainingTable, _finish +from kumo.engine.pquery.training_table import ( + TrainingTable, + finish_generated, +) logger = logging.getLogger(__name__) @@ -340,7 +343,8 @@ def generate_training_table( required_stages=task_plan.task_dataset.required_stages, generation_plan=generation_plan, ) - return _finish(table, snapshot, task_plan.train_table, mlops) + return finish_generated(table, snapshot, task_plan.train_table, + mlops) def generate_prediction_table( self, diff --git a/packages/kumo-engine/kumo/engine/pquery/training_table.py b/packages/kumo-engine/kumo/engine/pquery/training_table.py index 30bfb343d..6a772b315 100644 --- a/packages/kumo-engine/kumo/engine/pquery/training_table.py +++ b/packages/kumo-engine/kumo/engine/pquery/training_table.py @@ -11,6 +11,7 @@ from kumo_core.artifacts import PROVENANCE_FILE from kumo_core.model_plan import TrainingTableGenerationPlan +from kumo_core.training.table import TrainTableRequest from kumoapi.model_plan import \ TrainingTableGenerationPlan as PublicGenerationPlan from kumoapi.train import TrainingTableSpec @@ -279,7 +280,7 @@ def from_source_table( required_stages=task_plan.task_dataset.required_stages, generation_plan=generation_plan, ) - return _finish( + return finish_generated( table, snapshot, task_plan.train_table, mlops, source_table=provenance.source_table( source_table.connector, source_table.table_name, @@ -295,14 +296,14 @@ def _record_sources(snapshot: GraphSnapshot) -> bool: return manifest is not None and manifest.record_sources is True -def _finish(table: TrainingTable, snapshot: GraphSnapshot, request: Any, - mlops: Optional[Any], *, - source_table: Optional[SourceTableEntry] = None, - source_version: Any = None) -> TrainingTable: +def finish_generated(table: TrainingTable, snapshot: GraphSnapshot, + request: TrainTableRequest, mlops: Optional[Any], *, + source_table: Optional[SourceTableEntry] = None, + source_version: Any = None) -> TrainingTable: from kumo.engine._mlops import resolve seed = None - if getattr(request, 'split_family', None) in _RANDOM_SPLITS: + if request.split_family in _RANDOM_SPLITS: seed = (request.split_config or {}).get('rnd_seed') or 0 table = replace(table, id=_table_id(table, snapshot, seed, source_version, custom=source_table is not None)) @@ -333,7 +334,7 @@ def _table_id(table: TrainingTable, snapshot: GraphSnapshot, def _record_manifest(table: TrainingTable, snapshot: GraphSnapshot, - request: Any, seed: Optional[int], *, + request: TrainTableRequest, seed: Optional[int], *, source_table: Optional[SourceTableEntry] = None) -> None: from kumo.engine._manifest import embedded, stub From 709ea9ca2676f9297aa2dae8d2379a4c04febc15 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 21:50:36 +0000 Subject: [PATCH 109/278] Import finish_generated on one line Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/pquery/predictive_query.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/pquery/predictive_query.py b/packages/kumo-engine/kumo/engine/pquery/predictive_query.py index e1ed2c1f8..fcafb625e 100644 --- a/packages/kumo-engine/kumo/engine/pquery/predictive_query.py +++ b/packages/kumo-engine/kumo/engine/pquery/predictive_query.py @@ -94,10 +94,7 @@ PredictionTable, _compile_prediction_table, ) -from kumo.engine.pquery.training_table import ( - TrainingTable, - finish_generated, -) +from kumo.engine.pquery.training_table import TrainingTable, finish_generated logger = logging.getLogger(__name__) From 7878a8d253df27021bfcd73e86202ba9e0a1f834 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 21:50:36 +0000 Subject: [PATCH 110/278] Check manifest round trips once, through a model manifest with embedded parents Co-Authored-By: Claude Opus 5.5 --- .../test/test_bundle_provenance.py | 6 +++++ .../test/test_snapshot_manifest.py | 24 +---------------- .../test/test_training_table_export.py | 26 +------------------ 3 files changed, 8 insertions(+), 48 deletions(-) diff --git a/packages/kumo-engine/test/test_bundle_provenance.py b/packages/kumo-engine/test/test_bundle_provenance.py index 5210581c8..1a1954d30 100644 --- a/packages/kumo-engine/test/test_bundle_provenance.py +++ b/packages/kumo-engine/test/test_bundle_provenance.py @@ -50,6 +50,9 @@ def test_manifest_round_trips_unchanged(fitted, tmp_path): (trainer.result.bundle_dir / '_provenance.json').read_text()) newer['future'] = True newer['fit']['gpu_seconds'] = 1.5 + table = newer['parents']['training_table'] + table['parents']['label_source'] = 'x' + table['parents']['snapshot']['environment']['cuda'] = '13.0' source = tmp_path / 'newer.json' source.write_text(json.dumps(newer, indent=2, sort_keys=True) + '\n') @@ -59,6 +62,9 @@ def test_manifest_round_trips_unchanged(fitted, tmp_path): assert (tmp_path / 'out.json').read_bytes() == source.read_bytes() assert manifest.extra == {'future': True} assert manifest.fit.extra == {'gpu_seconds': 1.5} + table = manifest.parents.training_table + assert table.parents.extra == {'label_source': 'x'} + assert table.parents.snapshot.environment.extra == {'cuda': '13.0'} def test_bundle_without_a_manifest_loads_and_predicts_identically( diff --git a/packages/kumo-engine/test/test_snapshot_manifest.py b/packages/kumo-engine/test/test_snapshot_manifest.py index 192128c50..3f8acedb1 100644 --- a/packages/kumo-engine/test/test_snapshot_manifest.py +++ b/packages/kumo-engine/test/test_snapshot_manifest.py @@ -30,12 +30,7 @@ SnowflakeConnector, Table, ) -from kumo.engine._manifest import ( - SnapshotManifest, - SourceTableEntry, - read_manifest, - write_manifest, -) +from kumo.engine._manifest import SourceTableEntry from kumo.engine.connector.identity import ( DatabricksIdentity, SnowflakeIdentity, @@ -238,20 +233,3 @@ def test_toggling_record_sources_keeps_unknown_parent_fields(make_snapshot): redacted = make_snapshot(record_sources=False) assert redacted.manifest.record_sources is False assert redacted.manifest.parents.extra == {'future': {'kept': True}} - - -def test_snapshot_manifest_round_trips_unchanged(make_snapshot, tmp_path): - written = make_snapshot(30).index_path / PROVENANCE_FILE - newer = json.loads(written.read_text()) - newer['future'] = {'kept': [1, 2]} - newer['environment']['cuda'] = '13.0' - newer['parents']['source_tables']['ents']['row_count'] = 5 - source = tmp_path / 'newer.json' - source.write_text(json.dumps(newer, indent=2, sort_keys=True) + '\n') - - manifest = read_manifest(source, SnapshotManifest) - write_manifest(tmp_path / 'out.json', manifest) - - assert (tmp_path / 'out.json').read_bytes() == source.read_bytes() - assert manifest.extra == {'future': {'kept': [1, 2]}} - assert manifest.environment.extra == {'cuda': '13.0'} diff --git a/packages/kumo-engine/test/test_training_table_export.py b/packages/kumo-engine/test/test_training_table_export.py index 484f56481..5bd5486e3 100644 --- a/packages/kumo-engine/test/test_training_table_export.py +++ b/packages/kumo-engine/test/test_training_table_export.py @@ -19,13 +19,7 @@ from kumoapi.typing import Stype import kumo.engine as kumo -from kumo.engine._manifest import ( - SourceTableEntry, - TrainingTableManifest, - TrainingTableParents, - read_manifest, - write_manifest, -) +from kumo.engine._manifest import SourceTableEntry, TrainingTableParents from kumo.engine._scratch import get_scratch_dir from kumo.engine.trainer.output_config import TrainingTableOutputConfig @@ -349,24 +343,6 @@ def test_manifest_and_export_sidecar(graph, tmp_path): assert _seeds(reloaded)[1] == _seeds(table)[1] -def test_training_table_manifest_round_trips_unchanged(graph, tmp_path): - _, table, _ = _export(graph, 'PREDICT ents.LABEL FOR EACH ents.ID', - 'round_trip') - newer = json.loads( - (table.staging_dir / '_provenance.json').read_text()) - newer['future'] = 1 - newer['parents']['label_source'] = 'x' - source = tmp_path / 'newer.json' - source.write_text(json.dumps(newer, indent=2, sort_keys=True) + '\n') - - manifest = read_manifest(source, TrainingTableManifest) - write_manifest(tmp_path / 'out.json', manifest) - - assert (tmp_path / 'out.json').read_bytes() == source.read_bytes() - assert manifest.extra == {'future': 1} - assert manifest.parents.extra == {'label_source': 'x'} - - def test_custom_table_follows_the_snapshots_record_sources( graph, monkeypatch): graph_, snapshot, connector, _ = graph From a02c64747cf255ebf37941b3507bebb56648f0a6 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 21:46:57 +0000 Subject: [PATCH 111/278] Route volume model publish, fetch and preflight through FileStoreRouting Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_databricks/sink.py | 55 ++----------------- 1 file changed, 4 insertions(+), 51 deletions(-) diff --git a/packages/kumo-databricks/src/kumo_databricks/sink.py b/packages/kumo-databricks/src/kumo_databricks/sink.py index 47734fddf..fd0a7a353 100644 --- a/packages/kumo-databricks/src/kumo_databricks/sink.py +++ b/packages/kumo-databricks/src/kumo_databricks/sink.py @@ -22,26 +22,22 @@ import os import tempfile from pathlib import Path -from typing import Mapping, Optional, Sequence +from typing import Optional, Sequence import pandas as pd from databricks.sql.client import Connection -from kumo_core.object_store_sink import FilesystemSink +from kumo_core.object_store_sink import FileStoreRouting, FilesystemSink from kumo_databricks.registry import DatabricksModelRegistryMixin from kumo_databricks.session import _target_config from kumo_databricks.volume import is_volume_path, volume_filesystem from kumo_databricks.write import RunVolume from kumoapi.data_source import DataSourceType from kumoapi.jobs import WriteMode -from kumoapi.plugin import ( - ModelRegistryLocation, - PublishedModel, - SourceConnector, -) +from kumoapi.plugin import SourceConnector from pyarrow.fs import FileSystem -class DatabricksSinkBackend(DatabricksModelRegistryMixin): +class DatabricksSinkBackend(FileStoreRouting, DatabricksModelRegistryMixin): r"""Writes a produced artifact to a Databricks Delta table. Registered by default beside @@ -201,49 +197,6 @@ def write_files(self, connector, table_name, files, prepare_file=prepare) return sql.qualified_name(config.catalog, schema, table) - def publish_model( - self, - connector: SourceConnector, - bundle_dir: str, - *, - path: Optional[str] = None, - location: Optional[ModelRegistryLocation] = None, - metadata: Mapping[str, str] = {}, - metrics: Mapping[str, float] = {}, - project_id: Optional[str] = None, - ) -> PublishedModel: - """Register in Unity Catalog, or copy to a volume ``path``.""" - if location is None and path and self.holds_files(connector, path): - return FilesystemSink(self).publish_model( - connector, bundle_dir, path=path, metadata=metadata) - return super().publish_model(connector, bundle_dir, path=path, - location=location, metadata=metadata, - metrics=metrics, project_id=project_id) - - def fetch_model( - self, - connector: SourceConnector, - local_dir: str, - *, - path: Optional[str] = None, - location: Optional[ModelRegistryLocation] = None, - project_id: Optional[str] = None, - ) -> PublishedModel: - """Download from Unity Catalog, or from a volume ``path``.""" - if location is None and path and self.holds_files(connector, path): - return FilesystemSink(self).fetch_model(connector, local_dir, - path=path) - return super().fetch_model(connector, local_dir, path=path, - location=location, project_id=project_id) - - def preflight_write_files(self, connector: SourceConnector, - table_name: str, - write_mode: Optional[WriteMode] = None) -> None: - """Refuse an occupied volume directory; tables are not checked.""" - if self.holds_files(connector, table_name): - FilesystemSink(self).preflight_write_files( - connector, table_name, write_mode) - def holds_files(self, connector: SourceConnector, path: str) -> bool: """Volume paths hold files; anything else is a table.""" return is_volume_path(path) From 85fdff0947c0ebda83b8de38844f710bc3a7c122 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 21:47:28 +0000 Subject: [PATCH 112/278] Share one bounded put loop between a volume upload's writes and abort Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_databricks/volume.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/packages/kumo-databricks/src/kumo_databricks/volume.py b/packages/kumo-databricks/src/kumo_databricks/volume.py index ae7def3fe..2e14c31d8 100644 --- a/packages/kumo-databricks/src/kumo_databricks/volume.py +++ b/packages/kumo-databricks/src/kumo_databricks/volume.py @@ -206,13 +206,18 @@ def write(self, data) -> int: self._put(piece) return len(piece) - def _put(self, piece: Optional[bytes]) -> None: + def _offer(self, piece: Any) -> bool: while self._thread.is_alive(): try: self._feed.pieces.put(piece, timeout=1) - return + return True except queue.Full: continue + return False + + def _put(self, piece: Optional[bytes]) -> None: + if self._offer(piece): + return self._thread.join() if self._error is not None: raise self._error @@ -222,12 +227,7 @@ def _put(self, piece: Optional[bytes]) -> None: def abort(self) -> None: if self.closed: return - while self._thread.is_alive(): - try: - self._feed.pieces.put(_CANCEL, timeout=1) - break - except queue.Full: - continue + self._offer(_CANCEL) self._thread.join() super().close() From 4685a33c57af22b967a184c5590e6fac1f338b69 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 21:47:28 +0000 Subject: [PATCH 113/278] Drop the fake volume's unused directory set, bad-request branch and stream log Co-Authored-By: Claude Opus 5.5 --- packages/kumo-databricks/tests/fakes.py | 37 ++++++------------------- 1 file changed, 9 insertions(+), 28 deletions(-) diff --git a/packages/kumo-databricks/tests/fakes.py b/packages/kumo-databricks/tests/fakes.py index b629e404d..d42545f71 100644 --- a/packages/kumo-databricks/tests/fakes.py +++ b/packages/kumo-databricks/tests/fakes.py @@ -11,7 +11,7 @@ from urllib.parse import unquote import pyarrow as pa -from databricks.sdk.errors import BadRequest, NotFound +from databricks.sdk.errors import NotFound from databricks.sql.exc import DataError from sources.recording import RecordingConnection, RecordingCursor @@ -199,29 +199,23 @@ def delete_directory(self, path: str) -> None: class FakeVolume(FakeFiles): - """A volume with real directories, uploads and ranged downloads.""" + """A volume with uploads and ranged downloads.""" def __init__(self, objects: Mapping[str, bytes] | None = None) -> None: super().__init__(objects) - self.directories: set[str] = set() self.ignore_range = False - def _is_directory(self, path: str) -> bool: - return path in self.directories or any( - name.startswith(path + "/") for name in self.objects) - def list_directory_contents(self, path: str) -> list[FakeEntry]: path = path.rstrip("/") - if not self._is_directory(path): - raise NotFound(f"No such directory: {path}") children: dict[str, FakeEntry] = {} - for name in sorted(self.objects) + sorted(self.directories): + for name in sorted(self.objects): if not name.startswith(path + "/"): continue - child = name[len(path) + 1:].split("/")[0] - full = f"{path}/{child}" - children[full] = (FakeEntry(full, len(self.objects[full])) - if full in self.objects else FakeEntry( - full + "/", 0, is_directory=True)) + child = f"{path}/{name[len(path) + 1:].split('/')[0]}" + children[child] = (FakeEntry(child, len(self.objects[child])) + if child in self.objects else FakeEntry( + child + "/", 0, is_directory=True)) + if not children: + raise NotFound(f"No such directory: {path}") return list(children.values()) def upload(self, path: str, contents: Any, *, overwrite: bool) -> None: @@ -229,22 +223,10 @@ def upload(self, path: str, contents: Any, *, overwrite: bool) -> None: self.objects[path] = contents.read() def get_metadata(self, path: str) -> Any: - if path.endswith("/"): - raise BadRequest("Bad Request") if path not in self.objects: raise NotFound(f"No such file: {path}") return FakeMetadata(len(self.objects[path])) - def get_directory_metadata(self, path: str) -> None: - if not self._is_directory(path): - raise NotFound(f"No such directory: {path}") - - def create_directory(self, path: str) -> None: - self.directories.add(path) - - def delete_directory(self, path: str) -> None: - self.directories.discard(path.rstrip("/")) - def do(self, method: str, path: str, *, raw: bool, response_headers: list[str], headers: Mapping[str, str]) -> dict: assert method == "GET" and raw @@ -258,7 +240,6 @@ def do(self, method: str, path: str, *, raw: bool, last = last or str(size - 1) stream = FakeStreamingResponse( self.objects[remote][int(first):int(last) + 1]) - self.streams.append(stream) return { "contents": stream, "content-range": f"bytes {first}-{last}/{size}" From 1b7cb4e1d37b19942f14543f057d1edee0b1be97 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 21:54:58 +0000 Subject: [PATCH 114/278] Turn the Tracking protocol into a base class with defaults Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_tracking.py | 158 ++++++++++++++---- 1 file changed, 127 insertions(+), 31 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_tracking.py b/packages/kumo-engine/kumo/engine/_tracking.py index 4e489151b..cff3bfd63 100644 --- a/packages/kumo-engine/kumo/engine/_tracking.py +++ b/packages/kumo-engine/kumo/engine/_tracking.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -r"""Platform-neutral run records and the `Tracking` protocol.""" +r"""Platform-neutral run records and the `Tracking` base class.""" import dataclasses import json @@ -11,6 +11,7 @@ import threading import time import warnings +from abc import ABC, abstractmethod from dataclasses import dataclass, field from pathlib import Path from typing import ( @@ -22,7 +23,6 @@ List, Mapping, Optional, - Protocol, Sequence, Tuple, ) @@ -208,36 +208,136 @@ def trial_spec(events: Sequence[Event]) -> RunSpec: ) -class Tracking(Protocol): +class Tracking(ABC): r"""Records one `Trainer.fit` in an experiment tracker. Calls come from one background thread. `start` gets the fit run's record without metrics, `on_trial` a trial attempt's record so far with 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; - ``log_metrics(metrics, step)``, series of the fit run as they come (a - single trial's epochs, or a search's best validation metric after each - trial); ``tags`` (extra fit-run tags); ``log_model(bundle_dir, manifest, - 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``; ``connector``, a connector whose model - registry ``registry`` registers in when there is no ``log_model``; and - ``platform``, ``run_id`` and ``run_url``, the tracker's own run once - `start` returned, recorded as the model manifest's ``tracking``. + Override the optional members a tracker supports. ``log_model(bundle_dir, + manifest, snapshot, registry)`` is 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 returns the version it registered at + ``registry``, a ``ModelRegistryLocation``. Without it, ``registry`` + registers in the model registry of ``connector``. ``platform``, + ``run_id`` and ``run_url`` are the tracker's own run once `start` + returned, recorded as the model manifest's ``tracking``. """ + + run_id: Optional[str] = None + run_url: Optional[str] = None + tags: Optional[Mapping[str, str]] = None + connector: Any = None + log_model: Optional[Callable[..., Optional['PublishedModel']]] = None + register: Optional[Callable[..., 'PublishedModel']] = None + load_model: Optional[Callable[..., Any]] = None + + @property + def platform(self) -> str: + return type(self).__name__ + + @abstractmethod def start(self, record: RunRecord) -> None: ... + @abstractmethod def on_trial(self, record: RunRecord, event: Event) -> None: ... + @abstractmethod def finish(self, record: Optional[RunRecord], error: Optional[BaseException]) -> None: ... + def prepare(self) -> None: + r"""Called first, on the fit's own thread.""" + + def check(self, registry: Optional['ModelRegistryLocation'] = None + ) -> None: + r"""Called when `kumo.MLOps` is built.""" + + def log_metrics(self, metrics: Mapping[str, float], step: int) -> None: + r"""The fit run's series: a single trial's epochs, or a search's + best validation metric after each trial.""" + + +class _Foreign(Tracking): + r"""A tracker that does not subclass `Tracking`, with its defaults.""" + + def __init__(self, tracker: Any) -> None: + self._tracker = tracker + + def start(self, record: RunRecord) -> None: + self._tracker.start(record) + + def on_trial(self, record: RunRecord, event: Event) -> None: + self._tracker.on_trial(record, event) + + def finish(self, record: Optional[RunRecord], + error: Optional[BaseException]) -> None: + self._tracker.finish(record, error) + + def prepare(self) -> None: + getattr(self._tracker, 'prepare', _none)() + + def check(self, registry: Optional['ModelRegistryLocation'] = None + ) -> None: + check = getattr(self._tracker, 'check', None) + if callable(check): + check(registry) + + def log_metrics(self, metrics: Mapping[str, float], step: int) -> None: + getattr(self._tracker, 'log_metrics', _none)(metrics, step) + + @property + def platform(self) -> str: + return getattr(self._tracker, 'platform', + type(self._tracker).__name__) + + @property + def run_id(self) -> Optional[str]: + return getattr(self._tracker, 'run_id', None) + + @property + def run_url(self) -> Optional[str]: + return getattr(self._tracker, 'run_url', None) + + @property + def tags(self) -> Optional[Mapping[str, str]]: + return getattr(self._tracker, 'tags', None) + + @property + def connector(self) -> Any: + return getattr(self._tracker, 'connector', None) + + @property + def log_model(self) -> Optional[Callable[..., Any]]: + return _callable(getattr(self._tracker, 'log_model', None)) + + @property + def register(self) -> Optional[Callable[..., Any]]: + return _callable(getattr(self._tracker, 'register', None)) + + @property + def load_model(self) -> Optional[Callable[..., Any]]: + return getattr(self._tracker, 'load_model', None) + + def __repr__(self) -> str: + return repr(self._tracker) + + +def as_tracking(tracking: Any) -> Optional[Tracking]: + r"""``tracking``, wrapped unless it subclasses `Tracking`.""" + if tracking is None or isinstance(tracking, Tracking): + return tracking + return _Foreign(tracking) + + +def _callable(value: Any) -> Any: + return value if callable(value) else None + @dataclass class _Search: @@ -277,9 +377,8 @@ class RegistrationError(Exception): class TrackingSession: - def __init__(self, tracking: Optional[Tracking], - warn: bool = True) -> None: - self._tracker = tracking if _rank() == 0 else None + def __init__(self, tracking: Any, warn: bool = True) -> None: + self._tracker = as_tracking(tracking) if _rank() == 0 else None self._warn = warn self.spec: Optional[RunSpec] = None self._trials: Dict[int, List[Event]] = {} @@ -308,14 +407,14 @@ def start(self, spec: Callable[[], RunSpec]) -> None: return self.spec = value self._search = _Search.of(value.params) - self._call(lambda tracker: getattr(tracker, 'prepare', _none)()) + self._call(lambda tracker: tracker.prepare()) self._worker = threading.Thread(target=self._work, name='kumo-tracking', daemon=True) self._worker.start() def begin(tracker: Tracking) -> None: tracker.start(run_record(dataclasses.replace(value, tags={ - **value.tags, **(getattr(tracker, 'tags', None) or {}) + **value.tags, **(tracker.tags or {}) }))) started = self._submit(lambda: self._call(begin)) @@ -329,12 +428,11 @@ def pointer(self) -> Optional['TrackingRef']: if self.spec is None or not self._started.wait(FINISH_SECONDS): return None tracker = self._tracker - run_id = getattr(tracker, 'run_id', None) + run_id = None if tracker is None else tracker.run_id if tracker is None or not run_id: return None - return TrackingRef( - platform=str(getattr(tracker, 'platform', type(tracker).__name__)), - run_id=str(run_id), url=getattr(tracker, 'run_url', None)) + return TrackingRef(platform=str(tracker.platform), run_id=str(run_id), + url=tracker.run_url) def on_event(self, event: Event) -> None: if self.spec is None or event.trial is None: @@ -356,8 +454,7 @@ def on_event(self, event: Event) -> None: if series is not None: metrics, step = logged_names(series[0]), series[1] self._submit(lambda: self._call( - lambda tracker: getattr(tracker, 'log_metrics', _none)( - metrics, step))) + lambda tracker: tracker.log_metrics(metrics, step))) def log_model( self, @@ -367,8 +464,8 @@ def log_model( registry: Optional['ModelRegistryLocation'] = None, ) -> Tuple[Optional['PublishedModel'], Optional[BaseException]]: r"""The tracker's ``log_model``, run on the tracking thread.""" - if self._worker is None or (self._tracker is not None and not callable( - getattr(self._tracker, 'log_model', None))): + if self._worker is None or (self._tracker is not None + and self._tracker.log_model is None): return None, None done = threading.Event() failed: List[BaseException] = [] @@ -382,9 +479,8 @@ def log() -> None: 'tracking stopped after an earlier failure' + ( '' if self.error is None else f' ({full_error(self.error)})')) - registered.append( - getattr(tracker, 'log_model')(bundle_dir, manifest, - snapshot, registry)) + registered.append(tracker.log_model(bundle_dir, manifest, + snapshot, registry)) except BaseException as error: failed.append(error) finally: From 9fe3d75435ff5ef794929b2e54db5960c4711322 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 21:55:09 +0000 Subject: [PATCH 115/278] Make MlflowTracking and SnowflakeTracking subclass Tracking Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/integrations/mlflow.py | 4 ++-- packages/kumo-engine/kumo/engine/integrations/snowflake.py | 5 ++--- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index cedad020d..87c1b428d 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -53,6 +53,7 @@ DROPPED_PARAMS, RegistrationError, RunRecord, + Tracking, by_step, ) from kumo.engine._version import __version__ @@ -89,7 +90,7 @@ _UC_FORBIDDEN = re.compile(r'[ /\x00-\x1f\x7f]') -class MlflowTracking: +class MlflowTracking(Tracking): r"""Records `Trainer.fit` in MLflow: one run per fit (under the active run, if any, on the same server) with a child run per trial attempt. Source tables are logged as run inputs: Databricks tables as Delta @@ -154,7 +155,6 @@ def __init__( self.tags = dict(tags or {}) self.logs_model = log_model self._outer: Optional[str] = None - self.run_id: Optional[str] = None self._reset() def _reset(self) -> None: diff --git a/packages/kumo-engine/kumo/engine/integrations/snowflake.py b/packages/kumo-engine/kumo/engine/integrations/snowflake.py index cc2667786..bcf0c6e7c 100644 --- a/packages/kumo-engine/kumo/engine/integrations/snowflake.py +++ b/packages/kumo-engine/kumo/engine/integrations/snowflake.py @@ -43,6 +43,7 @@ DROPPED_METRICS, DROPPED_PARAMS, RunRecord, + Tracking, by_step, ) from kumo.engine.connector._extras import import_provider @@ -56,7 +57,7 @@ _FINISH_PARAMS = 20 -class SnowflakeTracking: +class SnowflakeTracking(Tracking): r"""Records `Trainer.fit` in Snowflake Experiment Tracking as one run; 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 @@ -85,7 +86,6 @@ class SnowflakeTracking: """ platform = 'snowflake' - run_url = None def __init__( self, @@ -113,7 +113,6 @@ def __init__( self.run_name = run_name self.tags = dict(tags or {}) self.native_model_lineage = native_model_lineage - self.run_id: Optional[str] = None self._reset() def _reset(self) -> None: From e16a252f0497e85d49923d238122e2d0ddc28ca4 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 22:05:48 +0000 Subject: [PATCH 116/278] Call tracker members directly in MLOps and Trainer Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_mlops.py | 16 +++++++------- .../kumo/engine/trainer/trainer.py | 21 ++++++++++++------- 2 files changed, 23 insertions(+), 14 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_mlops.py b/packages/kumo-engine/kumo/engine/_mlops.py index ff0ba7224..2f6d3f68d 100644 --- a/packages/kumo-engine/kumo/engine/_mlops.py +++ b/packages/kumo-engine/kumo/engine/_mlops.py @@ -116,15 +116,17 @@ def __post_init__(self) -> None: f'{list(KINDS)}.') if kinds: _require_writable(root) + from kumo.engine._tracking import as_tracking + + tracker = as_tracking(self.tracking) if self.registry is not None: - self._check_registry() - check = getattr(self.tracking, 'check', None) - if callable(check): - check(self.registry) + self._check_registry(tracker) + if tracker is not None: + tracker.check(self.registry) object.__setattr__(self, 'root', root) object.__setattr__(self, 'kinds', kinds) - def _check_registry(self) -> None: + def _check_registry(self, tracker: Any) -> None: from kumo.engine.trainer.output_config import ( ModelOutputConfig, check_model_location, @@ -137,12 +139,12 @@ def _check_registry(self) -> None: f'{type(registry).__name__}; the tracker registers it.') if registry.alias is not None: check_model_alias(registry.alias) - if callable(getattr(self.tracking, 'log_model', None)): + if tracker is not None and tracker.log_model is not None: if registry.version is not None: raise ValueError('MLflow assigns registered model versions; ' 'leave ModelRegistryLocation.version unset.') return - connector = getattr(self.tracking, 'connector', None) + connector = None if tracker is None else tracker.connector if connector is None: raise ValueError( 'registry= registers through the tracker: pass ' diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index c0cc8115b..ee307aa59 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -81,7 +81,12 @@ from kumo.engine._mlops import full_error, warn_failure from kumo.engine._predictor import TRAINER_PREDICTOR from kumo.engine._scratch import new_run_id, trainer_scratch_dir -from kumo.engine._tracking import RunSpec, TrackingSession, fit_inputs +from kumo.engine._tracking import ( + RunSpec, + TrackingSession, + as_tracking, + fit_inputs, +) from kumo.engine.connector import model_sink from kumo.engine.connector.base import Connector from kumo.engine.connector.source_table import SourceTable @@ -482,10 +487,10 @@ def _register_result(result: TrainingResult, settings: Any) -> PublishedModel: if settings.registry is None: raise ValueError('export(cfg) registers at cfg.registry, which is ' 'not set.') - register = getattr(settings.tracking, 'register', None) - if callable(register): + tracker = as_tracking(settings.tracking) + if tracker is not None and tracker.register is not None: return cast(PublishedModel, _with_metrics( - register(result, settings.registry), result.metrics)) + tracker.register(result, settings.registry), result.metrics)) registry = cast(ModelOutputConfig, _connector_registry(settings)) _require_publisher(registry) return _publish(result.bundle_dir, result.bundle_id, result.metrics, @@ -510,12 +515,13 @@ def _connector_registry(settings: Any) -> Optional[ModelOutputConfig]: if (settings is None or settings.registry is None or _logs_models(settings.tracking)): return None - return ModelOutputConfig(settings.tracking.connector, + return ModelOutputConfig(as_tracking(settings.tracking).connector, location=settings.registry) def _logs_models(tracking: Any) -> bool: - return callable(getattr(tracking, 'log_model', None)) + tracker = as_tracking(tracking) + return tracker is not None and tracker.log_model is not None def _publish(bundle_dir: Path, bundle_id: Optional[str], @@ -1434,7 +1440,8 @@ def load_from_connector( from kumo.engine._mlops import MLOps, on_root if isinstance(model_registry, MLOps) and location is not None: - load = getattr(model_registry.tracking, 'load_model', None) + tracker = as_tracking(model_registry.tracking) + load = None if tracker is None else tracker.load_model if path is not None or load is None: raise TypeError( 'Loading a registered model from a kumo.MLOps takes ' From 67060169ee97ad9ad8a265cdaf5ceaef0ae50e1b Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 21:46:54 +0000 Subject: [PATCH 117/278] Score frames in the trainer predictor, which then needs no mlflow extra Co-Authored-By: Claude Opus 5.5 --- .../kumo-engine/kumo/engine/_predictor.py | 96 ++++++++++++++++++- .../kumo/engine/integrations/_mlflow_model.py | 94 +----------------- .../kumo-engine/test/test_mlflow_model.py | 16 ++-- 3 files changed, 102 insertions(+), 104 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_predictor.py b/packages/kumo-engine/kumo/engine/_predictor.py index 8fd5f0977..97d6f1237 100644 --- a/packages/kumo-engine/kumo/engine/_predictor.py +++ b/packages/kumo-engine/kumo/engine/_predictor.py @@ -3,14 +3,18 @@ # SPDX-License-Identifier: Apache-2.0 r"""The predictor a model bundle's manifest names.""" +import logging +import tempfile from pathlib import Path -from typing import Any, Mapping, Protocol, Tuple, Union +from typing import Any, Mapping, Optional, Protocol, Tuple, Union import pandas as pd from kumo.engine._deps import imported TRAINER_PREDICTOR = 'kumo.engine._predictor:TrainerPredictor' +_ANCHOR = 'ANCHOR_TIMESTAMP' +logger = logging.getLogger(__name__) class Predictor(Protocol): @@ -52,10 +56,96 @@ def load(cls, bundle_dir: Union[str, Path]) -> 'TrainerPredictor': def predict(self, rows: Any, snapshot: Any, params: Mapping[str, Any]) -> Union[pd.DataFrame, Tuple[str, ...]]: - from kumo.engine.integrations._mlflow_model import predict from kumo.engine.pquery.prediction_table import PredictionTable if isinstance(rows, PredictionTable): return self.trainer.predict(snapshot, prediction_table=rows, **params).files - return predict(self.trainer, snapshot, rows, params) + return predict_frame(self.trainer, snapshot, rows, params) + + +def predict_frame(trainer: Any, snapshot: Any, rows: pd.DataFrame, + params: Mapping[str, Any]) -> pd.DataFrame: + from kumo.engine._scratch import scoped_scratch_dir + from kumo.engine.connector.file import FileConnector + + time = params.get('prediction_time') + time = None if time is None else pd.Timestamp(time) + options = dict( + binary_classification_threshold=params.get( + 'binary_classification_threshold'), + num_classes_to_return=params.get('num_classes_to_return'), + ) + _require_inputs(trainer, snapshot, rows, time) + trainer.scratch_dir.mkdir(parents=True, exist_ok=True) + with tempfile.TemporaryDirectory(dir=trainer.scratch_dir) as tmp, \ + scoped_scratch_dir(Path(tmp) / 'scratch'): + if len(rows) == 0: + return trainer.predict( + snapshot, prediction_time=None + if time is None else time.to_pydatetime(), **options).frame + rows = rows.rename(columns={'ENTITY': 'ENTITY1'}) + if time is not None and 'TIMESTAMP' not in rows: + rows = rows.assign(TIMESTAMP=time) + (Path(tmp) / 'rows').mkdir() + rows.to_parquet(Path(tmp) / 'rows' / 'rows.parquet', index=False) + entities = FileConnector(root_dir=str(Path(tmp) / 'rows'))[ + 'rows.parquet'] + timed = 'TIMESTAMP' in rows and _temporal(trainer, snapshot) + scored = trainer._predict( + trainer._prediction_table(snapshot, entities, None), + snapshot=snapshot, extra_fields=[_ANCHOR] if timed else [], + **options).frame + return _in_input_order(rows, scored) + + +def _in_input_order(rows: pd.DataFrame, scored: pd.DataFrame) -> pd.DataFrame: + r"""``scored`` in the order of ``rows``.""" + keys = ['ENTITY'] + left = pd.DataFrame({'ENTITY': rows['ENTITY1'].to_numpy()}) + right = scored[['ENTITY']].copy() + if _ANCHOR in scored: + keys.append(_ANCHOR) + left[_ANCHOR] = _seconds(rows['TIMESTAMP']) + right[_ANCHOR] = _seconds(scored[_ANCHOR]) + try: + left = left.astype({'ENTITY': right['ENTITY'].dtype}) + left['_n'] = left.groupby(keys, sort=False, dropna=False).cumcount() + right['_n'] = right.groupby(keys, sort=False, dropna=False).cumcount() + right['_row'] = range(len(right)) + order = left.merge(right, on=[*keys, '_n'])['_row'] + except (TypeError, ValueError): + order = None + scored = scored.drop(columns=_ANCHOR, errors='ignore') + if order is None or len(order) != len(scored): + logger.warning('Could not match the scores to the input rows; they ' + 'are returned in scoring order: join them on ENTITY') + return scored + return scored.iloc[order.to_numpy()].reset_index(drop=True) + + +def _temporal(trainer: Any, snapshot: Any) -> bool: + from kumo_core.pquery import query_is_temporal + + from kumo.engine.pquery import PredictiveQuery + + return query_is_temporal( + PredictiveQuery._from_snapshot(snapshot, + trainer._training_query)._parsed, + snapshot._native_graph) + + +def _seconds(values: pd.Series) -> pd.Series: + stamps = pd.to_datetime(values) + if stamps.dt.tz is not None: + stamps = stamps.dt.tz_convert(None) + return stamps.astype('datetime64[s]') + + +def _require_inputs(trainer: Any, snapshot: Any, rows: pd.DataFrame, + time: Optional[pd.Timestamp]) -> None: + if len(rows) == 0 or time is not None or 'TIMESTAMP' in rows: + return + if _temporal(trainer, snapshot): + raise ValueError("temporal query: pass params={'prediction_time': " + "...} or a TIMESTAMP column") diff --git a/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py b/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py index a2c2e4ba1..c80d8958d 100644 --- a/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py +++ b/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py @@ -1,10 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -import logging -import tempfile -from pathlib import Path -from typing import Any, Dict, Mapping, Optional, Tuple +from typing import Any, Dict, Optional, Tuple import pandas as pd @@ -16,8 +13,6 @@ _schema = import_provider('mlflow.types.schema', 'mlflow') _signature = import_provider('mlflow.models.signature', 'mlflow') -logger = logging.getLogger(__name__) -_ANCHOR = 'ANCHOR_TIMESTAMP' _DTYPES = { 'PREDICTED': 'boolean', 'SCORE': 'double', @@ -58,93 +53,6 @@ def _snapshot(self, uri: str) -> GraphSnapshot: return self._snapshots[uri] -def predict(trainer: Any, snapshot: GraphSnapshot, rows: pd.DataFrame, - params: Mapping[str, Any]) -> pd.DataFrame: - from kumo.engine._scratch import scoped_scratch_dir - from kumo.engine.connector.file import FileConnector - - time = params.get('prediction_time') - time = None if time is None else pd.Timestamp(time) - options = dict( - binary_classification_threshold=params.get( - 'binary_classification_threshold'), - num_classes_to_return=params.get('num_classes_to_return'), - ) - _require_inputs(trainer, snapshot, rows, time) - trainer.scratch_dir.mkdir(parents=True, exist_ok=True) - with tempfile.TemporaryDirectory(dir=trainer.scratch_dir) as tmp, \ - scoped_scratch_dir(Path(tmp) / 'scratch'): - if len(rows) == 0: - return trainer.predict( - snapshot, prediction_time=None - if time is None else time.to_pydatetime(), **options).frame - rows = rows.rename(columns={'ENTITY': 'ENTITY1'}) - if time is not None and 'TIMESTAMP' not in rows: - rows = rows.assign(TIMESTAMP=time) - (Path(tmp) / 'rows').mkdir() - rows.to_parquet(Path(tmp) / 'rows' / 'rows.parquet', index=False) - entities = FileConnector(root_dir=str(Path(tmp) / 'rows'))[ - 'rows.parquet'] - timed = 'TIMESTAMP' in rows and _temporal(trainer, snapshot) - scored = trainer._predict( - trainer._prediction_table(snapshot, entities, None), - snapshot=snapshot, extra_fields=[_ANCHOR] if timed else [], - **options).frame - return _in_input_order(rows, scored) - - -def _in_input_order(rows: pd.DataFrame, scored: pd.DataFrame) -> pd.DataFrame: - r"""``scored`` in the order of ``rows``.""" - keys = ['ENTITY'] - left = pd.DataFrame({'ENTITY': rows['ENTITY1'].to_numpy()}) - right = scored[['ENTITY']].copy() - if _ANCHOR in scored: - keys.append(_ANCHOR) - left[_ANCHOR] = _seconds(rows['TIMESTAMP']) - right[_ANCHOR] = _seconds(scored[_ANCHOR]) - try: - left = left.astype({'ENTITY': right['ENTITY'].dtype}) - left['_n'] = left.groupby(keys, sort=False, dropna=False).cumcount() - right['_n'] = right.groupby(keys, sort=False, dropna=False).cumcount() - right['_row'] = range(len(right)) - order = left.merge(right, on=[*keys, '_n'])['_row'] - except (TypeError, ValueError): - order = None - scored = scored.drop(columns=_ANCHOR, errors='ignore') - if order is None or len(order) != len(scored): - logger.warning('Could not match the scores to the input rows; they ' - 'are returned in scoring order: join them on ENTITY') - return scored - return scored.iloc[order.to_numpy()].reset_index(drop=True) - - -def _temporal(trainer: Any, snapshot: GraphSnapshot) -> bool: - from kumo_core.pquery import query_is_temporal - - from kumo.engine.pquery import PredictiveQuery - - return query_is_temporal( - PredictiveQuery._from_snapshot(snapshot, - trainer._training_query)._parsed, - snapshot._native_graph) - - -def _seconds(values: pd.Series) -> pd.Series: - stamps = pd.to_datetime(values) - if stamps.dt.tz is not None: - stamps = stamps.dt.tz_convert(None) - return stamps.astype('datetime64[s]') - - -def _require_inputs(trainer: Any, snapshot: GraphSnapshot, rows: pd.DataFrame, - time: Optional[pd.Timestamp]) -> None: - if len(rows) == 0 or time is not None or 'TIMESTAMP' in rows: - return - if _temporal(trainer, snapshot): - raise ValueError("temporal query: pass params={'prediction_time': " - "...} or a TIMESTAMP column") - - def snapshot_location(uri: str) -> Tuple[Connector, str]: from kumo.engine.connector import connector_for_root diff --git a/packages/kumo-engine/test/test_mlflow_model.py b/packages/kumo-engine/test/test_mlflow_model.py index 1aaa2a192..fbf42eb36 100644 --- a/packages/kumo-engine/test/test_mlflow_model.py +++ b/packages/kumo-engine/test/test_mlflow_model.py @@ -185,26 +185,26 @@ def predict_beside_another_call(self, *args, **kwargs): def test_pyfunc_names_a_missing_prediction_time_and_needs_no_threshold( fitted): - from kumo.engine.integrations._mlflow_model import predict + from kumo.engine._predictor import predict_frame trainer, table, snapshot = fitted rows = pd.DataFrame({'ENTITY': range(5)}) with pytest.raises(ValueError) as error: - predict(trainer, snapshot, rows, - {'binary_classification_threshold': .5}) + predict_frame(trainer, snapshot, rows, + {'binary_classification_threshold': .5}) assert str(error.value) == ("temporal query: pass params={" "'prediction_time': ...} or a TIMESTAMP " "column") time = str(pd.Timestamp(table.timestamps[0])) - scores = predict(trainer, snapshot, rows, {'prediction_time': time}) + scores = predict_frame(trainer, snapshot, rows, {'prediction_time': time}) assert len(scores) == 5 assert scores['TARGET_PRED'].isna().all() assert scores['True_PROB'].between(0, 1).all() def test_pyfunc_returns_rows_in_input_order(fitted): - from kumo.engine.integrations._mlflow_model import predict + from kumo.engine._predictor import predict_frame trainer, table, snapshot = fitted time = pd.Timestamp(table.timestamps[0]) @@ -214,14 +214,14 @@ def test_pyfunc_returns_rows_in_input_order(fitted): 'TIMESTAMP': [time - pd.Timedelta(days=i % 3) for i in range(len(entities))], }) - scored = predict(trainer, snapshot, rows, - {'binary_classification_threshold': .5}) + scored = predict_frame(trainer, snapshot, rows, + {'binary_classification_threshold': .5}) assert scored['ENTITY'].tolist() == [e for e in entities if e != 999] assert 'ANCHOR_TIMESTAMP' not in scored def test_scores_pair_with_their_rows_when_scoring_reorders(): - from kumo.engine.integrations._mlflow_model import _in_input_order + from kumo.engine._predictor import _in_input_order day = pd.Timestamp('2024-05-01 12:00:00.250') rows = pd.DataFrame({ From 9f9ca7e4bc4636d2bd8e4e1c28721d7f2a5f06c8 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 21:47:14 +0000 Subject: [PATCH 118/278] Return the alias error from _publish instead of appending it to a list Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/trainer/trainer.py | 25 ++++++++++--------- 1 file changed, 13 insertions(+), 12 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index ee307aa59..39536e804 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -186,7 +186,8 @@ def export(self, config: Any) -> PublishedModel: if isinstance(config, MLOps): return _register_result(self, config) _require_publisher(config) - return _publish(self.bundle_dir, self.bundle_id, self.metrics, config) + return _publish(self.bundle_dir, self.bundle_id, self.metrics, + config)[0] def with_scoring_defaults(self, defaults: ScoringDefaults) -> 'TrainingResult': @@ -394,17 +395,17 @@ def _store_and_register( registry = None if stored else _connector_registry(settings) if registry is None: return registered, error - errors: List[BaseException] = [] try: - registered = _publish(bundle_dir, None if manifest is None else - manifest.id, metrics, registry, errors, warn) + registered, alias_error = _publish( + bundle_dir, None if manifest is None else manifest.id, metrics, + registry, warn) except Exception as failure: if warn: warn_failure( f"Could not register the model: {full_error(failure)}", failure, 3) return registered, error or failure - return registered, error or (errors or [None])[0] + return registered, error or alias_error def _store_model(bundle_dir: Path, manifest: Optional[ModelManifest], @@ -494,7 +495,7 @@ def _register_result(result: TrainingResult, settings: Any) -> PublishedModel: registry = cast(ModelOutputConfig, _connector_registry(settings)) _require_publisher(registry) return _publish(result.bundle_dir, result.bundle_id, result.metrics, - registry) + registry)[0] def _with_metrics(model: Optional[PublishedModel], @@ -526,8 +527,9 @@ def _logs_models(tracking: Any) -> bool: def _publish(bundle_dir: Path, bundle_id: Optional[str], metrics: Mapping[str, float], export_config: ModelOutputConfig, - errors: Optional[List[BaseException]] = None, - warn: bool = True) -> PublishedModel: + warn: bool = True + ) -> Tuple[PublishedModel, Optional[BaseException]]: + r"""The published model, and the error setting its alias, if any.""" metadata = {} if (bundle_dir / SPEC_FILE).is_file(): spec = read_training_spec(bundle_dir) @@ -559,9 +561,8 @@ def _publish(bundle_dir: Path, bundle_id: Optional[str], f"Published {published.name} version " f"{published.version} without alias {alias}: " f"{full_error(error)}", error, 4) - if errors is not None: - errors.append(error) - return published + return published, error + return published, None def _stage_model_bundle(artifacts: Mapping[str, str], dest: Path) -> Path: @@ -1090,7 +1091,7 @@ def _finish_fit( training_result = replace( training_result, published_model=_publish( training_result.bundle_dir, training_result.bundle_id, - training_result.metrics, export_config)) + training_result.metrics, export_config)[0]) training_result = replace(training_result, _trainer=self) self._training_result = training_result From 2c20101afd1cadad5e287317a52a3b1c4090ce49 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 21:47:28 +0000 Subject: [PATCH 119/278] Read the encoders for the pip requirement with the trainer's loader Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/integrations/mlflow.py | 17 ++++------------- packages/kumo-engine/test/test_mlflow_model.py | 7 ++----- 2 files changed, 6 insertions(+), 18 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index 87c1b428d..38e59a10b 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -812,12 +812,8 @@ def _log_model(target: Any, client: Any, bundle_dir: Path, from kumo_core.artifacts import PROVENANCE_FILE from kumo.engine.integrations import _mlflow_model, _mlflow_model_code - from kumo.engine.trainer._spec import ( - SPEC_FILE, - listed_file, - read_training_spec, - ) - from kumo.engine.trainer.trainer import registry_pointers + from kumo.engine.trainer._spec import SPEC_FILE, read_training_spec + from kumo.engine.trainer.trainer import _encoders, registry_pointers mlflow = import_provider('mlflow', 'mlflow') if snapshot is None: @@ -860,7 +856,7 @@ def _log_model(target: Any, client: Any, bundle_dir: Path, input_example=example, metadata=pointers, pip_requirements=[ - _requirement(listed_file(bundle_dir, 'encoders')) + _requirement(_encoders(bundle_dir)) ], ) with target.scope(): @@ -888,15 +884,10 @@ def _id(manifest: Optional['ModelManifest']) -> Optional[str]: return None if manifest is None else manifest.id -def _requirement(encoders: Optional[Path]) -> str: - from kumo_core.encoders.config import EncodersConfig - from kumo_core.io.json import load_json +def _requirement(config: Any) -> str: from kumoapi.encoder import TextEmbedding extras = ['s3', 'gcs'] - path = None if encoders is None else str(encoders) - config = (load_json(path, EncodersConfig) - if path and Path(path).is_file() else None) if config is not None and any( isinstance(encoder, TextEmbedding) for table in config.encoders.values() diff --git a/packages/kumo-engine/test/test_mlflow_model.py b/packages/kumo-engine/test/test_mlflow_model.py index fbf42eb36..3b8459573 100644 --- a/packages/kumo-engine/test/test_mlflow_model.py +++ b/packages/kumo-engine/test/test_mlflow_model.py @@ -387,17 +387,14 @@ def log_model(target, client, bundle_dir, manifest, snapshot, **kwargs): assert 'the model is not registered' in caplog.text -def test_a_model_with_text_columns_requires_the_text_extra(tmp_path): +def test_a_model_with_text_columns_requires_the_text_extra(): from kumo_core.encoders.config import EncodersConfig - from kumo_core.io.json import save_json from kumoapi.encoder import Numerical, TextEmbedding from kumo.engine.integrations.mlflow import _requirement def requirement(encoders): - path = tmp_path / 'encoders.json' - save_json(EncodersConfig(encoders={'users': encoders}), str(path)) - return _requirement(path) + return _requirement(EncodersConfig(encoders={'users': encoders})) assert requirement({'age': Numerical()}) == ( f'kumo-engine[s3,gcs]=={kumo.__version__}') From 5b865d05d13c459f41075d980b62242f428953c4 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 21:47:44 +0000 Subject: [PATCH 120/278] Name the warning filter for what it hides and test it once Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/trainer/trainer.py | 17 +++++-------- packages/kumo-engine/test/test_fit_events.py | 24 ++----------------- .../kumo-engine/test/test_trainer_warnings.py | 11 ++++++--- 3 files changed, 16 insertions(+), 36 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 39536e804..1eab17b71 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -573,15 +573,10 @@ def _stage_model_bundle(artifacts: Mapping[str, str], dest: Path) -> Path: _F = TypeVar('_F', bound=Callable[..., Any]) -def _without_torch_scatter_hint(method: _F) -> _F: - r"""Runs ``method`` without PyTorch Geometric's GPU hint that min/max - aggregation "can be accelerated via the 'torch-scatter' package", - Lightning's tips and ``__len__`` note, and Python's fork warning. - - kumo-engine uses PyTorch's own scatter instead of that optional extension - (see ``test_wheel_uses_torch_fallbacks_without_native_extensions``), so - the hint, which PyG raises for every such aggregation, is not actionable. - """ +def _without_unactionable_warnings(method: _F) -> _F: + r"""Runs ``method`` without third-party notes users cannot act on: + PyG's torch-scatter hint, Lightning's tips and ``__len__`` note, and + Python's fork warning.""" @functools.wraps(method) def wrapper(*args: Any, **kwargs: Any) -> Any: tips = logging.getLogger('pytorch_lightning.utilities.rank_zero') @@ -683,7 +678,7 @@ def export(self, config: Any) -> PublishedModel: :meth:`TrainingResult.export`.""" return self._fitted('export').export(config) - @_without_torch_scatter_hint + @_without_unactionable_warnings def fit( self, train_table: TrainingTable, @@ -1114,7 +1109,7 @@ def _dependencies(self, result: TrainingResult, ['text'] if text else [])) return deps.records - @_without_torch_scatter_hint + @_without_unactionable_warnings def predict( self, source: Union[GraphSnapshot, PredictionTable, None] = None, diff --git a/packages/kumo-engine/test/test_fit_events.py b/packages/kumo-engine/test/test_fit_events.py index 5e881c759..77ecb676f 100644 --- a/packages/kumo-engine/test/test_fit_events.py +++ b/packages/kumo-engine/test/test_fit_events.py @@ -221,9 +221,9 @@ def test_a_snapshot_reports_progress_events(graph, tmp_path): def test_kumo_calls_drop_lightning_tips_and_only_those(caplog): from pytorch_lightning.utilities.rank_zero import rank_zero_info - from kumo.engine.trainer.trainer import _without_torch_scatter_hint + from kumo.engine.trainer.trainer import _without_unactionable_warnings - @_without_torch_scatter_hint + @_without_unactionable_warnings def fit(): rank_zero_info('💡 Tip: inside') rank_zero_info('GPU available: False') @@ -247,23 +247,3 @@ def test_an_event_prints_as_its_name_and_fields(): assert str(Event(name='epoch_ended', time=1.0, trial=1, epoch=2, metrics={'val_auroc': 0.5})) == ( 'epoch_ended trial=1 epoch=2 val_auroc=0.5') - - -def test_kumo_calls_drop_the_fork_and_iterable_dataset_notes(): - import warnings - - from kumo.engine.trainer.trainer import _without_torch_scatter_hint - - @_without_torch_scatter_hint - def fit(): - warnings.warn('This process (pid=1) is multi-threaded, use of fork() ' - 'may lead to deadlocks in the child.', - DeprecationWarning) - warnings.warn('Your `IterableDataset` has `__len__` defined.', - UserWarning) - warnings.warn('something else', UserWarning) - - with warnings.catch_warnings(record=True) as caught: - warnings.simplefilter('always') - fit() - assert [str(warning.message) for warning in caught] == ['something else'] diff --git a/packages/kumo-engine/test/test_trainer_warnings.py b/packages/kumo-engine/test/test_trainer_warnings.py index af488eacb..b95b7ab75 100644 --- a/packages/kumo-engine/test/test_trainer_warnings.py +++ b/packages/kumo-engine/test/test_trainer_warnings.py @@ -8,16 +8,21 @@ from kumoapi.model_plan import ModelPlan from kumo.engine import Trainer -from kumo.engine.trainer.trainer import _without_torch_scatter_hint +from kumo.engine.trainer.trainer import _without_unactionable_warnings HINT = ("The usage of `scatter(reduce='min')` can be accelerated via the " "'torch-scatter' package, but it was not found") -def test_torch_scatter_hint_is_hidden_and_other_warnings_are_not() -> None: - @_without_torch_scatter_hint +def test_unactionable_warnings_are_hidden_and_other_warnings_are_not() -> None: + @_without_unactionable_warnings def run() -> str: warnings.warn(HINT, UserWarning) + warnings.warn('This process (pid=1) is multi-threaded, use of fork() ' + 'may lead to deadlocks in the child.', + DeprecationWarning) + warnings.warn('Your `IterableDataset` has `__len__` defined.', + UserWarning) warnings.warn('an unrelated warning', UserWarning) return 'done' From d67ac288aad7e152a93e011eb69e5e6007a7578c Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 21:48:10 +0000 Subject: [PATCH 121/278] Name the fit's start parameters and run summaries as steps of Trainer.fit Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/trainer/trainer.py | 43 ++++++++++++------- 1 file changed, 28 insertions(+), 15 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 1eab17b71..d4a8b7355 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -40,6 +40,7 @@ build_generated_training_task_dataset, resolve_query_forecast, ) +from kumo_core.summaries import Summary from kumo_core.train.artifact_manifest import PredictionSelectionRef from kumo_core.training.explain_aggregation import aggregate_subgraph_features from kumo_core.training.workflow import TrainingWorkflowRequest @@ -408,6 +409,29 @@ def _store_and_register( return registered, error or alias_error +def _fit_identity(snapshot: GraphSnapshot, + train_table: TrainingTable) -> Dict[str, Any]: + try: + return { + 'snapshot_id': snapshot.id, + 'query_sha256': query_sha256(train_table.query), + } + except Exception: + return {} + + +def _fit_summaries(snapshot: GraphSnapshot, train_table: TrainingTable, + result: TrainingResult, + settings: Any) -> Tuple[Summary, ...]: + index = Path(snapshot.index_path) + return ( + *snapshot_summaries(index, snapshot.graph_config), + *([] if result.manifest is None else artifact_summaries( + result.manifest, 'training', settings, table_details(index))), + *split_summaries(train_table.staging_dir), + ) + + def _store_model(bundle_dir: Path, manifest: Optional[ModelManifest], settings: Any) -> None: r"""Copy the bundle to ``/models//`` unless a complete @@ -746,14 +770,7 @@ def fit( [*callbacks, session.on_event] if session.enabled else callbacks) deps = Dependencies(settings=settings) fit_start = time.monotonic() - try: - started = { - 'snapshot_id': snapshot.id, - 'query_sha256': query_sha256(train_table.query), - } - except Exception: - started = {} - events.emit('fit_started', params=started) + events.emit('fit_started', params=_fit_identity(snapshot, train_table)) try: result = self._fit(train_table, snapshot=snapshot, export_config=export_config, @@ -787,13 +804,9 @@ def fit( settings.registry.alias) session.finish( lambda: fitted_spec( - cast(RunSpec, session.spec), result, seconds, alias, ( - *snapshot_summaries(Path(snapshot.index_path), - snapshot.graph_config), - *([] if result.manifest is None else artifact_summaries( - result.manifest, 'training', settings, - table_details(Path(snapshot.index_path)))), - *split_summaries(train_table.staging_dir))), error) + cast(RunSpec, session.spec), result, seconds, alias, + _fit_summaries(snapshot, train_table, result, settings)), + error) raise_on_error(settings, session.error or error, 'Trainer.fit (its model is in trainer.result)') return result From 151aa6cad9c74beed0936766b8a24a0a9388d549 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 22:12:13 +0000 Subject: [PATCH 122/278] Check publish and file-export support with the sink protocols, not getattr Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/trainer/output_config.py | 8 ++++++-- packages/kumo-engine/kumo/engine/trainer/trainer.py | 8 ++++---- 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/output_config.py b/packages/kumo-engine/kumo/engine/trainer/output_config.py index d947fe775..5a44c180e 100644 --- a/packages/kumo-engine/kumo/engine/trainer/output_config.py +++ b/packages/kumo-engine/kumo/engine/trainer/output_config.py @@ -8,7 +8,11 @@ from kumo_core.sources.object_store_sessions import ObjectStoreSessions from kumoapi.jobs import MetadataField, PredictionArtifactType, WriteMode -from kumoapi.plugin import ModelRegistryLocation, SourceConnector +from kumoapi.plugin import ( + FileArtifactSink, + ModelRegistryLocation, + SourceConnector, +) from kumoapi.prediction_outputs import resolve_prediction_output_paths from kumo.engine.connector import is_file_destination, model_sink @@ -66,7 +70,7 @@ def destinations(self) -> dict[str, str]: def writer(self, destination: str): sink = model_sink(self.connector, destination) - if not callable(getattr(sink, 'write_files', None)): + if not isinstance(sink, FileArtifactSink): raise NotImplementedError( 'Prediction export requires FileArtifactSink.write_files; ' 'this sink only supports in-memory table output') diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 2f91d02bb..4038597a2 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -40,6 +40,7 @@ from kumoapi.model_plan import ModelPlan, TrainingTableGenerationPlan from kumoapi.plugin import ( ModelFetchingSink, + ModelPublishingSink, ModelRegistryLocation, PublishedModel, SourceConnector, @@ -358,10 +359,9 @@ def fit( from kumo.engine.trainer._runtime import local_training_bindings - if export_config is not None and not callable( - getattr(model_sink(export_config.connector, - export_config.path), 'publish_model', - None)): + if export_config is not None and not isinstance( + model_sink(export_config.connector, export_config.path), + ModelPublishingSink): raise NotImplementedError( f'{type(export_config.connector).__name__} cannot publish ' f'models') From 11b407bb4cf21a5785168eb5f717af9977b75ab6 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 22:12:13 +0000 Subject: [PATCH 123/278] Prune overlapping file-store tests Co-Authored-By: Claude Opus 5.5 --- packages/kumo-core/tests/test_chunked_fs.py | 21 +++------- .../kumo-engine/test/test_filestore_sink.py | 41 ++++++------------- 2 files changed, 18 insertions(+), 44 deletions(-) diff --git a/packages/kumo-core/tests/test_chunked_fs.py b/packages/kumo-core/tests/test_chunked_fs.py index 4fa2ee3f0..e569d4bef 100644 --- a/packages/kumo-core/tests/test_chunked_fs.py +++ b/packages/kumo-core/tests/test_chunked_fs.py @@ -249,13 +249,16 @@ def test_deletes(store): def test_parts_spool_under_the_spool_dir_and_are_removed( store, tmp_path, monkeypatch): handler, fs = store + root = [tmp_path / 'scratch'] monkeypatch.setattr(chunked_fs, '_spool_dir', None) - chunked_fs.set_spool_dir(lambda: str(tmp_path / 'scratch')) + chunked_fs.set_spool_dir(lambda: str(root[0])) _write(fs, 'root/a.bin', b'a' * (2 * PART + 1)) - handler.spool_dir = str(tmp_path / 'own') + root[0] = tmp_path / 'moved' _write(fs, 'root/b.bin', b'b') + handler.spool_dir = str(tmp_path / 'own') + _write(fs, 'root/c.bin', b'c') assert handler.spooled == [str(tmp_path / 'scratch')] * 4 + [ - str(tmp_path / 'own') + str(tmp_path / 'moved'), str(tmp_path / 'own') ] assert not os.listdir(tmp_path / 'scratch') assert not os.listdir(tmp_path / 'own') @@ -366,18 +369,6 @@ def test_a_known_size_streams_without_spooling(size): assert len(handler.spooled) == (0 if plain else 1) -def test_the_spool_dir_can_be_read_at_each_write(store, tmp_path, - monkeypatch): - handler, fs = store - root = [tmp_path / 'one'] - monkeypatch.setattr(chunked_fs, '_spool_dir', None) - chunked_fs.set_spool_dir(lambda: str(root[0])) - _write(fs, 'root/a.bin', b'a') - root[0] = tmp_path / 'two' - _write(fs, 'root/b.bin', b'b') - assert handler.spooled == [str(tmp_path / 'one'), str(tmp_path / 'two')] - - def test_reads_spanning_parts_fill_the_callers_buffer(store): _, fs = store data = os.urandom(3 * PART + 3) diff --git a/packages/kumo-engine/test/test_filestore_sink.py b/packages/kumo-engine/test/test_filestore_sink.py index 170a2d056..d0960e69c 100644 --- a/packages/kumo-engine/test/test_filestore_sink.py +++ b/packages/kumo-engine/test/test_filestore_sink.py @@ -38,7 +38,7 @@ class MemoryStore(ChunkedFileHandler): part_bytes = 4096 def __init__(self): - self.objects, self.spooled = {}, set() + self.objects = {} def get_type_name(self): return 'memory' @@ -55,7 +55,6 @@ def _list(self, directory): } def _put(self, local, path): - self.spooled.add(Path(local).parent.parent) self.objects[path] = Path(local).read_bytes() def _part(self, path, size): @@ -111,26 +110,16 @@ def test_output_configs_take_a_sinks_file_paths(connector): 'PREDS') -def test_bundles_round_trip_through_a_filesystem_only_sink( - tmp_path, connector): +def test_bundles_upload_provenance_last(tmp_path, connector): bundle = tmp_path / 'bundle' - (bundle / 'sub').mkdir(parents=True) - (bundle / '_provenance.json').write_text('{}') - (bundle / 'sub' / 'weights.bin').write_bytes(os.urandom(10000)) - files = FilesystemSink(connector.sink_backend) - published = files.publish_model(connector, str(bundle), - path='@stage/models/m') - assert published.uri == '@stage/models/m' - store = connector.sink_backend.store - assert list(store.objects)[-1] == '@stage/models/m/_provenance.json' - assert any( - name.startswith('@stage/models/m/sub/weights.bin.chunks/') - and name.endswith('/_SUCCESS') for name in store.objects) - files.fetch_model(connector, str(tmp_path / 'fetched'), - path='@stage/models/m') - for name in ('_provenance.json', 'sub/weights.bin'): - assert (tmp_path / 'fetched' / name).read_bytes() == ( - bundle / name).read_bytes() + bundle.mkdir() + for name in ('_provenance.json', 'z.bin'): + (bundle / name).write_text('{}') + FilesystemSink(connector.sink_backend).publish_model( + connector, str(bundle), path='@stage/models/m') + assert list(connector.sink_backend.store.objects) == [ + '@stage/models/m/z.bin', '@stage/models/m/_provenance.json' + ] def test_a_trainer_publishes_and_loads_through_a_filesystem_only_sink( @@ -170,7 +159,7 @@ def test_predictions_export_through_a_filesystem_only_sink(fitted, connector): 'DB.S.T').writer('DB.S.T') -def test_parquet_writes_stay_plain_up_to_one_part(tmp_path, connector): +def test_parquet_writes_stay_plain_up_to_one_part(connector): files = FilesystemSink(connector.sink_backend) frame = pd.DataFrame({'x': range(10)}) files.write_table(connector, '@stage/small.parquet', frame) @@ -185,14 +174,8 @@ def test_parquet_writes_stay_plain_up_to_one_part(tmp_path, connector): '@stage/big.parquet', filesystem=PyFileSystem(store)).shape == (100, 1) - part = tmp_path / 'part-0.parquet' - frame.to_parquet(part) - assert files.write_files(connector, '@stage/out/', [str(part)]) - assert pd.read_parquet('@stage/out/part-0.parquet', - filesystem=PyFileSystem(store)).equals(frame) - -def test_an_occupied_directory_is_refused_before_scoring(tmp_path, connector): +def test_an_occupied_directory_is_refused_before_scoring(connector): from kumo_core.artifact_export import ArtifactDestinationConflict files = FilesystemSink(connector.sink_backend) From fab1a11e51bb2d2d66f044992a620934734101c1 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 21:46:57 +0000 Subject: [PATCH 124/278] Add FileStoreRouting to send a sink's file paths to FilesystemSink Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_core/object_store_sink.py | 27 +++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/packages/kumo-core/src/kumo_core/object_store_sink.py b/packages/kumo-core/src/kumo_core/object_store_sink.py index d49153b95..dcbbd7908 100644 --- a/packages/kumo-core/src/kumo_core/object_store_sink.py +++ b/packages/kumo-core/src/kumo_core/object_store_sink.py @@ -434,3 +434,30 @@ def preflight_write_files(self, connector, table_name, write_mode=None): raise ArtifactDestinationConflict( f"Output directory '{path}/' already contains data from " f"another export. Choose an empty destination path.") + + +class FileStoreRouting: + """Sends ``holds_files`` paths to :class:`FilesystemSink`, the rest on.""" + def publish_model(self, connector, bundle_dir, *, path=None, location=None, + metadata={}, metrics={}, project_id=None): + if location is None and path is not None and self.holds_files( + connector, path): + return FilesystemSink(self).publish_model( + connector, bundle_dir, path=path, metadata=metadata) + return super().publish_model(connector, bundle_dir, path=path, + location=location, metadata=metadata, + metrics=metrics, project_id=project_id) + + def fetch_model(self, connector, local_dir, *, path=None, location=None, + project_id=None): + if location is None and path is not None and self.holds_files( + connector, path): + return FilesystemSink(self).fetch_model(connector, local_dir, + path=path) + return super().fetch_model(connector, local_dir, path=path, + location=location, project_id=project_id) + + def preflight_write_files(self, connector, table_name, write_mode=None): + if self.holds_files(connector, table_name): + FilesystemSink(self).preflight_write_files( + connector, table_name, write_mode) From 33b7f0dbd2f5e2a54f6d23852b47455a738749f0 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 22:16:55 +0000 Subject: [PATCH 125/278] Give the fake train-table requests a split_family, as the real one has Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/test/test_predictive_query.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/packages/kumo-engine/test/test_predictive_query.py b/packages/kumo-engine/test/test_predictive_query.py index 2f77622a3..92628ce4c 100644 --- a/packages/kumo-engine/test/test_predictive_query.py +++ b/packages/kumo-engine/test/test_predictive_query.py @@ -293,7 +293,8 @@ def test_generated_weight_rejected_before_artifacts(connector_kind, tmp_path, assert not run_dir.exists() native_plan = SimpleNamespace( - train_table=SimpleNamespace(lag_timesteps=None, compiler_options=None), + train_table=SimpleNamespace(lag_timesteps=None, compiler_options=None, + split_family=None), task_dataset=SimpleNamespace(required_stages=('train', 'val')), ) generated = SimpleNamespace( @@ -370,6 +371,7 @@ def test_forecast_generation_retains_effective_plan(graph_snapshot, tmp_path, ) native_plan = SimpleNamespace( train_table=SimpleNamespace( + split_family=None, lag_timesteps=3, compiler_options={ 'lag_offsets_secs': (-1, -2, -3), @@ -460,6 +462,7 @@ def test_custom_training_table_materializes_weight_and_manual_split( train_table=SimpleNamespace( lag_timesteps=None, compiler_options=None, + split_family=None, ), task_dataset=SimpleNamespace(required_stages=('train', 'val', 'test')), ) @@ -521,7 +524,8 @@ def build(**kwargs): captured.update(kwargs) return SimpleNamespace( train_table=SimpleNamespace(lag_timesteps=None, - compiler_options=None), + compiler_options=None, + split_family=None), task_dataset=SimpleNamespace(required_stages=('train', 'val')), ) From f752ceac0e5ef33ae0772c5272b97ccaf8f9024a Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 22:16:55 +0000 Subject: [PATCH 126/278] Test re-export over a missing file and a damaged export in one case Co-Authored-By: Claude Opus 5.5 --- .../kumo-engine/test/test_snapshot_export.py | 30 +++++++------------ 1 file changed, 11 insertions(+), 19 deletions(-) diff --git a/packages/kumo-engine/test/test_snapshot_export.py b/packages/kumo-engine/test/test_snapshot_export.py index 0e67c3286..3d0f7c715 100644 --- a/packages/kumo-engine/test/test_snapshot_export.py +++ b/packages/kumo-engine/test/test_snapshot_export.py @@ -383,34 +383,26 @@ def test_reexport_after_redaction_publishes_the_redacted_manifest( 'record_sources'] is False -def test_reexport_replaces_an_export_missing_a_remote_file( - tmp_path, scratch, graph, exports): +@pytest.mark.parametrize('damage', ['missing_file', 'damaged']) +def test_reexport_replaces_an_export_that_is_not_an_intact_copy( + tmp_path, scratch, graph, exports, damage): snapshot = graph.snapshot() location, exported = _export(snapshot, exports, tmp_path) first = _generation(exported) - (first / 'index' / 'col_stats.json').unlink() + stats = first / 'index/col_stats.json' + if damage == 'missing_file': + stats.unlink() + else: + stats.write_bytes(stats.read_bytes()[:3]) + (first / 'index/index_metadata.json').unlink() + _edit_manifest(exported, lambda m: m.update(snapshot_id='other')) assert snapshot.export(SnapshotOutputConfig(exports, 'one')) == location - assert _generation(exported) != first + assert not first.exists() scratch('empty') assert _index_files(_load(exports, location).index_path) == _index_files( snapshot.index_path) -def test_reexport_supersedes_a_damaged_export(tmp_path, scratch, graph, - exports): - snapshot = graph.snapshot() - location, exported = _export(snapshot, exports, tmp_path) - damaged = _generation(exported) - stats = damaged / 'index/col_stats.json' - stats.write_bytes(stats.read_bytes()[:3]) - (damaged / 'index/index_metadata.json').unlink() - _edit_manifest(exported, lambda m: m.update(snapshot_id='other')) - assert snapshot.export(SnapshotOutputConfig(exports, 'one')) == location - assert not damaged.exists() - scratch('gpu') - assert _load(exports, location).id == snapshot.id - - def test_failed_load_leaves_nothing_and_restarts_cleanly( tmp_path, scratch, graph, exports): snapshot = graph.snapshot() From 819404e6af97b3164e07d2db34e7b428f9f31ca7 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 22:19:53 +0000 Subject: [PATCH 127/278] Let file-holding sinks route their own paths; drop the engine's model_sink Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/connector/__init__.py | 14 -------- .../kumo/engine/trainer/output_config.py | 6 ++-- .../kumo/engine/trainer/trainer.py | 23 +++++------- .../kumo-engine/test/test_filestore_sink.py | 36 +++++++++++-------- 4 files changed, 33 insertions(+), 46 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/connector/__init__.py b/packages/kumo-engine/kumo/engine/connector/__init__.py index aee04c531..e5ac09e6f 100644 --- a/packages/kumo-engine/kumo/engine/connector/__init__.py +++ b/packages/kumo-engine/kumo/engine/connector/__init__.py @@ -61,20 +61,6 @@ def file_store(connector: Any, path: str) -> Optional[Tuple[Any, str]]: return None if sink is None else sink.filesystem(connector, path) -def model_sink(connector: Any, path: Optional[str]) -> Any: - r"""The sink that writes files at ``path``.""" - from kumo_core.object_store_sink import ( - FilesystemSink, - ObjectStoreSinkBackend, - ) - - sink = connector.sink_backend - if (path is None or isinstance(sink, ObjectStoreSinkBackend) - or _file_sink(connector, path) is None): - return sink - return FilesystemSink(sink) - - __all__ = [ "ADLSConnector", "ADLSIdentity", diff --git a/packages/kumo-engine/kumo/engine/trainer/output_config.py b/packages/kumo-engine/kumo/engine/trainer/output_config.py index 5a44c180e..32813bd33 100644 --- a/packages/kumo-engine/kumo/engine/trainer/output_config.py +++ b/packages/kumo-engine/kumo/engine/trainer/output_config.py @@ -15,7 +15,7 @@ ) from kumoapi.prediction_outputs import resolve_prediction_output_paths -from kumo.engine.connector import is_file_destination, model_sink +from kumo.engine.connector import is_file_destination from kumo.engine.connector.base import Connector @@ -68,8 +68,8 @@ def destinations(self) -> dict[str, str]: assert self.artifact_type is not None and self.table_name is not None return {self.artifact_type.value.lower(): self.table_name} - def writer(self, destination: str): - sink = model_sink(self.connector, destination) + def writer(self) -> FileArtifactSink: + sink = self.connector.sink_backend if not isinstance(sink, FileArtifactSink): raise NotImplementedError( 'Prediction export requires FileArtifactSink.write_files; ' diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 4038597a2..53173692f 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -54,7 +54,6 @@ require_trained_embedding_models, ) from kumo.engine._scratch import new_run_id, trainer_scratch_dir -from kumo.engine.connector import model_sink from kumo.engine.connector.base import Connector from kumo.engine.graph import GraphSnapshot from kumo.engine.pquery.prediction_table import PredictionTable @@ -158,17 +157,14 @@ def export(self, output_config: PredictionOutputConfig) -> None: self._export(output_config) def _export(self, output_config: PredictionOutputConfig) -> None: - writers = { - destination: output_config.writer(destination) - for destination in output_config.destinations().values() - } + writer = output_config.writer() outputs = [] for name, destination in output_config.destinations().items(): kind = PredictionArtifactType(name.upper()) files = PredictionResult(self.selection, kind).files outputs.append((kind, destination, files)) for kind, destination, files in outputs: - writers[destination].write_files( + writer.write_files( output_config.connector, destination, list(files), output_config.write_mode, export_id=f'{self.selection.job_id}-{kind.value}') @@ -360,7 +356,7 @@ def fit( from kumo.engine.trainer._runtime import local_training_bindings if export_config is not None and not isinstance( - model_sink(export_config.connector, export_config.path), + export_config.connector.sink_backend, ModelPublishingSink): raise NotImplementedError( f'{type(export_config.connector).__name__} cannot publish ' @@ -549,8 +545,7 @@ def _finish_fit( ) if export_config is not None: spec = read_training_spec(training_result.bundle_dir) - sink = model_sink(export_config.connector, export_config.path) - published = sink.publish_model( + published = export_config.connector.sink_backend.publish_model( export_config.connector, str(training_result.bundle_dir), path=export_config.path, @@ -655,10 +650,10 @@ def predict( ) if output_config is not None: - for destination in destinations.values(): - preflight = getattr(output_config.writer(destination), - 'preflight_write_files', None) - if callable(preflight): + preflight = getattr(output_config.writer(), + 'preflight_write_files', None) + if callable(preflight): + for destination in destinations.values(): preflight(output_config.connector, destination, write_mode=output_config.write_mode) @@ -815,7 +810,7 @@ def load_from_connector( address = ModelOutputConfig(model_registry, path=path, location=location) bundle_dir = trainer_scratch_dir(new_run_id()) / 'fetched-model' - sink = model_sink(model_registry, address.path) + sink = getattr(model_registry, 'sink_backend') if not isinstance(sink, ModelFetchingSink): raise TypeError( 'model_registry.sink_backend must implement ModelFetchingSink') diff --git a/packages/kumo-engine/test/test_filestore_sink.py b/packages/kumo-engine/test/test_filestore_sink.py index d0960e69c..c5d62ce51 100644 --- a/packages/kumo-engine/test/test_filestore_sink.py +++ b/packages/kumo-engine/test/test_filestore_sink.py @@ -9,7 +9,7 @@ import pandas as pd import pytest from kumo_core.chunked_fs import ChunkedFileHandler -from kumo_core.object_store_sink import FilesystemSink +from kumo_core.object_store_sink import FileStoreRouting, FilesystemSink from kumoapi.jobs import PredictionArtifactType from kumoapi.plugin import FileStoreSink from pyarrow.fs import PyFileSystem @@ -66,7 +66,7 @@ def _remove(self, path, *, tree): del self.objects[name] -class StageSink: +class StageSink(FileStoreRouting): def __init__(self): self.store = MemoryStore() @@ -77,6 +77,14 @@ def filesystem(self, connector, path): assert self.holds_files(connector, path) return PyFileSystem(self.store), path.rstrip('/') + def write_table(self, connector, table_name, *args, **kwargs): + return FilesystemSink(self).write_table(connector, table_name, *args, + **kwargs) + + def write_files(self, connector, table_name, *args, **kwargs): + return FilesystemSink(self).write_files(connector, table_name, *args, + **kwargs) + @pytest.fixture def connector(tmp_path): @@ -115,7 +123,7 @@ def test_bundles_upload_provenance_last(tmp_path, connector): bundle.mkdir() for name in ('_provenance.json', 'z.bin'): (bundle / name).write_text('{}') - FilesystemSink(connector.sink_backend).publish_model( + connector.sink_backend.publish_model( connector, str(bundle), path='@stage/models/m') assert list(connector.sink_backend.store.objects) == [ '@stage/models/m/z.bin', '@stage/models/m/_provenance.json' @@ -137,9 +145,6 @@ def test_a_trainer_publishes_and_loads_through_a_filesystem_only_sink( loaded = kumo.Trainer.load_from_connector(connector, path='@stage/models/m') assert loaded._snapshot_id == trainer._snapshot_id - with pytest.raises(NotImplementedError, match='cannot publish'): - trainer.fit(None, snapshot=None, - export_config=kumo.ModelOutputConfig(connector, path='T')) def test_predictions_export_through_a_filesystem_only_sink(fitted, connector): @@ -154,17 +159,18 @@ def test_predictions_export_through_a_filesystem_only_sink(fitted, connector): PredictionOutputConfig(connector, output_paths={'predictions': '@stage/again/'})) assert any(name.startswith('@stage/again/') for name in store.objects) + tables = SimpleNamespace(sink_backend=object()) with pytest.raises(NotImplementedError, match='requires'): - PredictionOutputConfig(connector, PredictionArtifactType.PREDICTIONS, - 'DB.S.T').writer('DB.S.T') + PredictionOutputConfig(tables, PredictionArtifactType.PREDICTIONS, + 'DB.S.T').writer() def test_parquet_writes_stay_plain_up_to_one_part(connector): - files = FilesystemSink(connector.sink_backend) + sink = connector.sink_backend frame = pd.DataFrame({'x': range(10)}) - files.write_table(connector, '@stage/small.parquet', frame) - files.write_table(connector, '@stage/big.parquet', - pd.DataFrame({'x': [os.urandom(50).hex() for _ in range(100)]})) + sink.write_table(connector, '@stage/small.parquet', frame) + sink.write_table(connector, '@stage/big.parquet', + pd.DataFrame({'x': [os.urandom(50).hex() for _ in range(100)]})) store = connector.sink_backend.store assert '@stage/small.parquet' in store.objects assert any( @@ -178,9 +184,9 @@ def test_parquet_writes_stay_plain_up_to_one_part(connector): def test_an_occupied_directory_is_refused_before_scoring(connector): from kumo_core.artifact_export import ArtifactDestinationConflict - files = FilesystemSink(connector.sink_backend) - files.preflight_write_files(connector, '@stage/out/') + sink = connector.sink_backend + sink.preflight_write_files(connector, '@stage/out/') connector.sink_backend.store.objects['@stage/out/old.parquet'] = b'x' for mode in (None, 'OVERWRITE', 'APPEND'): with pytest.raises(ArtifactDestinationConflict, match='already'): - files.preflight_write_files(connector, '@stage/out/', mode) + sink.preflight_write_files(connector, '@stage/out/', mode) From 52487e65a4c3809e374c178f7114a10b7a8ee52d Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 22:21:04 +0000 Subject: [PATCH 128/278] Inline the single-use download_bundle and drop the unused bundle upload callback Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_core/object_store_sink.py | 55 ++++++++----------- 1 file changed, 22 insertions(+), 33 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/object_store_sink.py b/packages/kumo-core/src/kumo_core/object_store_sink.py index dcbbd7908..aef94cf88 100644 --- a/packages/kumo-core/src/kumo_core/object_store_sink.py +++ b/packages/kumo-core/src/kumo_core/object_store_sink.py @@ -6,9 +6,8 @@ import posixpath import shutil from abc import ABC, abstractmethod -from functools import partial from pathlib import Path -from typing import Any, Callable, Mapping, Optional, Tuple +from typing import Any, Mapping, Optional, Tuple import pandas as pd from kumoapi.jobs import WriteMode @@ -55,44 +54,18 @@ def download_file(filesystem: FileSystem, source: str, shutil.copyfileobj(reader, writer, _CHUNK) -def upload_bundle(filesystem: FileSystem, destination: str, bundle_dir: str, - upload: Optional[Callable[[str, str], None]] = None) -> None: +def upload_bundle(filesystem: FileSystem, destination: str, + bundle_dir: str) -> None: """Copy ``bundle_dir`` under ``destination``, provenance last.""" from kumo_core.artifacts import PROVENANCE_FILE - upload = upload or partial(upload_file, filesystem) root = Path(bundle_dir) for source in sorted((item for item in root.rglob('*') if item.is_file()), key=lambda item: (item.name == PROVENANCE_FILE, item)): target = f'{destination}/{source.relative_to(root).as_posix()}' _prepared(filesystem, posixpath.dirname(target)) - upload(str(source), target) - - -def download_bundle( - filesystem: FileSystem, source: str, local_dir: str, - download: Optional[Callable[[str, str], None]] = None) -> bool: - """Replace ``local_dir`` with ``source``; False if it is empty.""" - from kumo_core.sources.layout import confined - from pyarrow.fs import FileSelector, FileType - - download = download or partial(download_file, filesystem) - destination = Path(local_dir) - if destination.exists(): - shutil.rmtree(destination) - destination.mkdir(parents=True) - found = False - selector = FileSelector(source, recursive=True, allow_not_found=True) - for info in filesystem.get_file_info(selector): - if info.type is not FileType.File: - continue - found = True - # Remote keys are untrusted. - target = destination / confined(info.path[len(source) + 1:]) - target.parent.mkdir(parents=True, exist_ok=True) - download(info.path, str(target)) - return found + upload_file(filesystem, str(source), target) class ObjectStoreSinkBackend(ABC): @@ -214,6 +187,9 @@ def fetch_model( ) -> PublishedModel: """Replace the local directory with the published bundle.""" del project_id + from kumo_core.sources.layout import confined + from pyarrow.fs import FileSelector, FileType + if location is not None: raise NotImplementedError( f"The local backend does not import a model by name from " @@ -227,8 +203,21 @@ def fetch_model( f"config.") filesystem, source = self.filesystem(connector, path) - if not download_bundle(filesystem, source, local_dir, - partial(self._download_file, filesystem)): + destination = Path(local_dir) + if destination.exists(): + shutil.rmtree(destination) + destination.mkdir(parents=True) + found = False + selector = FileSelector(source, recursive=True, allow_not_found=True) + for info in filesystem.get_file_info(selector): + if info.type is not FileType.File: + continue + found = True + # Remote keys are untrusted. + target = destination / confined(info.path[len(source) + 1:]) + target.parent.mkdir(parents=True, exist_ok=True) + self._download_file(filesystem, info.path, str(target)) + if not found: raise ValueError( f"Nothing is published at {self._uri(source)}. Check the " f"path the model was exported to.") From 4143035b021193cac38d36c9a26cd28e5463758b Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 22:25:19 +0000 Subject: [PATCH 129/278] Route a file-holding sink's write_table and write_files through FileStoreRouting Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_core/object_store_sink.py | 17 +++++++++++++++++ .../kumo-engine/test/test_filestore_sink.py | 10 +--------- 2 files changed, 18 insertions(+), 9 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/object_store_sink.py b/packages/kumo-core/src/kumo_core/object_store_sink.py index aef94cf88..ce4e7322d 100644 --- a/packages/kumo-core/src/kumo_core/object_store_sink.py +++ b/packages/kumo-core/src/kumo_core/object_store_sink.py @@ -427,6 +427,23 @@ def preflight_write_files(self, connector, table_name, write_mode=None): class FileStoreRouting: """Sends ``holds_files`` paths to :class:`FilesystemSink`, the rest on.""" + def write_table(self, connector, table_name, frame, + write_mode=WriteMode.OVERWRITE, project_id=None, *, + export_id=None): + sink = FilesystemSink(self) if self.holds_files( + connector, table_name) else super() + return sink.write_table(connector, table_name, frame, write_mode, + project_id, export_id=export_id) + + def write_files(self, connector, table_name, files, + write_mode=WriteMode.OVERWRITE, project_id=None, *, + export_id=None, extra_columns=None): + sink = FilesystemSink(self) if self.holds_files( + connector, table_name) else super() + return sink.write_files(connector, table_name, files, write_mode, + project_id, export_id=export_id, + extra_columns=extra_columns) + def publish_model(self, connector, bundle_dir, *, path=None, location=None, metadata={}, metrics={}, project_id=None): if location is None and path is not None and self.holds_files( diff --git a/packages/kumo-engine/test/test_filestore_sink.py b/packages/kumo-engine/test/test_filestore_sink.py index c5d62ce51..380640c90 100644 --- a/packages/kumo-engine/test/test_filestore_sink.py +++ b/packages/kumo-engine/test/test_filestore_sink.py @@ -9,7 +9,7 @@ import pandas as pd import pytest from kumo_core.chunked_fs import ChunkedFileHandler -from kumo_core.object_store_sink import FileStoreRouting, FilesystemSink +from kumo_core.object_store_sink import FileStoreRouting from kumoapi.jobs import PredictionArtifactType from kumoapi.plugin import FileStoreSink from pyarrow.fs import PyFileSystem @@ -77,14 +77,6 @@ def filesystem(self, connector, path): assert self.holds_files(connector, path) return PyFileSystem(self.store), path.rstrip('/') - def write_table(self, connector, table_name, *args, **kwargs): - return FilesystemSink(self).write_table(connector, table_name, *args, - **kwargs) - - def write_files(self, connector, table_name, *args, **kwargs): - return FilesystemSink(self).write_files(connector, table_name, *args, - **kwargs) - @pytest.fixture def connector(tmp_path): From 42a8dc6512161cbd96a3e791f002736109088e02 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 22:25:44 +0000 Subject: [PATCH 130/278] Keep Tracking a protocol and put its defaults in TrackingBase Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_tracking.py | 74 ++++++++++--------- .../kumo/engine/integrations/mlflow.py | 4 +- .../kumo/engine/integrations/snowflake.py | 4 +- 3 files changed, 42 insertions(+), 40 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_tracking.py b/packages/kumo-engine/kumo/engine/_tracking.py index cff3bfd63..b83266984 100644 --- a/packages/kumo-engine/kumo/engine/_tracking.py +++ b/packages/kumo-engine/kumo/engine/_tracking.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -r"""Platform-neutral run records and the `Tracking` base class.""" +r"""Platform-neutral run records and the `Tracking` protocol.""" import dataclasses import json @@ -11,7 +11,6 @@ import threading import time import warnings -from abc import ABC, abstractmethod from dataclasses import dataclass, field from pathlib import Path from typing import ( @@ -23,6 +22,7 @@ List, Mapping, Optional, + Protocol, Sequence, Tuple, ) @@ -208,23 +208,39 @@ def trial_spec(events: Sequence[Event]) -> RunSpec: ) -class Tracking(ABC): +class Tracking(Protocol): r"""Records one `Trainer.fit` in an experiment tracker. Calls come from one background thread. `start` gets the fit run's record without metrics, `on_trial` a trial attempt's record so far with 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). - Override the optional members a tracker supports. ``log_model(bundle_dir, - manifest, snapshot, registry)`` is 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 returns the version it registered at - ``registry``, a ``ModelRegistryLocation``. Without it, ``registry`` - registers in the model registry of ``connector``. ``platform``, - ``run_id`` and ``run_url`` are the tracker's own run once `start` - returned, recorded as the model manifest's ``tracking``. + Optional members: ``prepare()``, called first on the fit's own thread; + ``log_metrics(metrics, step)``, series of the fit run as they come (a + single trial's epochs, or a search's best validation metric after each + trial); ``tags`` (extra fit-run tags); ``log_model(bundle_dir, manifest, + 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``; ``connector``, a connector whose model + registry ``registry`` registers in when there is no ``log_model``; and + ``platform``, ``run_id`` and ``run_url``, the tracker's own run once + `start` returned, recorded as the model manifest's ``tracking``. """ + def start(self, record: RunRecord) -> None: + ... + + def on_trial(self, record: RunRecord, event: Event) -> None: + ... + + def finish(self, record: Optional[RunRecord], + error: Optional[BaseException]) -> None: + ... + + +class TrackingBase(Tracking): + r"""`Tracking` with defaults for its optional members.""" run_id: Optional[str] = None run_url: Optional[str] = None @@ -238,33 +254,19 @@ class Tracking(ABC): def platform(self) -> str: return type(self).__name__ - @abstractmethod - def start(self, record: RunRecord) -> None: - ... - - @abstractmethod - def on_trial(self, record: RunRecord, event: Event) -> None: - ... - - @abstractmethod - def finish(self, record: Optional[RunRecord], - error: Optional[BaseException]) -> None: - ... - def prepare(self) -> None: - r"""Called first, on the fit's own thread.""" + pass def check(self, registry: Optional['ModelRegistryLocation'] = None ) -> None: - r"""Called when `kumo.MLOps` is built.""" + pass def log_metrics(self, metrics: Mapping[str, float], step: int) -> None: - r"""The fit run's series: a single trial's epochs, or a search's - best validation metric after each trial.""" + pass -class _Foreign(Tracking): - r"""A tracker that does not subclass `Tracking`, with its defaults.""" +class _Foreign(TrackingBase): + r"""A tracker that does not subclass `TrackingBase`, with its defaults.""" def __init__(self, tracker: Any) -> None: self._tracker = tracker @@ -328,9 +330,9 @@ def __repr__(self) -> str: return repr(self._tracker) -def as_tracking(tracking: Any) -> Optional[Tracking]: - r"""``tracking``, wrapped unless it subclasses `Tracking`.""" - if tracking is None or isinstance(tracking, Tracking): +def as_tracking(tracking: Any) -> Optional[TrackingBase]: + r"""``tracking``, wrapped unless it subclasses `TrackingBase`.""" + if tracking is None or isinstance(tracking, TrackingBase): return tracking return _Foreign(tracking) @@ -412,7 +414,7 @@ def start(self, spec: Callable[[], RunSpec]) -> None: name='kumo-tracking', daemon=True) self._worker.start() - def begin(tracker: Tracking) -> None: + def begin(tracker: TrackingBase) -> None: tracker.start(run_record(dataclasses.replace(value, tags={ **value.tags, **(tracker.tags or {}) }))) @@ -583,7 +585,7 @@ def _work(self) -> None: logger.warning('Experiment tracking failed: %s', full_error(error), exc_info=_debug()) - def _call(self, call: Callable[[Tracking], Any]) -> None: + def _call(self, call: Callable[[TrackingBase], Any]) -> None: tracker = self._tracker if tracker is None: return diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index 38e59a10b..e86cf95f0 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -53,7 +53,7 @@ DROPPED_PARAMS, RegistrationError, RunRecord, - Tracking, + TrackingBase, by_step, ) from kumo.engine._version import __version__ @@ -90,7 +90,7 @@ _UC_FORBIDDEN = re.compile(r'[ /\x00-\x1f\x7f]') -class MlflowTracking(Tracking): +class MlflowTracking(TrackingBase): r"""Records `Trainer.fit` in MLflow: one run per fit (under the active run, if any, on the same server) with a child run per trial attempt. Source tables are logged as run inputs: Databricks tables as Delta diff --git a/packages/kumo-engine/kumo/engine/integrations/snowflake.py b/packages/kumo-engine/kumo/engine/integrations/snowflake.py index bcf0c6e7c..c31e265f5 100644 --- a/packages/kumo-engine/kumo/engine/integrations/snowflake.py +++ b/packages/kumo-engine/kumo/engine/integrations/snowflake.py @@ -43,7 +43,7 @@ DROPPED_METRICS, DROPPED_PARAMS, RunRecord, - Tracking, + TrackingBase, by_step, ) from kumo.engine.connector._extras import import_provider @@ -57,7 +57,7 @@ _FINISH_PARAMS = 20 -class SnowflakeTracking(Tracking): +class SnowflakeTracking(TrackingBase): r"""Records `Trainer.fit` in Snowflake Experiment Tracking as one run; 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 From 4320db4d825eecdd90205ae81b2bc76d210b1f93 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 22:32:10 +0000 Subject: [PATCH 131/278] Forward a foreign tracker's members and fall back to TrackingBase defaults Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_tracking.py | 69 ++++--------------- 1 file changed, 12 insertions(+), 57 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_tracking.py b/packages/kumo-engine/kumo/engine/_tracking.py index b83266984..39efd03d1 100644 --- a/packages/kumo-engine/kumo/engine/_tracking.py +++ b/packages/kumo-engine/kumo/engine/_tracking.py @@ -271,60 +271,20 @@ class _Foreign(TrackingBase): def __init__(self, tracker: Any) -> None: self._tracker = tracker - def start(self, record: RunRecord) -> None: - self._tracker.start(record) - - def on_trial(self, record: RunRecord, event: Event) -> None: - self._tracker.on_trial(record, event) - - def finish(self, record: Optional[RunRecord], - error: Optional[BaseException]) -> None: - self._tracker.finish(record, error) - - def prepare(self) -> None: - getattr(self._tracker, 'prepare', _none)() - - def check(self, registry: Optional['ModelRegistryLocation'] = None - ) -> None: - check = getattr(self._tracker, 'check', None) - if callable(check): - check(registry) - - def log_metrics(self, metrics: Mapping[str, float], step: int) -> None: - getattr(self._tracker, 'log_metrics', _none)(metrics, step) + def __getattribute__(self, name: str) -> Any: + if name.startswith('_'): + return object.__getattribute__(self, name) + if name in ('start', 'on_trial', 'finish'): + return getattr(self._tracker, name) + value = getattr(self._tracker, name, _MISSING) + called = name in ('check', 'log_model', 'register') + if value is _MISSING or called and not callable(value): + return object.__getattribute__(self, name) + return value @property def platform(self) -> str: - return getattr(self._tracker, 'platform', - type(self._tracker).__name__) - - @property - def run_id(self) -> Optional[str]: - return getattr(self._tracker, 'run_id', None) - - @property - def run_url(self) -> Optional[str]: - return getattr(self._tracker, 'run_url', None) - - @property - def tags(self) -> Optional[Mapping[str, str]]: - return getattr(self._tracker, 'tags', None) - - @property - def connector(self) -> Any: - return getattr(self._tracker, 'connector', None) - - @property - def log_model(self) -> Optional[Callable[..., Any]]: - return _callable(getattr(self._tracker, 'log_model', None)) - - @property - def register(self) -> Optional[Callable[..., Any]]: - return _callable(getattr(self._tracker, 'register', None)) - - @property - def load_model(self) -> Optional[Callable[..., Any]]: - return getattr(self._tracker, 'load_model', None) + return type(self._tracker).__name__ def __repr__(self) -> str: return repr(self._tracker) @@ -337,10 +297,6 @@ def as_tracking(tracking: Any) -> Optional[TrackingBase]: return _Foreign(tracking) -def _callable(value: Any) -> Any: - return value if callable(value) else None - - @dataclass class _Search: r"""The fit run's series: one trial's epochs, or a search's best.""" @@ -608,8 +564,7 @@ def _debug() -> bool: return logger.isEnabledFor(logging.DEBUG) -def _none(*args: Any) -> None: - return None +_MISSING = object() def _abandon(tracker: Tracking) -> None: From 76acbea84d50a6e52433887eed4734d693e24d1a Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 22:32:20 +0000 Subject: [PATCH 132/278] Let FileStoreRouting send Databricks volume paths to file writers; keep only table writes Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_databricks/sink.py | 113 +++++++++--------- 1 file changed, 54 insertions(+), 59 deletions(-) diff --git a/packages/kumo-databricks/src/kumo_databricks/sink.py b/packages/kumo-databricks/src/kumo_databricks/sink.py index fd0a7a353..80f32a6a6 100644 --- a/packages/kumo-databricks/src/kumo_databricks/sink.py +++ b/packages/kumo-databricks/src/kumo_databricks/sink.py @@ -26,7 +26,7 @@ import pandas as pd from databricks.sql.client import Connection -from kumo_core.object_store_sink import FileStoreRouting, FilesystemSink +from kumo_core.object_store_sink import FileStoreRouting from kumo_databricks.registry import DatabricksModelRegistryMixin from kumo_databricks.session import _target_config from kumo_databricks.volume import is_volume_path, volume_filesystem @@ -37,56 +37,7 @@ from pyarrow.fs import FileSystem -class DatabricksSinkBackend(FileStoreRouting, DatabricksModelRegistryMixin): - r"""Writes a produced artifact to a Databricks Delta table. - - Registered by default beside - :class:`~kumo_databricks.backend.DatabricksSourceBackend`, so a - run that read its graph from a catalog exports its predictions back to one - with no further setup. Installed as - ``pip install 'kumoai-studio[databricks]'``. - - A destination names its schema and its table in one string, joined by - ``__kumo__``, because that is the only field a Databricks output has -- - :func:`kumoai_studio.databricks.to_db_table_name` builds it from the - ``(schema, table)`` tuple an :class:`~kumoai_studio.artifact_export.OutputConfig` - takes. The schema must be the connector's own, which - :func:`kumo_databricks.output.within_schema` checks and the - hosted connector checks through as well. - - Every write stages through the connector's managed volume, - ``.._kumo_staging``, created on first use. That is not a - local choice: ``COPY INTO`` is how a Delta table is loaded, and hosted - states the alternatives -- S3, ADLS, an arbitrary volume, inline cloud - credentials -- are unsupported on purpose. The volume is emptied of what - this export put there once the table has landed. - - Args: - identity: A fallback identity, used only for connectors that name none - themselves. As with the reader, leaving it unset is the common - case, since a bound - :class:`~kumoai_studio.connector.DatabricksConnector` records the - identity it was created with. That identity needs ``CREATE TABLE`` - on the connector's schema and ``READ``/``WRITE VOLUME`` on its - staging volume, which a read-only role used for ingest will not - have. - - Example: - >>> import kumoai_studio as kumo # doctest: +SKIP - >>> from kumo_databricks import ( # doctest: +SKIP - ... DatabricksSinkBackend) - >>> kumo.set_backend(kumo.LocalBackend( # doctest: +SKIP - ... storage=paths, - ... sinks={kumo.DataSourceType.DATABRICKS: - ... DatabricksSinkBackend(identity='prod-dbx')}, - ... )) - """ - - source_type = DataSourceType.DATABRICKS - - _default_hint = ('LocalBackend(sinks={DataSourceType.DATABRICKS: ' - 'DatabricksSinkBackend(identity=identity)})') - +class _DatabricksTables(DatabricksModelRegistryMixin): def write_table( self, connector: SourceConnector, @@ -128,9 +79,6 @@ def write_table( names a schema other than the connector's, or if the staging volume is neither reachable nor creatable. """ - if self.holds_files(connector, table_name): - return FilesystemSink(self).write_table( - connector, table_name, frame, write_mode, export_id=export_id) del project_id, export_id from kumo_databricks import output, sql, write @@ -159,11 +107,7 @@ def write_table( def write_files(self, connector, table_name, files, write_mode=WriteMode.OVERWRITE, project_id=None, *, export_id=None, extra_columns=None): - """Copy Parquet to a volume path, or COPY it into a table.""" - if self.holds_files(connector, table_name): - return FilesystemSink(self).write_files( - connector, table_name, files, write_mode, - export_id=export_id, extra_columns=extra_columns) + """COPY selected Parquet into a table.""" del project_id, export_id from functools import partial @@ -197,6 +141,57 @@ def write_files(self, connector, table_name, files, prepare_file=prepare) return sql.qualified_name(config.catalog, schema, table) + +class DatabricksSinkBackend(FileStoreRouting, _DatabricksTables): + r"""Writes a produced artifact to a Databricks Delta table. + + Registered by default beside + :class:`~kumo_databricks.backend.DatabricksSourceBackend`, so a + run that read its graph from a catalog exports its predictions back to one + with no further setup. Installed as + ``pip install 'kumoai-studio[databricks]'``. + + A destination names its schema and its table in one string, joined by + ``__kumo__``, because that is the only field a Databricks output has -- + :func:`kumoai_studio.databricks.to_db_table_name` builds it from the + ``(schema, table)`` tuple an :class:`~kumoai_studio.artifact_export.OutputConfig` + takes. The schema must be the connector's own, which + :func:`kumo_databricks.output.within_schema` checks and the + hosted connector checks through as well. + + Every write stages through the connector's managed volume, + ``.._kumo_staging``, created on first use. That is not a + local choice: ``COPY INTO`` is how a Delta table is loaded, and hosted + states the alternatives -- S3, ADLS, an arbitrary volume, inline cloud + credentials -- are unsupported on purpose. The volume is emptied of what + this export put there once the table has landed. + + Args: + identity: A fallback identity, used only for connectors that name none + themselves. As with the reader, leaving it unset is the common + case, since a bound + :class:`~kumoai_studio.connector.DatabricksConnector` records the + identity it was created with. That identity needs ``CREATE TABLE`` + on the connector's schema and ``READ``/``WRITE VOLUME`` on its + staging volume, which a read-only role used for ingest will not + have. + + Example: + >>> import kumoai_studio as kumo # doctest: +SKIP + >>> from kumo_databricks import ( # doctest: +SKIP + ... DatabricksSinkBackend) + >>> kumo.set_backend(kumo.LocalBackend( # doctest: +SKIP + ... storage=paths, + ... sinks={kumo.DataSourceType.DATABRICKS: + ... DatabricksSinkBackend(identity='prod-dbx')}, + ... )) + """ + + source_type = DataSourceType.DATABRICKS + + _default_hint = ('LocalBackend(sinks={DataSourceType.DATABRICKS: ' + 'DatabricksSinkBackend(identity=identity)})') + def holds_files(self, connector: SourceConnector, path: str) -> bool: """Volume paths hold files; anything else is a table.""" return is_volume_path(path) From 31252b6edaf1d6a36efd544b01ddebcf9888a576 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 22:37:47 +0000 Subject: [PATCH 133/278] Ask whether MLOps' tracker logs models in one helper Co-Authored-By: Claude Opus 5.5 --- .../kumo-engine/kumo/engine/trainer/trainer.py | 14 ++++---------- 1 file changed, 4 insertions(+), 10 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index d4a8b7355..90789e07e 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -390,7 +390,7 @@ def _store_and_register( warn_failure(f'{stored}; the model is not registered', stored, 3) registered, error = (None, None) if session is None else session.log_model( bundle_dir, manifest, snapshot, - None if stored else _mlflow_registry(settings)) + None if stored or not _logs_models(settings) else settings.registry) registered = _with_metrics(registered, metrics) error = stored or error registry = None if stored else _connector_registry(settings) @@ -530,22 +530,16 @@ def _with_metrics(model: Optional[PublishedModel], return replace(model, metrics=dict(metrics)) -def _mlflow_registry(settings: Any) -> Optional[ModelRegistryLocation]: - if settings is None or not _logs_models(settings.tracking): - return None - return settings.registry - - def _connector_registry(settings: Any) -> Optional[ModelOutputConfig]: if (settings is None or settings.registry is None - or _logs_models(settings.tracking)): + or _logs_models(settings)): return None return ModelOutputConfig(as_tracking(settings.tracking).connector, location=settings.registry) -def _logs_models(tracking: Any) -> bool: - tracker = as_tracking(tracking) +def _logs_models(settings: Any) -> bool: + tracker = None if settings is None else as_tracking(settings.tracking) return tracker is not None and tracker.log_model is not None From a3d4929a9d12db05fc73dc4269258f77338d8582 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 22:37:47 +0000 Subject: [PATCH 134/278] List the warnings fit and predict hide in one table Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/trainer/trainer.py | 24 +++++++++---------- 1 file changed, 11 insertions(+), 13 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 90789e07e..eb21bcab8 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -591,6 +591,15 @@ def _stage_model_bundle(artifacts: Mapping[str, str], dest: Path) -> Path: _F = TypeVar('_F', bound=Callable[..., Any]) +_UNACTIONABLE = ( + (r"The usage of `scatter\(reduce='\w+'\)` can be accelerated via the " + r"'torch-scatter' package", UserWarning), + (r'Your `IterableDataset` has `__len__`', UserWarning), + (r'This process .* is multi-threaded, use of fork\(\) may lead to ' + r'deadlocks', DeprecationWarning), +) + + def _without_unactionable_warnings(method: _F) -> _F: r"""Runs ``method`` without third-party notes users cannot act on: PyG's torch-scatter hint, Lightning's tips and ``__len__`` note, and @@ -601,19 +610,8 @@ def wrapper(*args: Any, **kwargs: Any) -> Any: tips.addFilter(_no_lightning_tip) try: with warnings.catch_warnings(): - warnings.filterwarnings( - 'ignore', - message=r"The usage of `scatter\(reduce='\w+'\)` can be " - r"accelerated via the 'torch-scatter' package", - category=UserWarning, - ) - warnings.filterwarnings( - 'ignore', message=r'Your `IterableDataset` has `__len__`', - category=UserWarning) - warnings.filterwarnings( - 'ignore', message=r'This process .* is multi-threaded, ' - r'use of fork\(\) may lead to deadlocks', - category=DeprecationWarning) + for message, category in _UNACTIONABLE: + warnings.filterwarnings('ignore', message, category) return method(*args, **kwargs) finally: tips.removeFilter(_no_lightning_tip) From dd8eeaf38185260c9cf15393c39ec721117f924b Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 22:37:47 +0000 Subject: [PATCH 135/278] Inline the MLflow target's registry URI and SageMaker session helpers Co-Authored-By: Claude Opus 5.5 --- .../engine/integrations/_mlflow_target.py | 19 ++++++------------- 1 file changed, 6 insertions(+), 13 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py b/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py index 7982f5ba5..8a6ee9b9e 100644 --- a/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py +++ b/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py @@ -7,7 +7,7 @@ import os import re import threading -from contextlib import ExitStack, contextmanager +from contextlib import contextmanager from dataclasses import dataclass from typing import Any, Callable, ContextManager, Iterator, Mapping, Optional @@ -138,7 +138,9 @@ def resolve(target: Any) -> Target: if isinstance(target, str): if SAGEMAKER_ARN.fullmatch(target): return _sagemaker(SageMakerMLflow(target)) - return Target(target, registry_for(target)) + scheme, colon, rest = target.partition(':') + return Target(target, f'databricks-uc{colon}{rest}' + if scheme == 'databricks' else target) if isinstance(target, SageMakerMLflow): return _sagemaker(target) from kumo.engine.connector.databricks import DatabricksConnector @@ -165,13 +167,6 @@ def _workspace_host() -> Optional[str]: return None -def registry_for(tracking_uri: str) -> str: - r"""Unity Catalog on Databricks, else the tracking server itself.""" - scheme, colon, rest = tracking_uri.partition(':') - return f'databricks-uc{colon}{rest}' if scheme == 'databricks' else ( - tracking_uri) - - def _principal(identity: Any) -> Any: from kumo_core.sources.principals import store @@ -230,10 +225,8 @@ def _sagemaker(target: SageMakerMLflow) -> Target: def _sagemaker_session(arn: str, session: Any) -> Iterator[None]: r"""The SageMaker plugin's state, signing with ``session``.""" plugin = import_provider('sagemaker_mlflow.session', 'sagemaker') - with ExitStack() as stack: - stack.enter_context(environ({'SAGEMAKER_MLFLOW_ASSUME_ROLE_ARN': None})) - stack.enter_context(global_tracking_uri(arn)) - stack.enter_context(plugin.use_session(session)) + with environ({'SAGEMAKER_MLFLOW_ASSUME_ROLE_ARN': None}), \ + global_tracking_uri(arn), plugin.use_session(session): yield From af2546441febcfa614aea65afa96715b41c1c0a7 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 22:37:47 +0000 Subject: [PATCH 136/278] Drop Dependencies' unused locations and read manifest dependencies directly Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_deps.py | 32 +++++-------------- .../kumo-engine/kumo/engine/_predictor.py | 2 +- 2 files changed, 9 insertions(+), 25 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_deps.py b/packages/kumo-engine/kumo/engine/_deps.py index 64d76a591..65c1e7762 100644 --- a/packages/kumo-engine/kumo/engine/_deps.py +++ b/packages/kumo-engine/kumo/engine/_deps.py @@ -31,8 +31,6 @@ stamp, ) -Location = Union[str, Tuple[Any, str]] - class Asset(Protocol): r"""A content-addressed file or directory a model predicts with.""" @@ -113,18 +111,15 @@ def imported(entry: str) -> Any: class Dependencies: r"""Records a fit's dependencies and resolves a model's assets.""" def __init__(self, records: Sequence[Dependency] = (), - settings: Any = None, - locations: Sequence[Location] = ()) -> None: + settings: Any = None) -> None: self.records: List[Dependency] = list(records) self.errors: List[BaseException] = [] self._settings = settings - self._locations = list(locations) @classmethod - def of(cls, manifest: Optional[ModelManifest], settings: Any = None, - locations: Sequence[Location] = ()) -> 'Dependencies': - return cls(getattr(manifest, 'dependencies', None) or (), settings, - locations) + def of(cls, manifest: Optional[ModelManifest], + settings: Any = None) -> 'Dependencies': + return cls((manifest and manifest.dependencies) or (), settings) def add(self, dependency: Dependency) -> Dependency: self.records.append(dependency) @@ -178,14 +173,7 @@ def resolve(self, asset_kind: str, name: str) -> Path: def candidates(self, record: Dependency) -> List[Tuple[Any, str]]: from kumo.engine.connector import connector_for_root - relative = f'{record.asset_kind}/{record.sha256}' found: List[Tuple[Any, str]] = [] - for location in self._locations: - if isinstance(location, str): - found.append((connector_for_root(location), relative)) - else: - connector, path = location - found.append((connector, f"{path.rstrip('/')}/{relative}")) if record.location: parent, _, leaf = record.location.rstrip('/').rpartition('/') try: @@ -194,16 +182,11 @@ def candidates(self, record: Dependency) -> List[Tuple[Any, str]]: pass settings = self._settings if settings is not None: - found.append((settings.root.connector, - settings.root.location(f'assets/{relative}'))) + found.append((settings.root.connector, settings.root.location( + f'assets/{record.asset_kind}/{record.sha256}'))) return found -def asset_cache() -> Path: - base = os.environ.get('XDG_CACHE_HOME') or Path.home() / '.cache' - return Path(base) / 'kumo' / 'assets' - - def _store(asset: Asset, sha256: str, settings: Any) -> Optional[str]: r"""Writes ``asset`` once under the root; returns where.""" from kumo_core.object_store_sink import _prepared @@ -257,7 +240,8 @@ def _fetch(connector: Any, path: str, sha256: str) -> Optional[Path]: return None if isinstance(filesystem, LocalFileSystem): return Path(root) - target = asset_cache() / sha256 + cache = os.environ.get('XDG_CACHE_HOME') or Path.home() / '.cache' + target = Path(cache) / 'kumo' / 'assets' / sha256 if (target / PROVENANCE_FILE).is_file(): return target target.parent.mkdir(parents=True, exist_ok=True) diff --git a/packages/kumo-engine/kumo/engine/_predictor.py b/packages/kumo-engine/kumo/engine/_predictor.py index 97d6f1237..5a8720f36 100644 --- a/packages/kumo-engine/kumo/engine/_predictor.py +++ b/packages/kumo-engine/kumo/engine/_predictor.py @@ -31,7 +31,7 @@ def load_predictor(bundle_dir: Union[str, Path]) -> Predictor: manifest = read_provenance(Path(bundle_dir)) entry = next((record.name - for record in getattr(manifest, 'dependencies', None) or () + for record in (manifest and manifest.dependencies) or () if record.kind == 'predictor'), TRAINER_PREDICTOR) try: predictor = imported(entry) From b319b99e66843d121f9930c4197181663d727f20 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 22:38:14 +0000 Subject: [PATCH 137/278] Render summaries once in TrackingBase for the MLflow and Snowflake trackers Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_tracking.py | 67 ++++++++++++++++++- .../kumo/engine/integrations/mlflow.py | 51 +------------- .../kumo/engine/integrations/snowflake.py | 67 ++----------------- 3 files changed, 71 insertions(+), 114 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_tracking.py b/packages/kumo-engine/kumo/engine/_tracking.py index 39efd03d1..7b22f8585 100644 --- a/packages/kumo-engine/kumo/engine/_tracking.py +++ b/packages/kumo-engine/kumo/engine/_tracking.py @@ -28,12 +28,21 @@ ) from kumo_core.instrumentation import format_error +from kumo_core.summaries import ( + Confusion, + Curve, + Dataset, + Histogram, + Scalar, + Summary, + Table, + Tag, +) from kumo.engine._events import Event, _finite_floats from kumo.engine._mlops import full_error, warn_failure if TYPE_CHECKING: - from kumo_core.summaries import Summary from kumoapi.plugin import ModelRegistryLocation, PublishedModel from kumo.engine._manifest import ModelManifest, TrackingRef @@ -60,7 +69,7 @@ class RunRecord: artifacts: Mapping[str, str] = field(default_factory=dict) inputs: Mapping[str, Mapping[str, Optional[str]]] = field( default_factory=dict) - summaries: Tuple['Summary', ...] = () + summaries: Tuple[Summary, ...] = () @dataclass(frozen=True) @@ -78,7 +87,7 @@ class RunSpec: default_factory=dict) parents: Mapping[str, Optional[str]] = field(default_factory=dict) artifacts: Mapping[str, str] = field(default_factory=dict) - summaries: Sequence['Summary'] = () + summaries: Sequence[Summary] = () status: Optional[str] = None error: Optional[str] = None @@ -264,6 +273,58 @@ def check(self, registry: Optional['ModelRegistryLocation'] = None def log_metrics(self, metrics: Mapping[str, float], step: int) -> None: pass + def log_summaries(self, summaries: Sequence[Summary]) -> None: + for summary in summaries: + if isinstance(summary, Scalar): + self._log_scalar(summary) + elif isinstance(summary, Tag): + self._log_tag(summary) + elif isinstance(summary, Dataset): + self._log_dataset(summary) + elif isinstance(summary, (Table, Confusion)): + self._log_table(_as_table(summary)) + else: + self._log_points(_points(summary)) + + def _log_scalar(self, summary: Scalar) -> None: + raise NotImplementedError + + def _log_tag(self, summary: Tag) -> None: + raise NotImplementedError + + def _log_points(self, series: Mapping[str, Sequence[Tuple[int, float]]] + ) -> None: + raise NotImplementedError + + def _log_table(self, summary: Table) -> None: + raise NotImplementedError + + def _log_dataset(self, summary: Dataset) -> None: + pass + + +def _points(summary: Summary) -> Dict[str, Sequence[Tuple[int, float]]]: + if isinstance(summary, Curve): + return {summary.x.name: summary.x.points, + summary.y.name: summary.y.points} + if isinstance(summary, Histogram): + return {f'{summary.name}/count': tuple(enumerate(summary.counts)), + f'{summary.name}/lower_edge': tuple(enumerate(summary.edges, 1))} + return {summary.name: summary.points} + + +def _as_table(summary: Summary) -> Table: + if isinstance(summary, Confusion): + return Table(summary.name, ('actual', *summary.labels), + tuple((label, *row) for label, row in zip( + summary.labels, summary.counts))) + return summary + + +def table_json(table: Table) -> str: + return json.dumps({'columns': list(table.columns), + 'data': [list(row) for row in table.rows]}, default=str) + class _Foreign(TrackingBase): r"""A tracker that does not subclass `TrackingBase`, with its defaults.""" diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index e86cf95f0..c289e7b7b 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -34,17 +34,7 @@ from urllib.parse import quote from kumo_core.instrumentation import format_error -from kumo_core.summaries import ( - Confusion, - Curve, - Dataset, - Histogram, - Scalar, - Series, - Summary, - Table, - Tag, -) +from kumo_core.summaries import Dataset, Scalar, Table, Tag from kumoapi.plugin import BUNDLE_ID_KEY, ModelRegistryLocation, PublishedModel from kumo.engine._events import Event @@ -55,6 +45,7 @@ RunRecord, TrackingBase, by_step, + table_json, ) from kumo.engine._version import __version__ from kumo.engine.connector._extras import import_provider @@ -428,10 +419,6 @@ def finish(self, record: Optional[RunRecord], def log_metrics(self, metrics: Mapping[str, float], step: int) -> None: self._log_metrics(self._parent, metrics, step) - def log_summaries(self, summaries: Sequence[Summary]) -> None: - for summary in summaries: - _SINKS[type(summary)](self, summary) - def _log_scalar(self, summary: Scalar) -> None: model_id = self._model_id if summary.name.startswith('test/') else None self._log_metrics(self._parent, {summary.name: summary.value}, @@ -440,25 +427,9 @@ def _log_scalar(self, summary: Scalar) -> None: def _log_tag(self, summary: Tag) -> None: self._log_tags(self._parent, {summary.name: summary.value}) - def _log_series(self, summary: Series) -> None: - self._log_points({summary.name: summary.points}) - - def _log_curve(self, summary: Curve) -> None: - self._log_points({summary.x.name: summary.x.points, - summary.y.name: summary.y.points}) - - def _log_histogram(self, summary: Histogram) -> None: - self._log_points({ - f'{summary.name}/count': tuple(enumerate(summary.counts)), - f'{summary.name}/lower_edge': tuple(enumerate(summary.edges, 1)), - }) - def _log_table(self, summary: Table) -> None: path = f'{summary.name}.json' - self._log_artifacts({path: json.dumps({ - 'columns': list(summary.columns), - 'data': [list(row) for row in summary.rows], - }, default=str)}) + self._log_artifacts({path: table_json(summary)}) with self.target.scope(): tags = self._client.get_run(self._parent).data.tags logged = json.loads(tags.get('mlflow.loggedArtifacts', '[]')) @@ -467,12 +438,6 @@ def _log_table(self, summary: Table) -> None: self._client.set_tag(self._parent, 'mlflow.loggedArtifacts', json.dumps([*logged, entry])) - def _log_confusion(self, summary: Confusion) -> None: - self._log_table(Table( - summary.name, ('actual', *summary.labels), - tuple((label, *row) - for label, row in zip(summary.labels, summary.counts)))) - def _log_dataset(self, summary: Dataset) -> None: source = _dataset_source(summary.source) if source is None: @@ -564,16 +529,6 @@ def _log_metrics(self, run_id: Optional[str], ]) -_SINKS: Dict[type, Callable[[MlflowTracking, Any], None]] = { - Scalar: MlflowTracking._log_scalar, - Tag: MlflowTracking._log_tag, - Series: MlflowTracking._log_series, - Curve: MlflowTracking._log_curve, - Histogram: MlflowTracking._log_histogram, - Table: MlflowTracking._log_table, - Confusion: MlflowTracking._log_confusion, - Dataset: MlflowTracking._log_dataset, -} _BATCH_METRICS = 1000 # MLflow's per-request limit _SYSTEM_METRICS_FAILURES = 6 _COLUMN_TYPES = {'int64': 'long', 'int32': 'integer', 'float64': 'double', diff --git a/packages/kumo-engine/kumo/engine/integrations/snowflake.py b/packages/kumo-engine/kumo/engine/integrations/snowflake.py index c31e265f5..613ba2c73 100644 --- a/packages/kumo-engine/kumo/engine/integrations/snowflake.py +++ b/packages/kumo-engine/kumo/engine/integrations/snowflake.py @@ -4,7 +4,6 @@ r"""Snowflake Experiment Tracking for `Trainer.fit` (``pip install 'kumo-engine[snowflake]'``).""" -import json import logging import posixpath import re @@ -12,31 +11,11 @@ import uuid from contextlib import ExitStack from pathlib import Path -from typing import ( - Any, - Callable, - Dict, - List, - Mapping, - Optional, - Sequence, - Set, - Tuple, -) +from typing import Any, Dict, List, Mapping, Optional, Sequence, Set, Tuple from kumo_core.instrumentation import format_error from kumo_core.sources.sql import resolve_snowflake_identifier -from kumo_core.summaries import ( - Confusion, - Curve, - Dataset, - Histogram, - Scalar, - Series, - Summary, - Table, - Tag, -) +from kumo_core.summaries import Scalar, Table, Tag from kumo.engine._events import Event from kumo.engine._tracking import ( @@ -45,6 +24,7 @@ RunRecord, TrackingBase, by_step, + table_json, ) from kumo.engine.connector._extras import import_provider @@ -272,41 +252,14 @@ def _new_run_name(self, tracking: Any, experiment: Any, name or 'a generated run name') return name - def log_summaries(self, summaries: Sequence[Summary]) -> None: - if self._run is not None: - for summary in summaries: - _SINKS[type(summary)](self, summary) - def _log_scalar(self, summary: Scalar) -> None: self._log_metrics({summary.name: summary.value}, _MAX_METRICS) def _log_tag(self, summary: Tag) -> None: self._log_params({summary.name: summary.value}) - def _log_series(self, summary: Series) -> None: - self._log_points({summary.name: summary.points}) - - def _log_curve(self, summary: Curve) -> None: - self._log_points({summary.x.name: summary.x.points, - summary.y.name: summary.y.points}) - - def _log_histogram(self, summary: Histogram) -> None: - self._log_points({ - f'{summary.name}/count': tuple(enumerate(summary.counts)), - f'{summary.name}/lower_edge': tuple(enumerate(summary.edges, 1)), - }) - def _log_table(self, summary: Table) -> None: - self._log_artifacts({f'{summary.name}.json': json.dumps({ - 'columns': list(summary.columns), - 'data': [list(row) for row in summary.rows], - }, default=str)}) - - def _log_confusion(self, summary: Confusion) -> None: - self._log_table(Table( - summary.name, ('actual', *summary.labels), - tuple((label, *row) - for label, row in zip(summary.labels, summary.counts)))) + self._log_artifacts({f'{summary.name}.json': table_json(summary)}) def _log_points(self, series: Mapping[str, Sequence[Tuple[int, float]]] ) -> None: @@ -362,18 +315,6 @@ def _log_metrics(self, metrics: Mapping[str, float], limit: int, self._tracking.log_metrics(kept, step=step) -_SINKS: Dict[type, Callable[[SnowflakeTracking, Any], None]] = { - Scalar: SnowflakeTracking._log_scalar, - Tag: SnowflakeTracking._log_tag, - Series: SnowflakeTracking._log_series, - Curve: SnowflakeTracking._log_curve, - Histogram: SnowflakeTracking._log_histogram, - Table: SnowflakeTracking._log_table, - Confusion: SnowflakeTracking._log_confusion, - Dataset: lambda tracker, summary: None, -} - - def _run_identifier(name: str) -> str: r"""``name`` as an unquoted identifier, which artifact URIs need.""" name = re.sub(r'[^A-Za-z0-9_]', '_', name) From 1ccd3565af228022f394b28fc79927bdcbe3533b Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 22:38:14 +0000 Subject: [PATCH 138/278] Resolve Snowflake's experiment schema once and drop checks on calls made only after start Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/integrations/snowflake.py | 40 ++++++++----------- 1 file changed, 17 insertions(+), 23 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/integrations/snowflake.py b/packages/kumo-engine/kumo/engine/integrations/snowflake.py index 613ba2c73..11f71d48a 100644 --- a/packages/kumo-engine/kumo/engine/integrations/snowflake.py +++ b/packages/kumo-engine/kumo/engine/integrations/snowflake.py @@ -118,15 +118,19 @@ def load_model(self, location: Any, **kwargs: Any) -> Any: return Trainer.load_from_connector(self.connector, location=location, **kwargs) - def check(self, registry: Any = None) -> None: - r"""Raises before any work unless the session can see the - experiment's schema (one metadata query, no warehouse).""" + def _location(self) -> Tuple[Optional[str], str]: database = self.database or self.session.get_current_database() schema = self.schema or (None if self.database else self.session.get_current_schema()) + return database, schema or 'PUBLIC' + + def check(self, registry: Any = None) -> None: + r"""Raises before any work unless the session can see the + experiment's schema (one metadata query, no warehouse).""" + database, schema = self._location() if database is None: return - name = resolve_snowflake_identifier(schema or 'PUBLIC') + name = resolve_snowflake_identifier(schema) try: rows = self.session.sql( f"SHOW SCHEMAS LIKE '{name.replace(chr(39), chr(39) * 2)}' " @@ -138,7 +142,7 @@ def check(self, registry: Any = None) -> None: if not rows: raise ValueError( f'Snowflake Experiment Tracking needs schema {database}.' - f'{schema or "PUBLIC"}, which does not exist or this role ' + f'{schema}, which does not exist or this role ' f'cannot see.') def start(self, record: RunRecord) -> None: @@ -157,12 +161,10 @@ def start(self, record: RunRecord) -> None: f'Snowflake; end that run, or pass SnowflakeTracking the ' f'session it uses') self._tracking = tracking - database = self.database or self.session.get_current_database() - schema = self.schema or (None if self.database else - self.session.get_current_schema()) + database, schema = self._location() experiment = tracking.set_experiment(self.experiment_name, database_name=database, - schema_name=schema or 'PUBLIC') + schema_name=schema) self._run = tracking.start_run( self._new_run_name( tracking, experiment.name, @@ -175,8 +177,6 @@ def start(self, record: RunRecord) -> None: self._log_artifacts(record.artifacts) def on_trial(self, record: RunRecord, event: Event) -> None: - if self._run is None: - return prefix = f'{record.name}/' key = (event.trial or 0, int(event.params['attempt'])) if event.name == 'trial_started': @@ -210,7 +210,11 @@ def finish(self, record: Optional[RunRecord], if k not in self._start_tags}, _MAX_PARAMS) self._log_artifacts(record.artifacts) self.log_summaries(record.summaries) - self._log_dropped() + dropped = {DROPPED_METRICS: len(self._dropped_metrics), + DROPPED_PARAMS: self._dropped_params} + if any(dropped.values()): + self._tracking.log_params( + {name: n for name, n in dropped.items() if n}) if error is None: status = 'FINISHED' finally: @@ -225,8 +229,7 @@ def finish(self, record: Optional[RunRecord], self._reset() def log_metrics(self, metrics: Mapping[str, float], step: int) -> None: - if self._run is not None: - self._log_metrics(metrics, _EPOCH_METRICS, step=step) + self._log_metrics(metrics, _EPOCH_METRICS, step=step) def _new_run_name(self, tracking: Any, experiment: Any, run_name: Optional[str]) -> Optional[str]: @@ -277,15 +280,6 @@ def _log_artifacts(self, artifacts: Mapping[str, str]) -> None: str(path), posixpath.dirname(name) or None) self._artifacts.add(name) - def _log_dropped(self) -> None: - dropped = { - DROPPED_METRICS: len(self._dropped_metrics), - DROPPED_PARAMS: self._dropped_params, - } - dropped = {name: count for name, count in dropped.items() if count} - if dropped: - self._tracking.log_params(dropped) - def _log_params(self, params: Mapping[str, str], limit: int = _MAX_PARAMS - _FINISH_PARAMS) -> None: kept = {} From 4b88e05f754a8e174acf833a4c0f1ee37fd1f217 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 22:38:14 +0000 Subject: [PATCH 139/278] Inline MLflow's experiment creation and Delta source name and version helpers Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/integrations/mlflow.py | 27 ++++++------------- 1 file changed, 8 insertions(+), 19 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index c289e7b7b..db3912f01 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -207,10 +207,8 @@ def check(self, registry: Optional[ModelRegistryLocation] = None) -> None: except Exception as error: raise RuntimeError(f'MLflow at {uri} is not usable with these ' f'credentials: {full_error(error)}') from error - if missing: - self._create_experiment(client, owned) - - def _create_experiment(self, client: Any, owned: FrozenSet[str]) -> None: + if not missing: + return with _limited_http(owned), self.target.scope(): try: client.create_experiment(self.experiment_name) @@ -975,9 +973,13 @@ def _dataset_source(entry: Mapping[str, Optional[str]]) -> Any: return None if connector == 'DATABRICKS': delta = import_provider('mlflow.data.delta_dataset_source', 'mlflow') + parts = re.findall(r'`((?:[^`]|``)*)`', location) + head = (entry.get('version') or '').split('_', 1)[0] return delta.DeltaDatasetSource( - delta_table_name=_unquoted(location), - delta_table_version=_delta_version(entry)) + delta_table_name='.'.join(part.replace('``', '`') + for part in parts) or location, + delta_table_version=int(head) if entry.get('version_kind') + == 'exact' and head.isdigit() else None) if connector in _PATH_CONNECTORS: registry = import_provider('mlflow.data.dataset_source_registry', 'mlflow') @@ -987,19 +989,6 @@ def _dataset_source(entry: Mapping[str, Optional[str]]) -> Any: return None -def _unquoted(name: str) -> str: - parts = re.findall(r'`((?:[^`]|``)*)`', name) - return '.'.join(part.replace('``', '`') for part in parts) or name - - -def _delta_version(entry: Mapping[str, Optional[str]]) -> Optional[int]: - version = entry.get('version') or '' - head = version.split('_', 1)[0] - if entry.get('version_kind') == 'exact' and head.isdigit(): - return int(head) - return None - - def _mlflow_params(params: Mapping[str, str]) -> Dict[str, str]: return { key: value From dcf56c93c949fcb7a61deb1c70fb1f93e8af47d5 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 22:38:14 +0000 Subject: [PATCH 140/278] Batch MLflow metrics and series through one helper Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/integrations/mlflow.py | 38 +++++++++---------- 1 file changed, 19 insertions(+), 19 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index db3912f01..3130a1ad4 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -460,17 +460,9 @@ def _log_dataset(self, summary: Dataset) -> None: def _log_points(self, series: Mapping[str, Sequence[Tuple[int, float]]] ) -> None: - entities = self._mlflow.entities - timestamp = int(time.time() * 1000) - metrics = [ - entities.Metric(metric_key(name), float(value), timestamp, step) - for name, points in series.items() for step, value in points - ] - with self.target.scope(): - for start in range(0, len(metrics), _BATCH_METRICS): - self._client.log_batch( - self._parent, - metrics=metrics[start:start + _BATCH_METRICS]) + self._log_batch(self._parent, [(name, value, step) + for name, points in series.items() + for step, value in points]) def _log_artifacts(self, artifacts: Mapping[str, str]) -> None: for name, text in artifacts.items(): @@ -514,17 +506,25 @@ def _log_params(self, run_id: str, params: Mapping[str, str]) -> None: def _log_metrics(self, run_id: Optional[str], metrics: Mapping[str, float], step: int = 0, model_id: Optional[str] = None) -> None: - if run_id is None or not metrics: - return + if run_id is not None: + self._log_batch(run_id, [(name, value, step) + for name, value in metrics.items()], + model_id) + + def _log_batch(self, run_id: Optional[str], + points: Sequence[Tuple[str, float, int]], + model_id: Optional[str] = None) -> None: entities = self._mlflow.entities timestamp = int(time.time() * 1000) + metrics = [ + entities.Metric(metric_key(name), float(value), timestamp, step, + model_id=model_id) + for name, value, step in points + ] with self.target.scope(): - self._client.log_batch( - run_id, metrics=[ - entities.Metric(metric_key(key), float(value), timestamp, - step, model_id=model_id) - for key, value in metrics.items() - ]) + for start in range(0, len(metrics), _BATCH_METRICS): + self._client.log_batch( + run_id, metrics=metrics[start:start + _BATCH_METRICS]) _BATCH_METRICS = 1000 # MLflow's per-request limit From 43eb34c532229cd6083dedc5d1e50914fd671b5b Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 22:39:36 +0000 Subject: [PATCH 141/278] Register an MLOps export in TrainingResult.export, sharing its publish path Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/trainer/trainer.py | 23 +++++++------------ 1 file changed, 8 insertions(+), 15 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index eb21bcab8..3cb9ba08a 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -185,7 +185,14 @@ def export(self, config: Any) -> PublishedModel: from kumo.engine._mlops import MLOps if isinstance(config, MLOps): - return _register_result(self, config) + if config.registry is None: + raise ValueError('export(cfg) registers at cfg.registry, ' + 'which is not set.') + tracker = as_tracking(config.tracking) + if tracker is not None and tracker.register is not None: + return cast(PublishedModel, _with_metrics( + tracker.register(self, config.registry), self.metrics)) + config = _connector_registry(config) _require_publisher(config) return _publish(self.bundle_dir, self.bundle_id, self.metrics, config)[0] @@ -508,20 +515,6 @@ def _require_publisher(export_config: ModelOutputConfig) -> None: check_model_location(export_config) -def _register_result(result: TrainingResult, settings: Any) -> PublishedModel: - if settings.registry is None: - raise ValueError('export(cfg) registers at cfg.registry, which is ' - 'not set.') - tracker = as_tracking(settings.tracking) - if tracker is not None and tracker.register is not None: - return cast(PublishedModel, _with_metrics( - tracker.register(result, settings.registry), result.metrics)) - registry = cast(ModelOutputConfig, _connector_registry(settings)) - _require_publisher(registry) - return _publish(result.bundle_dir, result.bundle_id, result.metrics, - registry)[0] - - def _with_metrics(model: Optional[PublishedModel], metrics: Mapping[str, float]) -> Optional[PublishedModel]: r"""``model`` with ``metrics``, where its registry returned none.""" From 9d5851d32fb59938530c4f64077f767f67d406e9 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 22:39:36 +0000 Subject: [PATCH 142/278] Set a published model's alias inside _publish Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/trainer/trainer.py | 19 ++++++++----------- 1 file changed, 8 insertions(+), 11 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 3cb9ba08a..9a4eeef03 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -564,8 +564,14 @@ def _publish(bundle_dir: Path, bundle_id: Optional[str], ) if location and alias: try: - _set_alias(export_config.connector, - replace(location, version=published.version), alias) + aliasing = export_config.connector.sink_backend + if not isinstance(aliasing, ModelAliasingSink): + raise NotImplementedError( + f'{type(export_config.connector).__name__} cannot set ' + f'model aliases') + aliasing.set_model_alias( + cast(SourceConnector, export_config.connector), + replace(location, version=published.version), alias) except Exception as error: if warn: warn_failure( @@ -1511,15 +1517,6 @@ def _encoders(bundle_dir: Path) -> Any: if path is not None and path.is_file() else None) -def _set_alias(connector: Connector, location: ModelRegistryLocation, - alias: str) -> None: - sink = connector.sink_backend - if not isinstance(sink, ModelAliasingSink): - raise NotImplementedError( - f'{type(connector).__name__} cannot set model aliases') - sink.set_model_alias(cast(SourceConnector, connector), location, alias) - - def warn_if_other_bundle(metadata: Optional[Mapping[str, Any]], trainer: Trainer) -> None: recorded = (metadata or {}).get(BUNDLE_ID_KEY) From d2191c1a207066fb7c0c1ba50d466609f013c4d9 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 22:39:36 +0000 Subject: [PATCH 143/278] Read the fetched model's metadata directly Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/trainer/trainer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 9a4eeef03..d7674cc04 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -1468,7 +1468,7 @@ def load_from_connector( location=address.location) trainer = cls._load_bundle(bundle_dir, embedding_models, accelerator=accelerator) - warn_if_other_bundle(getattr(published, 'metadata', None), trainer) + warn_if_other_bundle(published.metadata, trainer) return trainer @classmethod From 2915e908f0e922f0dd0abc6446dc6e3adab6db76 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 22:40:09 +0000 Subject: [PATCH 144/278] Default _log_model's name and serving snapshot, and take the experiment its callers know Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/integrations/mlflow.py | 21 +++++++------------ .../kumo-engine/test/test_mlflow_model.py | 2 +- 2 files changed, 9 insertions(+), 14 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index 3130a1ad4..20cca1f45 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -280,9 +280,7 @@ def register(self, result: Any, tags={'kumo.bundle_id': str(result.id)}).info.run_id info = _log_model(self.target, client, result.bundle_dir, result.manifest or stub('model', result.id), - snapshot, name='model', run_id=run_id, - experiment_id=experiment_id, - serving_snapshot=None) + snapshot, run_id, experiment_id) with self.target.scope(): client.set_terminated(run_id, 'FINISHED') return _register(self.target, client, info, location, pointers) @@ -366,9 +364,7 @@ def log_model( os.environ.update(upload) try: info = _log_model(self.target, self._client, bundle_dir, manifest, - snapshot, name='model', run_id=run_id, - experiment_id=self._experiment_id, - serving_snapshot=None) + snapshot, run_id, self._experiment_id) finally: if self._parent == run_id: os.environ.update({name: _HTTP_LIMITS[name] for name in upload}) @@ -691,9 +687,8 @@ def log_model( _log_inputs(mlflow, client, run_id, fit_inputs(result.manifest)) info = _log_model(target, client, result.bundle_dir, result.manifest or (None if result.id is None else stub('model', result.id)), - snapshot, name=name, run_id=run_id, - experiment_id=experiment_id, - serving_snapshot=snapshot_uri(serving_snapshot)) + snapshot, run_id, experiment_id, name, + snapshot_uri(serving_snapshot)) with target.scope(): if result.metrics: client.log_batch(run_id, metrics=[ @@ -758,9 +753,9 @@ def _download(mlflow: Any, target: Any, model_uri: str, def _log_model(target: Any, client: Any, bundle_dir: Path, manifest: Optional['ModelManifest'], - snapshot: Optional['GraphSnapshot'], *, name: str, run_id: str, - experiment_id: Optional[str], - serving_snapshot: Optional[str]) -> Any: + snapshot: Optional['GraphSnapshot'], run_id: str, + experiment_id: str, name: str = 'model', + serving_snapshot: Optional[str] = None) -> Any: r"""``mlflow.pyfunc.log_model`` through ``client``.""" from kumo_core.artifacts import PROVENANCE_FILE @@ -781,7 +776,7 @@ def _log_model(target: Any, client: Any, bundle_dir: Path, provenance = bundle_dir / PROVENANCE_FILE with target.scope(): model = client.create_logged_model( - experiment_id or client.get_run(run_id).info.experiment_id, + experiment_id, name=name, source_run_id=run_id, params=dict(client.get_run(run_id).data.params)) mlflow_model = mlflow.models.Model( diff --git a/packages/kumo-engine/test/test_mlflow_model.py b/packages/kumo-engine/test/test_mlflow_model.py index 3b8459573..64a926126 100644 --- a/packages/kumo-engine/test/test_mlflow_model.py +++ b/packages/kumo-engine/test/test_mlflow_model.py @@ -363,7 +363,7 @@ def test_model_upload_gets_a_longer_http_timeout_and_a_late_one_is_dropped( location = kumo.ModelRegistryLocation('m') seen = [] - def log_model(target, client, bundle_dir, manifest, snapshot, **kwargs): + def log_model(*args): seen.append(os.environ.get('MLFLOW_HTTP_REQUEST_TIMEOUT')) if len(seen) == 2: tracker.finish(None, None) From d57699eb57182abfb6a91c5873e10a0faef1d982 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 22:41:23 +0000 Subject: [PATCH 145/278] Inline the run records' text helper and simplify the tracking pointer Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_tracking.py | 25 +++++++------------ 1 file changed, 9 insertions(+), 16 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_tracking.py b/packages/kumo-engine/kumo/engine/_tracking.py index 7b22f8585..38522dbff 100644 --- a/packages/kumo-engine/kumo/engine/_tracking.py +++ b/packages/kumo-engine/kumo/engine/_tracking.py @@ -108,8 +108,8 @@ def fit_inputs(manifest: Any) -> Dict[str, Dict[str, Optional[str]]]: if source is not None: tables['training_table'] = source return { - str(name): {key: _text(getattr(entry, key)) - for key in INPUT_FIELDS} + str(name): {key: None if getattr(entry, key) is None else str( + getattr(entry, key)) for key in INPUT_FIELDS} for name, entry in tables.items() } except Exception as error: @@ -124,10 +124,6 @@ def _parents(manifest: Any) -> Any: return getattr(manifest, 'parents', None) -def _text(value: Any) -> Optional[str]: - return None if value is None else str(value) - - def logged_name(name: str) -> str: r"""``val_auroc`` as ``val/auroc``, ``best_val_auroc`` as ``val/best_auroc``.""" @@ -174,18 +170,16 @@ def run_record(spec: RunSpec) -> RunRecord: ) -def _user_tags(tags: Optional[Mapping[str, Any]], +def _user_tags(tags: Mapping[str, Any], params: Mapping[str, str]) -> Dict[str, str]: - dropped = sorted(key for key in tags or {} + dropped = sorted(key for key in tags if key.startswith('kumo.') or key in params) if dropped: warnings.warn( f'Tracking tags {dropped} are reserved and not recorded.', UserWarning, stacklevel=3) - return { - key: str(value) - for key, value in (tags or {}).items() if key not in dropped - } + return {key: str(value) for key, value in tags.items() + if key not in dropped} def trial_spec(events: Sequence[Event]) -> RunSpec: @@ -447,11 +441,10 @@ def pointer(self) -> Optional['TrackingRef']: if self.spec is None or not self._started.wait(FINISH_SECONDS): return None tracker = self._tracker - run_id = None if tracker is None else tracker.run_id - if tracker is None or not run_id: + if tracker is None or not tracker.run_id: return None - return TrackingRef(platform=str(tracker.platform), run_id=str(run_id), - url=tracker.run_url) + return TrackingRef(platform=str(tracker.platform), + run_id=str(tracker.run_id), url=tracker.run_url) def on_event(self, event: Event) -> None: if self.spec is None or event.trial is None: From d75d1f917f476097d8c9125611204f5ae4c7c806 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 22:41:38 +0000 Subject: [PATCH 146/278] Shorten docstrings that repeat the code or docs/mlops.md Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_embeddings.py | 5 +---- packages/kumo-engine/kumo/engine/_manifest.py | 6 ++---- .../kumo-engine/kumo/engine/integrations/mlflow.py | 12 +++--------- packages/kumo-engine/kumo/engine/trainer/trainer.py | 4 +--- 4 files changed, 7 insertions(+), 20 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_embeddings.py b/packages/kumo-engine/kumo/engine/_embeddings.py index ffeb42ff7..de3a77638 100644 --- a/packages/kumo-engine/kumo/engine/_embeddings.py +++ b/packages/kumo-engine/kumo/engine/_embeddings.py @@ -116,10 +116,7 @@ def require_text_encoders(deps: Any, store: EmbeddingModelStore, def with_text_encoders(store: EmbeddingModelStore, encoders: Any, deps: Any) -> EmbeddingModelStore: - r"""``store`` plus the stored word vectors of the model's text encoders, - resolved where ``deps`` looks for assets. Unless a place was given to - look, ``store`` alone serves, as before; a missing model it does not - have either is an error.""" + r"""``store`` plus the model's stored word vectors that ``deps`` finds.""" from kumo.engine._deps import MissingDependencyError, loaded added: Dict[str, Any] = {} diff --git a/packages/kumo-engine/kumo/engine/_manifest.py b/packages/kumo-engine/kumo/engine/_manifest.py index cce259d42..c6b83ef1b 100644 --- a/packages/kumo-engine/kumo/engine/_manifest.py +++ b/packages/kumo-engine/kumo/engine/_manifest.py @@ -219,10 +219,8 @@ def _walk(value: Any, key: str, root: Manifest) -> Iterator[Artifact]: def artifact_summaries(manifest: Manifest, context: str, settings: Any = None, details: Optional[Details] = None) -> List[Summary]: - r"""Where ``manifest``'s artifacts are, for its tracked run: each - artifact's tag (the nearest of each), ``kumo.artifact_root`` under MLOps - ``settings``, and a `Dataset` per stored dataset, read in ``context`` - or written, with the schema and profile ``details`` gives its id.""" + r"""Tags locating ``manifest``'s artifacts, and a `Dataset` per stored + dataset.""" located = [item for item in artifacts(manifest) if item.location] tags: Dict[str, str] = {} for item in located: diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index 20cca1f45..3c207cd56 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -109,15 +109,9 @@ class MlflowTracking(TrackingBase): server's tracking URI. `None` uses MLflow's own tracking and registry URIs, read once here. - With explicit credentials (a Databricks service principal, an AWS - access key or role), each MLflow call sets the state MLflow reads them - from for its duration (Databricks: ``DATABRICKS_HOST`` and a short-lived - ``DATABRICKS_TOKEN``; SageMaker: the process-wide tracking URI), so do - not call MLflow or Databricks yourself on other threads during a fit. - Ambient credentials (a notebook, job, profile or instance role) set no - credentials, but a Databricks or SageMaker target still sets the - process-wide tracking URI for each call: Unity Catalog downloads and the - SageMaker plugin route with it. + With explicit credentials, each MLflow call sets process-wide state + for its duration (docs/mlops.md, "Credentials"): do not call MLflow or + Databricks yourself on other threads during a fit. ``run_id`` is the MLflow id of the latest fit run, once it has started, and ``run_url`` its link in the MLflow UI, where the server has one. diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index d7674cc04..3c63e24ef 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -600,9 +600,7 @@ def _stage_model_bundle(artifacts: Mapping[str, str], dest: Path) -> Path: def _without_unactionable_warnings(method: _F) -> _F: - r"""Runs ``method`` without third-party notes users cannot act on: - PyG's torch-scatter hint, Lightning's tips and ``__len__`` note, and - Python's fork warning.""" + r"""Runs ``method`` without third-party notes users cannot act on.""" @functools.wraps(method) def wrapper(*args: Any, **kwargs: Any) -> Any: tips = logging.getLogger('pytorch_lightning.utilities.rank_zero') From 6324879c3da4adc8617b48dadcc49d95da8ca758 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 23:03:14 +0000 Subject: [PATCH 147/278] Drop the Series, Curve, Histogram and Confusion summaries nothing produces yet Co-Authored-By: Claude Opus 5.5 --- packages/kumo-core/src/kumo_core/summaries.py | 42 +------------------ packages/kumo-engine/kumo/engine/_tracking.py | 37 +--------------- .../kumo/engine/integrations/mlflow.py | 6 --- .../kumo/engine/integrations/snowflake.py | 7 +--- .../test/test_tracking_contract.py | 34 +++------------ 5 files changed, 9 insertions(+), 117 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/summaries.py b/packages/kumo-core/src/kumo_core/summaries.py index bbb9e05dd..d7794b5e9 100644 --- a/packages/kumo-core/src/kumo_core/summaries.py +++ b/packages/kumo-core/src/kumo_core/summaries.py @@ -19,35 +19,6 @@ class Tag: value: str -@dataclass(frozen=True) -class Series: - r"""``(step, value)`` points; ``x_label`` says what a step counts.""" - - name: str - points: Tuple[Tuple[int, float], ...] - x_label: str = 'step' - - -@dataclass(frozen=True) -class Curve: - r"""``y`` against ``x``, two series over the same steps, such as ROC's - true-positive against false-positive rate over thresholds.""" - - name: str - x: Series - y: Series - - -@dataclass(frozen=True) -class Histogram: - r"""``counts[i]`` values in ``[edges[i - 1], edges[i])``: ``counts[0]`` - below ``edges[0]`` and ``counts[-1]`` at or above ``edges[-1]``.""" - - name: str - edges: Tuple[float, ...] - counts: Tuple[float, ...] - - @dataclass(frozen=True) class Table: name: str @@ -55,16 +26,6 @@ class Table: rows: Tuple[Tuple[Any, ...], ...] -@dataclass(frozen=True) -class Confusion: - r"""``counts[i][j]`` rows of actual ``labels[i]`` predicted - ``labels[j]``.""" - - name: str - labels: Tuple[str, ...] - counts: Tuple[Tuple[int, ...], ...] - - @dataclass(frozen=True) class Dataset: r"""A source table: ``source`` its ``connector``, ``location``, @@ -79,5 +40,4 @@ class Dataset: context: str -Summary = Union[Scalar, Tag, Series, Curve, Histogram, Table, Confusion, - Dataset] +Summary = Union[Scalar, Tag, Table, Dataset] diff --git a/packages/kumo-engine/kumo/engine/_tracking.py b/packages/kumo-engine/kumo/engine/_tracking.py index 38522dbff..af341fbbf 100644 --- a/packages/kumo-engine/kumo/engine/_tracking.py +++ b/packages/kumo-engine/kumo/engine/_tracking.py @@ -28,16 +28,7 @@ ) from kumo_core.instrumentation import format_error -from kumo_core.summaries import ( - Confusion, - Curve, - Dataset, - Histogram, - Scalar, - Summary, - Table, - Tag, -) +from kumo_core.summaries import Dataset, Scalar, Summary, Table, Tag from kumo.engine._events import Event, _finite_floats from kumo.engine._mlops import full_error, warn_failure @@ -275,10 +266,8 @@ def log_summaries(self, summaries: Sequence[Summary]) -> None: self._log_tag(summary) elif isinstance(summary, Dataset): self._log_dataset(summary) - elif isinstance(summary, (Table, Confusion)): - self._log_table(_as_table(summary)) else: - self._log_points(_points(summary)) + self._log_table(summary) def _log_scalar(self, summary: Scalar) -> None: raise NotImplementedError @@ -286,10 +275,6 @@ def _log_scalar(self, summary: Scalar) -> None: def _log_tag(self, summary: Tag) -> None: raise NotImplementedError - def _log_points(self, series: Mapping[str, Sequence[Tuple[int, float]]] - ) -> None: - raise NotImplementedError - def _log_table(self, summary: Table) -> None: raise NotImplementedError @@ -297,24 +282,6 @@ def _log_dataset(self, summary: Dataset) -> None: pass -def _points(summary: Summary) -> Dict[str, Sequence[Tuple[int, float]]]: - if isinstance(summary, Curve): - return {summary.x.name: summary.x.points, - summary.y.name: summary.y.points} - if isinstance(summary, Histogram): - return {f'{summary.name}/count': tuple(enumerate(summary.counts)), - f'{summary.name}/lower_edge': tuple(enumerate(summary.edges, 1))} - return {summary.name: summary.points} - - -def _as_table(summary: Summary) -> Table: - if isinstance(summary, Confusion): - return Table(summary.name, ('actual', *summary.labels), - tuple((label, *row) for label, row in zip( - summary.labels, summary.counts))) - return summary - - def table_json(table: Table) -> str: return json.dumps({'columns': list(table.columns), 'data': [list(row) for row in table.rows]}, default=str) diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index 3c207cd56..e55bbe2dc 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -448,12 +448,6 @@ def _log_dataset(self, summary: Dataset) -> None: entities.DatasetInput(dataset, tags=[entities.InputTag( 'mlflow.data.context', summary.context)])]) - def _log_points(self, series: Mapping[str, Sequence[Tuple[int, float]]] - ) -> None: - self._log_batch(self._parent, [(name, value, step) - for name, points in series.items() - for step, value in points]) - def _log_artifacts(self, artifacts: Mapping[str, str]) -> None: for name, text in artifacts.items(): if name in self._artifacts: diff --git a/packages/kumo-engine/kumo/engine/integrations/snowflake.py b/packages/kumo-engine/kumo/engine/integrations/snowflake.py index 11f71d48a..4dd6d8789 100644 --- a/packages/kumo-engine/kumo/engine/integrations/snowflake.py +++ b/packages/kumo-engine/kumo/engine/integrations/snowflake.py @@ -11,7 +11,7 @@ import uuid from contextlib import ExitStack from pathlib import Path -from typing import Any, Dict, List, Mapping, Optional, Sequence, Set, Tuple +from typing import Any, Dict, List, Mapping, Optional, Set, Tuple from kumo_core.instrumentation import format_error from kumo_core.sources.sql import resolve_snowflake_identifier @@ -264,11 +264,6 @@ def _log_tag(self, summary: Tag) -> None: def _log_table(self, summary: Table) -> None: self._log_artifacts({f'{summary.name}.json': table_json(summary)}) - def _log_points(self, series: Mapping[str, Sequence[Tuple[int, float]]] - ) -> None: - for step, metrics in by_step(series): - self._log_metrics(metrics, _SUMMARY_METRICS, step=step) - def _log_artifacts(self, artifacts: Mapping[str, str]) -> None: with tempfile.TemporaryDirectory() as directory: for name, text in artifacts.items(): diff --git a/packages/kumo-engine/test/test_tracking_contract.py b/packages/kumo-engine/test/test_tracking_contract.py index 22c72489a..b48ef5fd4 100644 --- a/packages/kumo-engine/test/test_tracking_contract.py +++ b/packages/kumo-engine/test/test_tracking_contract.py @@ -14,15 +14,7 @@ import pytest from kumo_core.artifacts import PROVENANCE_FILE -from kumo_core.summaries import ( - Confusion, - Curve, - Dataset, - Histogram, - Scalar, - Series, - Table, -) +from kumo_core.summaries import Dataset, Scalar, Table from kumo.engine._tracking import ( DROPPED_METRICS, @@ -415,12 +407,7 @@ def test_adapter_records_a_single_trial_fit_epoch_by_epoch(adapter): SUMMARIES = ( Scalar('test/brier', 0.25), - Series('search/trials', ((1, 3.0), (2, 4.0)), x_label='trial'), - Curve('test/roc', Series('test/roc/fpr', ((0, 1.0), (1, 0.5))), - Series('test/roc/tpr', ((0, 1.0), (1, 0.75)))), - Histogram('pred/True_PROB/hist', (0.5,), (3, 4)), Table('graph/col_stats', ('table', 'rows'), (('ents', 400),)), - Confusion('test/confusion', ('a', 'b'), ((5, 1), (2, 7))), Dataset('ents', INPUTS['ents'], {'ID': {'stype': 'id', 'dtype': 'int64'}}, {'rows': 400}, 'training'), ) @@ -430,18 +417,9 @@ def test_adapter_logs_each_summary_type(adapter, mlflow_adapter): replay(adapter.make(), summaries=SUMMARIES) fit = adapter.observe()[FIT] - assert {name: fit.steps[name] for name in ( - 'test/brier', 'search/trials', 'test/roc/fpr', 'test/roc/tpr', - 'pred/True_PROB/hist/count', 'pred/True_PROB/hist/lower_edge')} == { - 'test/brier': ((0, 0.25),), 'search/trials': ((1, 3.0), (2, 4.0)), - 'test/roc/fpr': ((0, 1.0), (1, 0.5)), - 'test/roc/tpr': ((0, 1.0), (1, 0.75)), - 'pred/True_PROB/hist/count': ((0, 3.0), (1, 4.0)), - 'pred/True_PROB/hist/lower_edge': ((1, 0.5),)} + assert fit.steps['test/brier'] == ((0, 0.25),) tables = {'col_stats.json': {'columns': ['table', 'rows'], - 'data': [['ents', 400]]}, - 'confusion.json': {'columns': ['actual', 'a', 'b'], - 'data': [['a', 5, 1], ['b', 2, 7]]}} + 'data': [['ents', 400]]}} if adapter is not mlflow_adapter: assert {name: json.loads(fit.artifacts[name]) for name in tables} == tables @@ -454,11 +432,9 @@ def test_adapter_logs_each_summary_type(adapter, mlflow_adapter): if r.data.tags['kumo.kind'] == 'fit') assert {name: json.loads(mlflow.artifacts.load_text( f'{run.info.artifact_uri}/{section}/{name}')) - for section, name in (('graph', 'col_stats.json'), - ('test', 'confusion.json'))} == tables + for section, name in (('graph', 'col_stats.json'),)} == tables assert json.loads(run.data.tags['mlflow.loggedArtifacts']) == [ - {'path': 'graph/col_stats.json', 'type': 'table'}, - {'path': 'test/confusion.json', 'type': 'table'}] + {'path': 'graph/col_stats.json', 'type': 'table'}] dataset, = run.inputs.dataset_inputs assert (dataset.dataset.name, dataset.dataset.digest, dataset.tags[0].value) == ('ents', '1790000000', 'training') From 64eac88507c63c22984edf492ba06ba1592fcfd2 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 00:17:03 +0000 Subject: [PATCH 148/278] Refuse a table path a file-only sink cannot write, with NotImplementedError Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_core/object_store_sink.py | 22 +++++++++++++------ .../kumo-engine/test/test_filestore_sink.py | 6 +++++ 2 files changed, 21 insertions(+), 7 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/object_store_sink.py b/packages/kumo-core/src/kumo_core/object_store_sink.py index ce4e7322d..bb4ec832d 100644 --- a/packages/kumo-core/src/kumo_core/object_store_sink.py +++ b/packages/kumo-core/src/kumo_core/object_store_sink.py @@ -427,11 +427,18 @@ def preflight_write_files(self, connector, table_name, write_mode=None): class FileStoreRouting: """Sends ``holds_files`` paths to :class:`FilesystemSink`, the rest on.""" + def _rest(self, method): + rest = super() + if not hasattr(rest, method): + raise NotImplementedError( + f'{type(self).__name__} cannot {method} outside its files') + return rest + def write_table(self, connector, table_name, frame, write_mode=WriteMode.OVERWRITE, project_id=None, *, export_id=None): sink = FilesystemSink(self) if self.holds_files( - connector, table_name) else super() + connector, table_name) else self._rest('write_table') return sink.write_table(connector, table_name, frame, write_mode, project_id, export_id=export_id) @@ -439,7 +446,7 @@ def write_files(self, connector, table_name, files, write_mode=WriteMode.OVERWRITE, project_id=None, *, export_id=None, extra_columns=None): sink = FilesystemSink(self) if self.holds_files( - connector, table_name) else super() + connector, table_name) else self._rest('write_files') return sink.write_files(connector, table_name, files, write_mode, project_id, export_id=export_id, extra_columns=extra_columns) @@ -450,9 +457,9 @@ def publish_model(self, connector, bundle_dir, *, path=None, location=None, connector, path): return FilesystemSink(self).publish_model( connector, bundle_dir, path=path, metadata=metadata) - return super().publish_model(connector, bundle_dir, path=path, - location=location, metadata=metadata, - metrics=metrics, project_id=project_id) + return self._rest('publish_model').publish_model( + connector, bundle_dir, path=path, location=location, + metadata=metadata, metrics=metrics, project_id=project_id) def fetch_model(self, connector, local_dir, *, path=None, location=None, project_id=None): @@ -460,8 +467,9 @@ def fetch_model(self, connector, local_dir, *, path=None, location=None, connector, path): return FilesystemSink(self).fetch_model(connector, local_dir, path=path) - return super().fetch_model(connector, local_dir, path=path, - location=location, project_id=project_id) + return self._rest('fetch_model').fetch_model( + connector, local_dir, path=path, location=location, + project_id=project_id) def preflight_write_files(self, connector, table_name, write_mode=None): if self.holds_files(connector, table_name): diff --git a/packages/kumo-engine/test/test_filestore_sink.py b/packages/kumo-engine/test/test_filestore_sink.py index 380640c90..8a7e36f52 100644 --- a/packages/kumo-engine/test/test_filestore_sink.py +++ b/packages/kumo-engine/test/test_filestore_sink.py @@ -182,3 +182,9 @@ def test_an_occupied_directory_is_refused_before_scoring(connector): for mode in (None, 'OVERWRITE', 'APPEND'): with pytest.raises(ArtifactDestinationConflict, match='already'): sink.preflight_write_files(connector, '@stage/out/', mode) + + +def test_a_file_only_sink_refuses_a_table(connector, tmp_path): + with pytest.raises(NotImplementedError, match='outside its files'): + connector.sink_backend.publish_model(connector, str(tmp_path), + path='DB.SCHEMA.T') From a66094bcfefced56cb7d5fe47b531bc252d04a13 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 00:19:55 +0000 Subject: [PATCH 149/278] Set MLflow's HTTP limits through the target's environ helper Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/integrations/mlflow.py | 27 +++++++------------ 1 file changed, 10 insertions(+), 17 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index e55bbe2dc..1ef983299 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -52,6 +52,7 @@ from kumo.engine.integrations._mlflow_target import ( SAGEMAKER_ARN, SageMakerMLflow, + environ, global_tracking_uri, ) from kumo.engine.integrations._mlflow_target import resolve as resolve_target @@ -182,10 +183,9 @@ def check(self, registry: Optional[ModelRegistryLocation] = None) -> None: if registry is not None: self.check_model_name(registry.model_name) uri = self.target.tracking_uri - owned = _unset_limits() missing = False try: - with _limited_http(owned), self.target.scope(): + with _limited_http(), self.target.scope(): client = self.target.client() if self.experiment_name is not None: missing = client.get_experiment_by_name( @@ -203,7 +203,7 @@ def check(self, registry: Optional[ModelRegistryLocation] = None) -> None: f'credentials: {full_error(error)}') from error if not missing: return - with _limited_http(owned), self.target.scope(): + with _limited_http(), self.target.scope(): try: client.create_experiment(self.experiment_name) except Exception as error: @@ -258,7 +258,7 @@ def register(self, result: Any, self._require_experiment() client = self.target.client() pointers = registry_pointers(result.id) - with _limited_http(_unset_limits()): + with _limited_http(): with self.target.scope(): held = _registered_version(client, location.model_name, pointers) @@ -283,8 +283,7 @@ def start(self, record: RunRecord) -> None: self._require_experiment() self._reset() self.run_id = None - self._http_owned = _unset_limits() - self._guard.enter_context(_limited_http(self._http_owned)) + self._http_owned = self._guard.enter_context(_limited_http()) try: autologging = import_provider('mlflow.utils.autologging_utils', 'mlflow') @@ -989,15 +988,9 @@ def _quietly(function: Callable[..., Any], *args: Any) -> None: logger.debug('MLflow cleanup call failed', exc_info=True) -def _unset_limits() -> FrozenSet[str]: - return frozenset(name for name in _HTTP_LIMITS if name not in os.environ) - - @contextmanager -def _limited_http(names: FrozenSet[str]) -> Iterator[None]: - os.environ.update({name: _HTTP_LIMITS[name] for name in names}) - try: - yield - finally: - for name in names: - os.environ.pop(name, None) +def _limited_http() -> Iterator[FrozenSet[str]]: + r"""MLflow's HTTP limits where unset; yields the names it set.""" + owned = frozenset(name for name in _HTTP_LIMITS if name not in os.environ) + with environ({name: _HTTP_LIMITS[name] for name in owned}): + yield owned From d879ab5fe8c48a26c3c703a81f10c6d90503a174 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 00:20:02 +0000 Subject: [PATCH 150/278] Read the alias of a Snowflake fetch location as a field Co-Authored-By: Claude Opus 5.5 --- packages/kumo-snowflake/src/kumo_snowflake/registry.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/packages/kumo-snowflake/src/kumo_snowflake/registry.py b/packages/kumo-snowflake/src/kumo_snowflake/registry.py index c4d88b580..dd47f3c8c 100644 --- a/packages/kumo-snowflake/src/kumo_snowflake/registry.py +++ b/packages/kumo-snowflake/src/kumo_snowflake/registry.py @@ -455,13 +455,12 @@ def fetch_model( # Resolved before anything on disk is touched: a name the registry # does not know must not cost the caller the bundle already there. model = registry.get_model(location.model_name) - alias = getattr(location, 'alias', None) - if alias is not None: - if resolved(alias) not in _alias_holders(model): + if location.alias is not None: + if resolved(location.alias) not in _alias_holders(model): raise ValueError( f'No version of {location.model_name} holds alias ' - f'{alias}.') - version = model.version(alias) + f'{location.alias}.') + version = model.version(location.alias) elif location.version is None: version = model.default else: From fef887f6bae526f1833ef71edd138f4cdab6d0ba Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 00:20:16 +0000 Subject: [PATCH 151/278] Drop the Snowflake hung-call test the tracking contract's hung-tracker test covers Co-Authored-By: Claude Opus 5.5 --- .../test/test_snowflake_tracking.py | 29 +------------------ 1 file changed, 1 insertion(+), 28 deletions(-) diff --git a/packages/kumo-engine/test/test_snowflake_tracking.py b/packages/kumo-engine/test/test_snowflake_tracking.py index 37239caac..43dc694da 100644 --- a/packages/kumo-engine/test/test_snowflake_tracking.py +++ b/packages/kumo-engine/test/test_snowflake_tracking.py @@ -5,7 +5,6 @@ import inspect import re import threading -import time from unittest import mock import _snowflake_experiment_fake @@ -13,12 +12,7 @@ from _snowflake_experiment_fake import FakeExperimentTracking, FakeSession from kumoapi.model_plan import ModelPlan from test_fit_events import table_and_snapshot # noqa: F401 -from test_tracking_contract import ( - CONTEXT, - EVENTS, - RESULT, - join_tracking_threads, -) +from test_tracking_contract import CONTEXT, EVENTS, RESULT from kumo.engine._tracking import TrackingSession from kumo.engine.trainer._run_spec import fit_spec, fitted_spec @@ -219,27 +213,6 @@ def test_the_users_session_is_not_altered(fake): assert session.statements == [] -def test_a_hung_snowflake_call_does_not_stall_the_fit(fake, monkeypatch): - import kumo.engine._tracking as tracking - - monkeypatch.setattr(tracking, 'FINISH_SECONDS', 0.5) - release = threading.Event() - monkeypatch.setattr(FakeExperimentTracking, 'log_metrics', - lambda *args, **kwargs: release.wait(30)) - began = time.monotonic() - try: - session = _replay(_tracking()) - elapsed = time.monotonic() - began - finally: - release.set() - join_tracking_threads() - - assert elapsed < 1.5 - assert isinstance(session.error, TimeoutError) - run, = fake._instance.runs.values() - assert run['status'] == 'FINISHED' - - class _SchemaSession: def __init__(self, schemas): self.schemas = schemas From bbad2a43aa09e010461aeddc07567344090b78c8 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 00:35:21 +0000 Subject: [PATCH 152/278] Build MLflow metric batches in _log_metrics Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/integrations/mlflow.py | 21 +++++++------------ 1 file changed, 7 insertions(+), 14 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index 1ef983299..5a0609d93 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -27,7 +27,6 @@ Literal, Mapping, Optional, - Sequence, Set, Tuple, ) @@ -489,25 +488,19 @@ def _log_params(self, run_id: str, params: Mapping[str, str]) -> None: def _log_metrics(self, run_id: Optional[str], metrics: Mapping[str, float], step: int = 0, model_id: Optional[str] = None) -> None: - if run_id is not None: - self._log_batch(run_id, [(name, value, step) - for name, value in metrics.items()], - model_id) - - def _log_batch(self, run_id: Optional[str], - points: Sequence[Tuple[str, float, int]], - model_id: Optional[str] = None) -> None: + if run_id is None or not metrics: + return entities = self._mlflow.entities timestamp = int(time.time() * 1000) - metrics = [ - entities.Metric(metric_key(name), float(value), timestamp, step, + batch = [ + entities.Metric(metric_key(key), float(value), timestamp, step, model_id=model_id) - for name, value, step in points + for key, value in metrics.items() ] with self.target.scope(): - for start in range(0, len(metrics), _BATCH_METRICS): + for start in range(0, len(batch), _BATCH_METRICS): self._client.log_batch( - run_id, metrics=metrics[start:start + _BATCH_METRICS]) + run_id, metrics=batch[start:start + _BATCH_METRICS]) _BATCH_METRICS = 1000 # MLflow's per-request limit From 5a3b43014b83447bc188bb9fac412c8173d3ca5c Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 00:35:28 +0000 Subject: [PATCH 153/278] Log log_model's metrics through MlflowTracking._log_metrics Co-Authored-By: Claude Opus 5.5 --- .../kumo-engine/kumo/engine/integrations/mlflow.py | 14 ++++---------- 1 file changed, 4 insertions(+), 10 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index 5a0609d93..28f1edbc1 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -653,7 +653,7 @@ def log_model( target = tracking.target if run_id is None: tracking._require_experiment() - client = target.client() + tracking._client = client = target.client() active = mlflow.active_run() if run_id is None and active is not None and ( mlflow.get_tracking_uri() == target.tracking_uri): @@ -669,15 +669,9 @@ def log_model( (None if result.id is None else stub('model', result.id)), snapshot, run_id, experiment_id, name, snapshot_uri(serving_snapshot)) - with target.scope(): - if result.metrics: - client.log_batch(run_id, metrics=[ - mlflow.entities.Metric(metric_key(key), float(value), - int(time.time() * 1000), 0, - model_id=info.model_id) - for key, value in result.metrics.items() - ]) - if created: + tracking._log_metrics(run_id, result.metrics, model_id=info.model_id) + if created: + with target.scope(): client.set_terminated(run_id, 'FINISHED') return info From e4d390b36187b21f3a66b15976fea498eaaa4b19 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 00:50:03 +0000 Subject: [PATCH 154/278] Trim docstrings, comments and messages in the registry and core packages Co-Authored-By: Claude Opus 5.5 --- packages/kumo-api/kumoapi/plugin.py | 25 ++++-------- .../src/kumo_core/artifact_export.py | 3 +- .../src/kumo_core/automl/optuna_searcher.py | 1 - .../src/kumo_core/diskgraph/lightning_data.py | 3 +- packages/kumo-core/src/kumo_core/summaries.py | 5 +-- .../src/kumo_core/train/formatting.py | 3 +- .../src/kumo_databricks/registry.py | 9 ++--- .../src/kumo_databricks/session.py | 3 +- .../kumo-databricks/tests/test_registry.py | 1 - .../src/kumo_snowflake/registry.py | 39 ++++++------------- 10 files changed, 27 insertions(+), 65 deletions(-) diff --git a/packages/kumo-api/kumoapi/plugin.py b/packages/kumo-api/kumoapi/plugin.py index 11153975c..633effb05 100644 --- a/packages/kumo-api/kumoapi/plugin.py +++ b/packages/kumo-api/kumoapi/plugin.py @@ -66,12 +66,8 @@ '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) @@ -170,8 +166,7 @@ class ModelRegistryLocation: 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`. + alias: Set on the new version on export; read on import. """ model_name: str version: Optional[str] = None @@ -740,8 +735,7 @@ def fetch_model( @runtime_checkable class ModelAliasingSink(ArtifactSink, Protocol): - r"""A registry sink that sets aliases and fetches by - :attr:`ModelRegistryLocation.alias`.""" + r"""A registry sink that sets and fetches by alias.""" def set_model_alias( self, connector: SourceConnector, @@ -753,22 +747,17 @@ def set_model_alias( 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`` in any case.""" + 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 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 ' - f'(for DEFAULT use ALTER MODEL ... SET DEFAULT_VERSION).') + 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.') 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`.""" + r"""A version's ``kumo_*`` registry metadata, without the marker.""" return { key: str(value) for key, value in recorded.items() diff --git a/packages/kumo-core/src/kumo_core/artifact_export.py b/packages/kumo-core/src/kumo_core/artifact_export.py index 16d80cca7..4c99e5927 100644 --- a/packages/kumo-core/src/kumo_core/artifact_export.py +++ b/packages/kumo-core/src/kumo_core/artifact_export.py @@ -304,8 +304,7 @@ def prepared_artifact_file(path, extra_columns=None, *, parquet_options=None, writer = None try: for batch in local_parquet_batches(path): - # In Arrow: a pandas round trip turns an all-null bool column - # into Arrow's null type. + # Pandas turns all-null bool columns into Arrow nulls. table = pa.Table.from_batches([batch]) for name, value in (extra_columns or {}).items(): column = pa.array([value] * table.num_rows, diff --git a/packages/kumo-core/src/kumo_core/automl/optuna_searcher.py b/packages/kumo-core/src/kumo_core/automl/optuna_searcher.py index c6e4b5fb8..a758e1c44 100644 --- a/packages/kumo-core/src/kumo_core/automl/optuna_searcher.py +++ b/packages/kumo-core/src/kumo_core/automl/optuna_searcher.py @@ -85,7 +85,6 @@ def sample_value(self, key: str, value: list[Any]) -> Any: return self.current_trial.suggest_categorical(key, value) def analyze(self, path: str) -> None: - # Optuna flags its matplotlib plots as experimental on every call. with warnings.catch_warnings(): warnings.simplefilter("ignore", optuna.exceptions.ExperimentalWarning) diff --git a/packages/kumo-core/src/kumo_core/diskgraph/lightning_data.py b/packages/kumo-core/src/kumo_core/diskgraph/lightning_data.py index aa28e8e73..229464e1c 100644 --- a/packages/kumo-core/src/kumo_core/diskgraph/lightning_data.py +++ b/packages/kumo-core/src/kumo_core/diskgraph/lightning_data.py @@ -237,7 +237,7 @@ def open(cls, process: Any) -> _WorkerHandle: # observation for a worker that won the startup/exit race. pass except OSError as error: - # Sandboxes such as SageMaker hosting refuse pidfd_open. + # SageMaker hosting refuses pidfd_open. if error.errno not in (errno.ENOSYS, errno.EPERM): os.close(sentinel_fd) raise @@ -268,7 +268,6 @@ def send_signal(self, sig: int) -> None: if self.pidfd is not None: _pidfd_send_signal(self.pidfd, sig) else: - # Without a pidfd the live sentinel above guards the pid. os.kill(self.pid, sig) except ProcessLookupError: pass diff --git a/packages/kumo-core/src/kumo_core/summaries.py b/packages/kumo-core/src/kumo_core/summaries.py index d7794b5e9..55064547e 100644 --- a/packages/kumo-core/src/kumo_core/summaries.py +++ b/packages/kumo-core/src/kumo_core/summaries.py @@ -28,10 +28,7 @@ class Table: @dataclass(frozen=True) class Dataset: - r"""A source table: ``source`` its ``connector``, ``location``, - ``version`` and ``version_kind``; ``schema`` each column's ``stype`` and - ``dtype``; ``profile`` its rows and column statistics; ``context`` - ``training`` or ``scoring``.""" + r"""A source table's location, schema, profile and context.""" name: str source: Mapping[str, Optional[str]] diff --git a/packages/kumo-core/src/kumo_core/train/formatting.py b/packages/kumo-core/src/kumo_core/train/formatting.py index 0d5585d73..ecc129b7c 100644 --- a/packages/kumo-core/src/kumo_core/train/formatting.py +++ b/packages/kumo-core/src/kumo_core/train/formatting.py @@ -77,8 +77,7 @@ def _get_col_names( 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``.""" + r"""The prediction columns a scored row carries after its entity.""" if task_type in _CLASS_TASKS: return ['PREDICTED', 'CLASS', 'SCORE'] if task_type.is_ranking: diff --git a/packages/kumo-databricks/src/kumo_databricks/registry.py b/packages/kumo-databricks/src/kumo_databricks/registry.py index 9c7f9e3ab..8f1e3d11e 100644 --- a/packages/kumo-databricks/src/kumo_databricks/registry.py +++ b/packages/kumo-databricks/src/kumo_databricks/registry.py @@ -47,9 +47,9 @@ def _address(connector: SourceConnector, path: Optional[str], schema=config.schema_name) if len(name) > 255 or _UC_FORBIDDEN.search(name): raise ValueError( - f'Unity Catalog cannot register a model named {name!r}: use at ' - f'most 255 characters with no space, "/", "." or control ' - f"character, e.g. 'users_churn'.") + f'Unity Catalog cannot register a model named {name!r}: use up to ' + f'255 characters without spaces, "/" or ".", e.g. ' + f"'users_churn'.") return f'{config.catalog}.{schema}.{name}' @@ -338,8 +338,7 @@ class DatabricksModelRegistryMixin(_DatabricksSessions): """Publish, discover and reload bundles; prediction remains in Kumo.""" def check_model_location(self, connector: SourceConnector, location: ModelRegistryLocation) -> None: - """Raise ``ValueError`` unless ``location`` names a model this - connector's catalog and schema can hold.""" + """Raises ``ValueError`` unless this schema can hold ``location``.""" _address(connector, None, location) def publish_model( diff --git a/packages/kumo-databricks/src/kumo_databricks/session.py b/packages/kumo-databricks/src/kumo_databricks/session.py index ceb299d0b..207877dde 100644 --- a/packages/kumo-databricks/src/kumo_databricks/session.py +++ b/packages/kumo-databricks/src/kumo_databricks/session.py @@ -91,8 +91,7 @@ def _client_for(principal: Principal) -> WorkspaceClient: f"Identity {principal.name!r} holds " f"{type(credentials).__name__} credentials; Databricks " f"authenticates a service principal with an OAuth secret.") - # Explicit: a DATABRICKS_TOKEN another thread lends MLflow must not - # make this client's configuration ambiguous. + # Explicit, so a DATABRICKS_TOKEN set for MLflow cannot leak in. return WorkspaceClient( host=config.host, client_id=config.client_id, diff --git a/packages/kumo-databricks/tests/test_registry.py b/packages/kumo-databricks/tests/test_registry.py index e0172b877..e42bc4019 100644 --- a/packages/kumo-databricks/tests/test_registry.py +++ b/packages/kumo-databricks/tests/test_registry.py @@ -227,7 +227,6 @@ def list_directory(self, path): relative = key[len(path) + 1:] child = relative.split('/')[0] directory = '/' in relative - # The Files API lists directories with a trailing slash. children[child] = SimpleNamespace( path=f"{path}/{child}{'/' if directory else ''}", is_directory=directory) diff --git a/packages/kumo-snowflake/src/kumo_snowflake/registry.py b/packages/kumo-snowflake/src/kumo_snowflake/registry.py index dd47f3c8c..cca96ef2d 100644 --- a/packages/kumo-snowflake/src/kumo_snowflake/registry.py +++ b/packages/kumo-snowflake/src/kumo_snowflake/registry.py @@ -123,11 +123,8 @@ class _BundleModel(custom_model.CustomModel): Deliberately not runnable. Scoring a Kumo model needs a materialized DiskGraph task over the graph being scored, built where the data is, so this refuses rather than answering. The version exists to carry, - version and expose the bundle. Load it with - ``Trainer.load_from_connector(SnowflakeConnector(...), - location=ModelRegistryLocation(...))`` and predict against a snapshot - (kumo-engine), or with :func:`kumoai_studio.import_model` (hosted - service). + version and expose the bundle; load it with + ``Trainer.load_from_connector``. """ # Unquoted, and evaluated here rather than deferred: the registry # compares this annotation against `pandas.DataFrame` itself, so a @@ -135,17 +132,13 @@ class _BundleModel(custom_model.CustomModel): @custom_model.inference_api def predict(self, df: pd.DataFrame) -> pd.DataFrame: raise NotImplementedError( - "This model version carries a Kumo model bundle; it does not " - "score in Snowflake. Load it with Trainer.load_from_connector(" - "SnowflakeConnector(...), location=ModelRegistryLocation(...)) " - "and predict against a snapshot (kumo-engine), or " - "kumoai_studio.import_model() (hosted service).") + "This model version holds a Kumo bundle and does not score in " + "Snowflake; load it with Trainer.load_from_connector.") return _BundleModel( custom_model.ModelContext(artifacts={'bundle': bundle_dir})) -#: Snowflake ML task types by Kumo task type; Model Monitor needs one. _TASKS = { 'binary_classification': 'TABULAR_BINARY_CLASSIFICATION', 'multiclass_classification': 'TABULAR_MULTI_CLASSIFICATION', @@ -154,8 +147,7 @@ def predict(self, df: pd.DataFrame) -> pd.DataFrame: def _task(bundle: Path) -> Dict[str, Any]: - r"""``task=`` for ``log_model`` from the bundle's prediction schema; - empty for other tasks, or when it cannot be read.""" + r"""``log_model``'s ``task=`` for the bundle, if Snowflake has one.""" try: from snowflake.ml.model import type_hints @@ -302,8 +294,7 @@ def _session( def check_model_location(self, connector: SourceConnector, location: ModelRegistryLocation) -> None: - r"""Raises :class:`ValueError` unless ``location.model_name`` is a - Snowflake identifier.""" + r"""Raises ``ValueError`` unless the model name is an identifier.""" del connector if not _MODEL_NAME.fullmatch(location.model_name): raise ValueError( @@ -350,8 +341,7 @@ def publish_model( the model, and optionally the version. A version of ``None`` lets the registry assign one -- the usual case, and the reason this reports back what it created. - metadata: Returned with the publication result; its - ``kumo_bundle_id`` is also recorded in the version's metrics. + metadata: Returned with the result; its bundle id is recorded. metrics: Returned with the publication result. Training metrics remain in the bundle's metrics.json. project_id: Ignored; a local runtime has no projects. @@ -397,7 +387,6 @@ def list_models( That marker makes discovery both cheap and precise: versions logged by another tool are omitted, and no bundle is downloaded merely to decide whether it belongs in this list. - ``metadata`` carries the recorded ``kumo_bundle_id``. """ del project_id found: List[PublishedModel] = [] @@ -433,15 +422,13 @@ def fetch_model( though the caller always hands this a fresh, empty ``mkdtemp``. path: Refused, as in :meth:`publish_model`. - location: The model, and optionally the version or an alias. With - neither, takes the registry's default version, which is what a - caller who names none means. + location: The model, and optionally a version or alias (else the + registry's default version). project_id: Ignored. Returns: The name, the version actually read -- resolved, when the caller - named none -- where it lives, metrics read from metrics.json, and - the version's recorded ``kumo_bundle_id`` as ``metadata``. + named none -- where it lives, metrics and the recorded bundle id. """ del project_id _refuse_path(path, 'fetching') @@ -497,11 +484,7 @@ def set_model_alias( location: ModelRegistryLocation, alias: str, ) -> None: - r"""Points :obj:`alias` at ``location.version``, moving it if held. - - Snowflake allows one alias per version. A failed move puts the alias - back on the version that held it. - """ + r"""Points :obj:`alias` at ``location.version``, moving it if held.""" if location.version is None: raise ValueError('Setting an alias needs a ModelRegistryLocation ' 'naming an exact version.') From e0907a2e6cd3b056c3aefc1a19c7ffc15d14dcd9 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 00:50:03 +0000 Subject: [PATCH 155/278] Trim docstrings, comments and messages in the tracker integrations Co-Authored-By: Claude Opus 5.5 --- .../engine/integrations/_mlflow_target.py | 13 +-- .../kumo/engine/integrations/mlflow.py | 108 +++++------------- .../kumo/engine/integrations/snowflake.py | 51 ++------- .../test/_snowflake_experiment_fake.py | 3 +- .../kumo-engine/test/test_mlflow_target.py | 16 +-- .../kumo-engine/test/test_mlflow_tracking.py | 4 +- .../test/test_snowflake_tracking.py | 6 +- 7 files changed, 55 insertions(+), 146 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py b/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py index 8a6ee9b9e..fb956759d 100644 --- a/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py +++ b/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py @@ -44,15 +44,10 @@ def _without_header_warnings() -> Iterator[None]: @dataclass(frozen=True) class SageMakerMLflow: - r"""A SageMaker managed MLflow tracking server or app, for - ``MlflowTracking(target=...)``. + r"""A SageMaker managed MLflow server, for ``MlflowTracking(target=)``. Args: - tracking_server_arn: the server's (or app's) ARN. - identity: a ``kumo.engine.connector.S3Identity`` with an access key - or a role to assume; `None` or ``S3Identity()`` uses the ambient - AWS credential chain (an instance or execution role, - ``aws configure``). + identity: an ``S3Identity``; `None` uses the ambient AWS credentials. """ tracking_server_arn: str @@ -148,7 +143,7 @@ def resolve(target: Any) -> Target: if isinstance(target, DatabricksConnector): return _databricks(target) raise TypeError(f'MlflowTracking(target=...) takes a DatabricksConnector, ' - f'a SageMakerMLflow or a tracking URI, not ' + f'SageMakerMLflow or tracking URI, not ' f'{type(target).__name__}.') @@ -190,7 +185,6 @@ def _databricks(connector: Any) -> Target: @contextmanager def _databricks_token(client: Any) -> Iterator[None]: - r"""A short-lived ``DATABRICKS_TOKEN``; the secret stays out.""" config_type = import_provider('databricks.sdk.config', 'databricks').Config authorization = client.config.authenticate()['Authorization'] @@ -223,7 +217,6 @@ def _sagemaker(target: SageMakerMLflow) -> Target: @contextmanager def _sagemaker_session(arn: str, session: Any) -> Iterator[None]: - r"""The SageMaker plugin's state, signing with ``session``.""" plugin = import_provider('sagemaker_mlflow.session', 'sagemaker') with environ({'SAGEMAKER_MLFLOW_ASSUME_ROLE_ARN': None}), \ global_tracking_uri(arn), plugin.use_session(session): diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index 28f1edbc1..ea1dd0e67 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -1,8 +1,7 @@ # 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 model logging for `Trainer.fit`.""" import hashlib import json @@ -69,7 +68,6 @@ _MAX_PARAM_KEY = 250 _MAX_PARAM_VALUE = 6000 -# MLflow's Windows-safe charset. _PARAM_KEY = re.compile(r'[/\w.\- ]+') _HTTP_LIMITS = { 'MLFLOW_HTTP_REQUEST_MAX_RETRIES': '2', @@ -82,39 +80,10 @@ class MlflowTracking(TrackingBase): - r"""Records `Trainer.fit` in MLflow: one run per fit (under the active - run, if any, on the same server) with a child run per trial attempt. - Source tables are logged as run inputs: Databricks tables as Delta - tables (Unity 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. + r"""Records `Trainer.fit` in MLflow, one run per fit. Args: - experiment_name: experiment for the fit run, created if missing. - `None` uses MLflow's current experiment; required to record - runs at a ``target``. - run_name: name of the fit run. `None` names it after its registry - model and UTC start, as ``fit-users-churn-20261027-0618``. - log_system_metrics: sample CPU, memory, disk, network and (with - ``nvidia-ml-py``) GPU usage on the fit run. - tags: extra tags for the fit run. - log_model: log the fitted model on the fit run (see `log_model`) - before the fit returns. ``kumo.MLOps(registry= - ModelRegistryLocation(...))`` also registers it, which implies - this. - target: where to record, with the credentials that reach it: a - ``DatabricksConnector`` (Databricks MLflow and Unity Catalog, - as its identity), a `SageMakerMLflow`, or another MLflow - server's tracking URI. `None` uses MLflow's own tracking and - registry URIs, read once here. - - With explicit credentials, each MLflow call sets process-wide state - for its duration (docs/mlops.md, "Credentials"): do not call MLflow or - Databricks yourself on other threads during a fit. - - ``run_id`` is the MLflow id of the latest fit run, once it has started, - and ``run_url`` its link in the MLflow UI, where the server has one. + target: a ``DatabricksConnector``, `SageMakerMLflow` or tracking URI. """ def __init__( @@ -164,10 +133,7 @@ def run_url(self) -> Optional[str]: @contextmanager def client(self) -> Iterator[Any]: - r"""An ``MlflowClient`` for this tracker's server, with its - credentials and tracking URI set for the block, so your own MLflow - calls (the client's, or ``mlflow.search_runs``) reach the server - kumo-engine records to, as the same identity.""" + r"""An ``MlflowClient`` with this tracker's credentials.""" with self.target.scope(), global_tracking_uri( self.target.tracking_uri): yield self.target.client() @@ -210,10 +176,8 @@ def check(self, registry: Optional[ModelRegistryLocation] = None) -> None: return raise RuntimeError( f'Could not create the MLflow experiment ' - f'{self.experiment_name!r}: {full_error(error)}. Its ' - f'parent folder must exist: on Databricks, create it ' - f'with `databricks workspace mkdirs ` or in the ' - f'workspace UI.') from error + f'{self.experiment_name!r}: {full_error(error)}; create ' + f'its parent folder first.') from error def prepare(self) -> None: active = self._mlflow.active_run() @@ -230,16 +194,15 @@ def _experiment_for(self, client: Any) -> str: def _require_experiment(self) -> None: if self._named and self.experiment_name is None: - raise ValueError('MlflowTracking(target=...) records runs in an ' - 'experiment: pass experiment_name.') + raise ValueError('MlflowTracking(target=...) needs ' + 'experiment_name.') def load_model(self, location: ModelRegistryLocation, **kwargs: Any) -> 'Trainer': r"""The `Trainer` of ``location``'s alias or version.""" if (location.alias is None) == (location.version is None): - raise ValueError('Loading from an MLflow registry takes a ' - 'ModelRegistryLocation with an alias or a ' - 'version.') + raise ValueError('Pass a ModelRegistryLocation with either an ' + 'alias or a version.') reference = (f'@{location.alias}' if location.alias else f'/{location.version}') return load_model(f'models:/{location.model_name}{reference}', @@ -247,8 +210,7 @@ def load_model(self, location: ModelRegistryLocation, def register(self, result: Any, location: ModelRegistryLocation) -> PublishedModel: - r"""Registers a `TrainingResult` at ``location``, reusing a version - that holds its bundle.""" + r"""Registers a `TrainingResult` at ``location``.""" from kumo.engine._manifest import stub from kumo.engine.graph import GraphSnapshot from kumo.engine.trainer._spec import read_training_spec @@ -312,7 +274,7 @@ def start(self, record: RunRecord) -> None: _SystemMetrics(self.log_metrics)) self._log_params(self._parent, record.params) self._log_artifacts(record.artifacts) - if record.tags.get('kumo.kind') != 'fit': # a fit logs Datasets + if record.tags.get('kumo.kind') != 'fit': with self.target.scope(): _log_inputs(self._mlflow, client, self._parent, record.inputs) @@ -503,7 +465,7 @@ def _log_metrics(self, run_id: Optional[str], run_id, metrics=batch[start:start + _BATCH_METRICS]) -_BATCH_METRICS = 1000 # MLflow's per-request limit +_BATCH_METRICS = 1000 _SYSTEM_METRICS_FAILURES = 6 _COLUMN_TYPES = {'int64': 'long', 'int32': 'integer', 'float64': 'double', 'float32': 'float', 'bool': 'boolean', 'date': 'datetime', @@ -511,8 +473,6 @@ def _log_metrics(self, run_id: Optional[str], class _SystemMetrics: - r"""MLflow's system monitors, logged with the tracker's credentials.""" - def __init__(self, log: Callable[[Mapping[str, float], int], None], interval: float = 10.0) -> None: self._log = log @@ -585,9 +545,7 @@ def check_model_name(name: str, registry_uri: str) -> None: if not _SAGEMAKER_MODEL.fullmatch(name): raise ValueError( f'SageMaker MLflow cannot register a model named {name!r}: ' - f'use letters, digits and hyphens, starting and ending with ' - f'a letter or digit, at most 57 letters and digits ' - f'(^[a-zA-Z0-9](-*[a-zA-Z0-9]){{0,56}}$), e.g. ' + f'use up to 57 letters and digits joined by hyphens, e.g. ' f"'users-churn'.") elif is_databricks_unity_catalog_uri(registry_uri or ''): parts = name.split('.') @@ -596,9 +554,7 @@ def check_model_name(name: str, registry_uri: str) -> None: for part in parts): raise ValueError( f'Unity Catalog cannot register a model named {name!r}: use ' - f'catalog.schema.model (on a cluster, a shorter name takes ' - f'its current catalog and schema), each part at most 255 ' - f'characters with no space, "/" or control character, e.g. ' + f'catalog.schema.model without spaces or "/", e.g. ' f"'main.ml.users_churn'.") @@ -631,16 +587,13 @@ def log_model( target: Any = None, experiment_name: Optional[str] = None, ) -> Any: - r"""Log a trainer's bundle as an MLflow pyfunc model in ``run_id`` (or - the active run, else a new one in ``experiment_name``) with its source - tables as run inputs; returns its `ModelInfo`. ``target`` says where, - with what credentials, as in `MlflowTracking`; `None` uses MLflow's own - tracking URI. Scores with ``predict(entities, params={'snapshot': - uri})``; with ``serving_snapshot`` (an exported snapshot URI, or a - `GraphSnapshot` or snapshot id stored under the artifact root of the - enclosing `kumo.mlops` block), each serving replica downloads that - snapshot once at startup and scores against it when no ``snapshot`` - param is given.""" + r"""Logs a trainer as an MLflow pyfunc model; returns its `ModelInfo`. + + Args: + serving_snapshot: snapshot to score against when the request names + none. + target: as in `MlflowTracking`. + """ from kumo.engine._manifest import stub from kumo.engine._mlops import snapshot_uri from kumo.engine._tracking import fit_inputs @@ -683,9 +636,11 @@ def load_model( target: Any = None, accelerator: Literal['cpu', 'gpu', 'mps', 'auto'] = 'auto', ) -> 'Trainer': - r"""Restore the `Trainer` of a model `log_model` logged, from the - server ``target`` reaches, as in `MlflowTracking` (`None`: MLflow's own - tracking URI).""" + r"""Restores the `Trainer` of a model `log_model` logged. + + Args: + target: as in `MlflowTracking`. + """ from kumo.engine._scratch import new_run_id, trainer_scratch_dir from kumo.engine.trainer.trainer import Trainer, warn_if_other_bundle @@ -730,7 +685,6 @@ def _log_model(target: Any, client: Any, bundle_dir: Path, snapshot: Optional['GraphSnapshot'], run_id: str, experiment_id: str, name: str = 'model', serving_snapshot: Optional[str] = None) -> Any: - r"""``mlflow.pyfunc.log_model`` through ``client``.""" from kumo_core.artifacts import PROVENANCE_FILE from kumo.engine.integrations import _mlflow_model, _mlflow_model_code @@ -873,7 +827,6 @@ def _register(target: Any, client: Any, info: Any, def _registered_version(client: Any, name: str, metadata: Mapping[str, str]) -> Optional[str]: - r"""The version of ``name`` that already holds this bundle, if any.""" bundle = metadata.get(BUNDLE_ID_KEY) if not bundle: return None @@ -887,7 +840,7 @@ def _registered_version(client: Any, name: str, return None raise for version in page: - # Unity Catalog's search results carry tags as a method. + # Unity Catalog returns tags as a method. tags = version.tags if not isinstance(tags, Mapping) or not tags: tags = client.get_model_version(name, version.version).tags @@ -924,9 +877,8 @@ def check_registry_schema(name: str, registry_uri: str) -> None: raise ValueError( f"Unity Catalog cannot register {name!r}: " f"{body.get('error_code', response.status_code)}: " - f"{body.get('message', response.reason)}. The catalog and schema " - f"must exist, and the registering principal needs USE CATALOG, " - f"USE SCHEMA and CREATE MODEL on them.") + f"{body.get('message', response.reason)}; create the schema or " + f"grant USE CATALOG, USE SCHEMA and CREATE MODEL on it.") response.raise_for_status() diff --git a/packages/kumo-engine/kumo/engine/integrations/snowflake.py b/packages/kumo-engine/kumo/engine/integrations/snowflake.py index 4dd6d8789..8dc4aa931 100644 --- a/packages/kumo-engine/kumo/engine/integrations/snowflake.py +++ b/packages/kumo-engine/kumo/engine/integrations/snowflake.py @@ -1,8 +1,7 @@ # 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`.""" import logging import posixpath @@ -38,31 +37,11 @@ class SnowflakeTracking(TrackingBase): - r"""Records `Trainer.fit` in Snowflake Experiment Tracking as one run; - 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. + r"""Records `Trainer.fit` in Snowflake Experiment Tracking as one run. Args: - source: a ``SnowflakeConnector``, whose identity opens the session - and whose Snowflake Model Registry ``kumo.MLOps(registry=...)`` - registers models in; or a Snowpark session, which tracks only. - It needs CREATE EXPERIMENT on the schema, or MODIFY on an - existing experiment. - experiment_name: experiment for the fit run, created if missing. - database: `None` uses the session's current database. - schema_name: `None` uses the session's current schema, or ``PUBLIC`` - when `database` is given or the session has none. - run_name: name of the fit run. `None` names it after its registry - model and UTC start, as ``fit_users_churn_20261027_0618``; a - taken name gets a random suffix. - tags: extra tags for the fit run. - 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. + source: a ``SnowflakeConnector``, or a Snowpark session (tracking only). + native_model_lineage: link models registered during the fit to its run. """ platform = 'snowflake' @@ -108,8 +87,7 @@ def _reset(self) -> None: self._dropped_params = 0 def load_model(self, location: Any, **kwargs: Any) -> Any: - r"""The `Trainer` of the Snowflake Model Registry version - ``location`` names, through this tracker's connector.""" + r"""The `Trainer` registered at ``location``.""" from kumo.engine.trainer.trainer import Trainer if self.connector is None: @@ -125,8 +103,7 @@ def _location(self) -> Tuple[Optional[str], str]: return database, schema or 'PUBLIC' def check(self, registry: Any = None) -> None: - r"""Raises before any work unless the session can see the - experiment's schema (one metadata query, no warehouse).""" + r"""Raises unless the session can see the experiment's schema.""" database, schema = self._location() if database is None: return @@ -141,9 +118,8 @@ def check(self, registry: Any = None) -> None: f'{type(error).__name__}: {error}') from error if not rows: raise ValueError( - f'Snowflake Experiment Tracking needs schema {database}.' - f'{schema}, which does not exist or this role ' - f'cannot see.') + f'Snowflake schema {database}.{schema} does not exist or this ' + f'role cannot see it.') def start(self, record: RunRecord) -> None: self._reset() @@ -155,11 +131,9 @@ def start(self, record: RunRecord) -> None: elif (not hasattr(tracking, '_run') or tracking._run is not None or getattr(tracking, '_session', None) is not self.session): raise RuntimeError( - f'Snowflake ExperimentTracking is in use in this process, by ' - f'an active run or another Snowpark session, so this ' - f'{record.tags.get("kumo.kind")} run is not tracked in ' - f'Snowflake; end that run, or pass SnowflakeTracking the ' - f'session it uses') + f'Snowflake ExperimentTracking is already in use in this ' + f'process, so this {record.tags.get("kumo.kind")} run is not ' + f'tracked; end that run or pass its session.') self._tracking = tracking database, schema = self._location() experiment = tracking.set_experiment(self.experiment_name, @@ -235,7 +209,7 @@ def _new_run_name(self, tracking: Any, experiment: Any, run_name: Optional[str]) -> Optional[str]: if run_name is None: return None - # start_run matches name as a LIKE pattern. + # start_run matches name with LIKE. existing = [ str(row['name']) for row in tracking._sql_client.show_runs_in_experiment( @@ -305,7 +279,6 @@ def _log_metrics(self, metrics: Mapping[str, float], limit: int, def _run_identifier(name: str) -> str: - r"""``name`` as an unquoted identifier, which artifact URIs need.""" name = re.sub(r'[^A-Za-z0-9_]', '_', name) return name if re.match(r'[A-Za-z_]', name) else f'_{name}' diff --git a/packages/kumo-engine/test/_snowflake_experiment_fake.py b/packages/kumo-engine/test/_snowflake_experiment_fake.py index 27b1cc778..3c914cedf 100644 --- a/packages/kumo-engine/test/_snowflake_experiment_fake.py +++ b/packages/kumo-engine/test/_snowflake_experiment_fake.py @@ -1,8 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""A recording stand-in for `snowflake.ml.experiment.ExperimentTracking` that -keeps its observable rules (singleton, one active run, LIKE-matched resume).""" +"""A recording fake of `snowflake.ml.experiment.ExperimentTracking`.""" import os import re diff --git a/packages/kumo-engine/test/test_mlflow_target.py b/packages/kumo-engine/test/test_mlflow_target.py index d052f0cb6..01f05dad2 100644 --- a/packages/kumo-engine/test/test_mlflow_target.py +++ b/packages/kumo-engine/test/test_mlflow_target.py @@ -1,9 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -r"""How `MlflowTracking` reaches its server: explicit credentials set exactly -the state MLflow reads, only around each call; ambient ones set nothing. The -last tests pin the MLflow and plugin internals this relies on.""" +r"""How `MlflowTracking` reaches its server.""" import os import threading @@ -62,8 +60,6 @@ def test_ambient_identities_set_no_credentials(monkeypatch, tracking_uri): target = resolve(_connector(DatabricksIdentity())) before = dict(os.environ) with target.scope(): - # A profile or notebook identity: Unity Catalog registration - # downloads through the process-wide URI, so it points at Databricks. assert mlflow.get_tracking_uri() == 'databricks' assert {k: v for k, v in os.environ.items() if k.startswith('DATABRICKS_')} == {'DATABRICKS_TOKEN': 'mine'} @@ -99,7 +95,6 @@ def test_a_databricks_service_principal_lends_a_token_for_each_call( with target.scope(): inside = {k: v for k, v in os.environ.items() if k.startswith('DATABRICKS_')} - # Unity Catalog registration downloads through the global URI. assert mlflow.get_tracking_uri() == 'databricks' assert mlflow.get_tracking_uri() == previous assert inside == {'DATABRICKS_HOST': 'https://dbc.example', @@ -194,7 +189,7 @@ def call(): def test_ambient_sagemaker_calls_see_its_uri_and_restore_the_global_one( tracking_uri, monkeypatch): - pytest.importorskip('sagemaker_mlflow') # registers the arn store + pytest.importorskip('sagemaker_mlflow') from kumo.engine.integrations.mlflow import MlflowTracking seen = [] @@ -222,7 +217,7 @@ def test_an_unreachable_server_fails_when_mlops_is_built( with pytest.raises(RuntimeError, match='MLflow at http://127.0.0.1:9 ' 'is not usable with these credentials'): kumo.MLOps(str(tmp_path), tracking=tracking) - with pytest.raises(ValueError, match='pass experiment_name'): + with pytest.raises(ValueError, match='needs experiment_name'): kumo.MLOps(str(tmp_path), tracking=MlflowTracking(target='http://127.0.0.1:9')) @@ -253,9 +248,6 @@ def http_request(creds, endpoint, method, **kwargs): assert len(calls) == 1 -# Pins: MLflow and plugin internals the code above relies on. - - def test_pin_set_tracking_uri_keeps_the_module_value(tracking_uri): utils = pytest.importorskip('mlflow.tracking._tracking_service.utils') assert utils._tracking_uri == tracking_uri @@ -384,7 +376,7 @@ def test_kumo_calls_drop_mlflow_request_header_errors_only(caplog): headers = logging.getLogger('mlflow.tracking.request_header.registry') message = ('Encountered unexpected error during resolving request ' 'headers: %s') - headers.addHandler(caplog.handler) # MLflow's loggers do not propagate + headers.addHandler(caplog.handler) try: with caplog.at_level('WARNING', logger=headers.name): with resolve('sqlite:///x.db').scope(): diff --git a/packages/kumo-engine/test/test_mlflow_tracking.py b/packages/kumo-engine/test/test_mlflow_tracking.py index 526726d0e..54beec57f 100644 --- a/packages/kumo-engine/test/test_mlflow_tracking.py +++ b/packages/kumo-engine/test/test_mlflow_tracking.py @@ -38,7 +38,7 @@ @pytest.fixture def tracking_uri(tmp_path, monkeypatch): - monkeypatch.chdir(tmp_path) # the default artifact root is ./mlruns + monkeypatch.chdir(tmp_path) previous = mlflow.get_tracking_uri() uri = f'sqlite:///{tmp_path / "mlflow.db"}' mlflow.set_tracking_uri(uri) @@ -281,7 +281,7 @@ def refuse(self, name, *args, **kwargs): monkeypatch.setattr(mlflow.MlflowClient, 'create_experiment', refuse) with pytest.raises(RuntimeError, match=r"experiment '/Shared/x/churn'.*" - r"folder does not exist.*parent folder must exist"): + r"folder does not exist.*create its parent folder"): kumo.MLOps(str(tmp_path), tracking=MlflowTracking( '/Shared/x/churn', log_system_metrics=False, target=uri)) diff --git a/packages/kumo-engine/test/test_snowflake_tracking.py b/packages/kumo-engine/test/test_snowflake_tracking.py index 43dc694da..b6de6819e 100644 --- a/packages/kumo-engine/test/test_snowflake_tracking.py +++ b/packages/kumo-engine/test/test_snowflake_tracking.py @@ -130,7 +130,7 @@ def test_an_instance_in_use_is_left_alone_and_fails_strict_mode( plan.optimization.max_epochs = 2 plan.optimization.max_steps_per_epoch = 2 trainer = kumo.Trainer(plan) - in_use = 'Snowflake ExperimentTracking is in use in this process' + in_use = 'Snowflake ExperimentTracking is already in use' with pytest.raises(RuntimeError, match=in_use): trainer.fit(table, snapshot=snapshot, mlops=kumo.MLOps( @@ -162,7 +162,7 @@ def test_an_existing_run_is_never_resumed(fake, taken, caplog): if taken == 'FIT': user.end_run() else: - user._run = None # still running, in another process + user._run = None run_name = 'fit' if taken == 'FIT' else 'my_run' tracking = _tracking(session, run_name=run_name) @@ -242,7 +242,7 @@ def test_a_connector_opens_the_session_and_the_schema_is_checked_first( assert session.statements == ["SHOW SCHEMAS LIKE 'RUNS' IN DATABASE DB"] tracking.session = _SchemaSession([]) - with pytest.raises(ValueError, match='needs schema DB.RUNS'): + with pytest.raises(ValueError, match='schema DB.RUNS does not exist'): kumo.MLOps('/tmp/root', tracking=tracking) with pytest.raises(ValueError, match='or a Snowpark session'): _tracking(None) From 5de9a0bcb24ced4bd3fa772ca1747a3ed000e4ca Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 00:50:03 +0000 Subject: [PATCH 156/278] Trim docstrings, comments and messages in kumo-engine MLOps Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/_data_summaries.py | 2 - packages/kumo-engine/kumo/engine/_deps.py | 7 +-- .../kumo-engine/kumo/engine/_embeddings.py | 7 +-- packages/kumo-engine/kumo/engine/_manifest.py | 19 +++----- packages/kumo-engine/kumo/engine/_mlops.py | 45 +++++-------------- .../kumo-engine/kumo/engine/_predictor.py | 1 - packages/kumo-engine/kumo/engine/_quiet.py | 3 +- packages/kumo-engine/kumo/engine/_scratch.py | 3 +- .../kumo-engine/kumo/engine/_snapshot_io.py | 3 +- packages/kumo-engine/kumo/engine/_tracking.py | 28 +----------- .../kumo/engine/connector/snowflake.py | 3 +- packages/kumo-engine/kumo/engine/graph.py | 6 +-- packages/kumo-engine/kumo/engine/tracking.py | 2 +- .../kumo/engine/trainer/_provenance.py | 3 +- .../kumo/engine/trainer/output_config.py | 3 +- .../kumo/engine/trainer/trainer.py | 35 ++++----------- packages/kumo-engine/test/test_fit_events.py | 2 +- packages/kumo-engine/test/test_mlops.py | 2 +- .../test/test_model_agnostic_mlops.py | 6 +-- .../test/test_tracking_contract.py | 9 +--- 20 files changed, 44 insertions(+), 145 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_data_summaries.py b/packages/kumo-engine/kumo/engine/_data_summaries.py index 8a589f003..9e24d9a64 100644 --- a/packages/kumo-engine/kumo/engine/_data_summaries.py +++ b/packages/kumo-engine/kumo/engine/_data_summaries.py @@ -23,7 +23,6 @@ _TOP_CATEGORIES = 10 _EDGE_STATS = {'FRAC_SRC_NONZERO_DEGREE': 'src_match_frac', 'FRAC_DST_NONZERO_DEGREE': 'dst_match_frac'} -# The Kumo platform's data checks. _BAD_TIME_FRACTION = 0.05 _INVALID_FRACTION = 0.10 _MIN_MATCH_FRACTION = 0.20 @@ -86,7 +85,6 @@ def _profiles(columns: Mapping[str, Mapping[str, Any]] def _relations(edge_stats: Mapping[str, Any], ingest: Mapping[str, Any]) -> Dict[str, Dict[str, float]]: - r"""Each relation's counts and match fractions, by DiskGraph name.""" relations: Dict[str, Dict[str, float]] = { name: {'count': count} for name, count in (ingest.get('edge_counts') or {}).items() diff --git a/packages/kumo-engine/kumo/engine/_deps.py b/packages/kumo-engine/kumo/engine/_deps.py index 65c1e7762..6c11abd07 100644 --- a/packages/kumo-engine/kumo/engine/_deps.py +++ b/packages/kumo-engine/kumo/engine/_deps.py @@ -165,10 +165,8 @@ def resolve(self, asset_kind: str, name: str) -> Path: raise MissingDependencyError( f"The model needs {asset_kind} {name!r} (sha256 " f"{record.sha256}), which is in none of " - f"{tried or 'no configured locations'}. Store it by fitting " - f"with MLOps on (it is kept under /assets/), " - f"predict inside a kumo.mlops block whose artifact root holds " - f"it, or register it where the model looks it up.") + f"{tried or 'no configured locations'}; predict in a kumo.mlops " + f"block whose artifact root holds it.") def candidates(self, record: Dependency) -> List[Tuple[Any, str]]: from kumo.engine.connector import connector_for_root @@ -188,7 +186,6 @@ def candidates(self, record: Dependency) -> List[Tuple[Any, str]]: def _store(asset: Asset, sha256: str, settings: Any) -> Optional[str]: - r"""Writes ``asset`` once under the root; returns where.""" from kumo_core.object_store_sink import _prepared from kumo.engine.connector import file_store diff --git a/packages/kumo-engine/kumo/engine/_embeddings.py b/packages/kumo-engine/kumo/engine/_embeddings.py index de3a77638..e06675647 100644 --- a/packages/kumo-engine/kumo/engine/_embeddings.py +++ b/packages/kumo-engine/kumo/engine/_embeddings.py @@ -15,9 +15,7 @@ def no_encoder_inference_warnings( overridden: Optional[Collection[str]] = None) -> Iterator[None]: # Only fit reports what inference chose: training tables do not use it, - # and a trained model encodes columns as it was trained to. Given the - # columns whose encoders the model plan overrides, fit stays silent on - # those alone. + # and a trained model encodes columns as it was trained to. def drop(record: logging.LogRecord) -> bool: if overridden is None: return record.levelno < logging.WARNING @@ -94,8 +92,7 @@ def _text_encoders(encoders: Any) -> List[Any]: def require_text_encoders(deps: Any, store: EmbeddingModelStore, encoders: Any) -> bool: - r"""Records the word vectors of a fitted model's text encoders in - ``deps``; returns whether it has any.""" + r"""Records the model's text-encoder vectors; returns whether any.""" from kumo_core.embedding_models import GloVeVectors, resolve_encoder from kumo.engine._manifest import Dependency diff --git a/packages/kumo-engine/kumo/engine/_manifest.py b/packages/kumo-engine/kumo/engine/_manifest.py index c6b83ef1b..0e82261cb 100644 --- a/packages/kumo-engine/kumo/engine/_manifest.py +++ b/packages/kumo-engine/kumo/engine/_manifest.py @@ -210,7 +210,6 @@ def _walk(value: Any, key: str, root: Manifest) -> Iterator[Artifact]: yield from _walk(item, key, root) -#: Connector kinds of the artifact root's URI schemes. _SCHEMES = {'s3': 'S3', 'gs': 'GCS', 'file': 'LOCAL_FILE', '': 'LOCAL_FILE'} Details = Mapping[str, Tuple[Mapping[str, Any], Mapping[str, Any]]] @@ -219,8 +218,7 @@ def _walk(value: Any, key: str, root: Manifest) -> Iterator[Artifact]: def artifact_summaries(manifest: Manifest, context: str, settings: Any = None, details: Optional[Details] = None) -> List[Summary]: - r"""Tags locating ``manifest``'s artifacts, and a `Dataset` per stored - dataset.""" + r"""Tags locating ``manifest``'s artifacts, and its `Dataset`s.""" located = [item for item in artifacts(manifest) if item.location] tags: Dict[str, str] = {} for item in located: @@ -348,9 +346,7 @@ class FitTiming(Record): @dataclass(frozen=True, kw_only=True) class TrackingRef(Record): - r"""A run in an experiment tracker: its ``platform`` (``mlflow``, - ``databricks``, ``sagemaker``, ``snowflake``), the platform's own - ``run_id`` and, where the platform has one, the run's ``url``.""" + r"""A run in an experiment tracker.""" platform: str run_id: str @@ -359,9 +355,7 @@ class TrackingRef(Record): @dataclass(frozen=True, kw_only=True) class Dependency(Record): - r"""One thing a model needs to predict: its ``predictor`` entry point, - a ``package`` requirement, or an ``asset`` (content-addressed by - ``sha256``, loaded by ``entry_point``, stored at ``location``).""" + r"""A predictor, package or asset a model needs to predict.""" kind: str name: str @@ -385,9 +379,7 @@ def as_artifact(self, key: str) -> Optional[Artifact]: class ModelManifest(Manifest): """How a model bundle was fitted: from which training table (embedded, with its snapshot) and warm-start model, when, and with which fit - mode. Its ``id`` is the bundle id; ``dependencies`` is what predicting - needs besides the bundle and a snapshot (none listed: the built-in - `Trainer` predictor).""" + mode. Its ``id`` is the bundle id.""" kind: str = 'model' URI_TAG: ClassVar[Optional[str]] = 'kumo.bundle_uri' @@ -402,8 +394,7 @@ class ModelManifest(Manifest): @dataclass(frozen=True, kw_only=True) class AssetManifest(Manifest): - """An asset a model depends on, stored once under - ``/assets///``; its ``id`` is the sha256.""" + """An asset a model depends on; its ``id`` is the sha256.""" kind: str = 'asset' asset_kind: Optional[str] = None diff --git a/packages/kumo-engine/kumo/engine/_mlops.py b/packages/kumo-engine/kumo/engine/_mlops.py index 2f6d3f68d..6d2c37b3c 100644 --- a/packages/kumo-engine/kumo/engine/_mlops.py +++ b/packages/kumo-engine/kumo/engine/_mlops.py @@ -46,32 +46,13 @@ def uri(self, relative: str) -> str: @dataclass(frozen=True) class MLOps: - r"""Where artifacts go and what records them, for :func:`mlops` blocks - and ``mlops=`` arguments. + r"""Where artifacts go and what records them. Args: - artifact_root: an ``s3://``, ``gs://``, ``file://`` or local path - (``/Volumes/...`` on a Databricks cluster), or a ``(connector, - path)`` pair for a store a URI cannot reach, such as a volume - through a ``DatabricksConnector`` or a Snowflake stage. - tracking: the experiment tracker, a - ``kumo.engine.tracking.Tracking``. Its ``check(registry)``, if - any, runs here, so a server or registry that is unreachable or - refuses these credentials fails before any work. - registry: a ``ModelRegistryLocation`` each fit registers its model - at, in the tracker's registry: MLflow's (Unity Catalog, SageMaker - MLflow) for ``MlflowTracking``, the Snowflake Model Registry for - ``SnowflakeTracking``. - store: the artifact kinds stored under the root when produced - (``snapshot``, ``training_table``, ``model``); - ``True`` for all, ``False`` for none. - on_error: what a failure to track a run or register a model does: - ``'warn'`` warns and the call returns; ``'raise'`` raises once - the call has finished and stored its result, with the failure as - the cause. - - In ``load_from_connector`` calls, an ``MLOps`` stands in for the - connector, with paths relative to its artifact root. + artifact_root: a URI, local path or ``(connector, path)`` pair. + registry: where each fit registers its model, through the tracker. + store: artifact kinds to store; ``True`` for all, ``False`` none. + on_error: ``'warn'`` or ``'raise'`` when tracking or registry fails. """ artifact_root: Union[str, Tuple[Any, str]] @@ -147,9 +128,8 @@ def _check_registry(self, tracker: Any) -> None: connector = None if tracker is None else tracker.connector if connector is None: raise ValueError( - 'registry= registers through the tracker: pass ' - 'tracking=MlflowTracking(...) or SnowflakeTracking(connector, ' - '...).') + 'registry= needs tracking=MlflowTracking(...) or ' + 'SnowflakeTracking(connector, ...).') check_model_location(ModelOutputConfig(connector, location=registry)) def snapshot_uri(self, snapshot: Any) -> str: @@ -232,11 +212,7 @@ def _require_writable(root: _Root) -> None: @contextlib.contextmanager def mlops(settings: MLOps) -> Iterator[MLOps]: - r"""Turns MLOps on for the calls in the block (in this thread or asyncio - task): each snapshot, training table and model is stored - under ``settings.artifact_root``. An inner block replaces an outer one; - an ``mlops=`` argument replaces both for its call. Runs are tracked and - models registered as ``settings`` says.""" + r"""Turns MLOps on, as ``settings`` says, for the calls in the block.""" if not isinstance(settings, MLOps): raise TypeError('kumo.mlops() takes a kumo.MLOps') token = _CURRENT.set(settings) @@ -308,7 +284,6 @@ def snapshot_uri(snapshot: Any, settings = resolve(settings) if settings is None: raise ValueError( - 'A snapshot given by object or id is read from the artifact ' - 'root: call this in a kumo.mlops block, or pass the location ' - 'snapshot.export(...) returned.') + 'A snapshot object or id needs a kumo.mlops block; otherwise ' + 'pass the URI snapshot.export(...) returned.') return settings.snapshot_uri(snapshot) diff --git a/packages/kumo-engine/kumo/engine/_predictor.py b/packages/kumo-engine/kumo/engine/_predictor.py index 5a8720f36..c237363f8 100644 --- a/packages/kumo-engine/kumo/engine/_predictor.py +++ b/packages/kumo-engine/kumo/engine/_predictor.py @@ -100,7 +100,6 @@ def predict_frame(trainer: Any, snapshot: Any, rows: pd.DataFrame, def _in_input_order(rows: pd.DataFrame, scored: pd.DataFrame) -> pd.DataFrame: - r"""``scored`` in the order of ``rows``.""" keys = ['ENTITY'] left = pd.DataFrame({'ENTITY': rows['ENTITY1'].to_numpy()}) right = scored[['ENTITY']].copy() diff --git a/packages/kumo-engine/kumo/engine/_quiet.py b/packages/kumo-engine/kumo/engine/_quiet.py index 3a407b1b2..e498a9e2e 100644 --- a/packages/kumo-engine/kumo/engine/_quiet.py +++ b/packages/kumo-engine/kumo/engine/_quiet.py @@ -1,8 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -r"""Imports PyTorch Geometric without the FutureWarning torch 2.15+ raises -for its import-time ``torch.jit.script`` calls, which users cannot act on.""" +r"""Imports PyTorch Geometric without its unactionable FutureWarning.""" import warnings diff --git a/packages/kumo-engine/kumo/engine/_scratch.py b/packages/kumo-engine/kumo/engine/_scratch.py index 5e62df71c..008165c94 100644 --- a/packages/kumo-engine/kumo/engine/_scratch.py +++ b/packages/kumo-engine/kumo/engine/_scratch.py @@ -76,8 +76,7 @@ def get_scratch_dir() -> Path: @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.""" + r"""Nests this context's new scratch directories under ``path``.""" token = _scoped_root.set(path) try: yield diff --git a/packages/kumo-engine/kumo/engine/_snapshot_io.py b/packages/kumo-engine/kumo/engine/_snapshot_io.py index 55f6b6513..f4450234c 100644 --- a/packages/kumo-engine/kumo/engine/_snapshot_io.py +++ b/packages/kumo-engine/kumo/engine/_snapshot_io.py @@ -192,8 +192,7 @@ def export_location(snapshot: Any, path: str) -> str: def store(snapshot: Any, settings: Any) -> None: - r"""Export ``snapshot`` under ``/snapshots/`` unless a complete - export of it is there.""" + r"""Exports ``snapshot`` under the root unless it is there.""" from kumo.engine._manifest import SnapshotManifest from kumo.engine._mlops import snapshot_path diff --git a/packages/kumo-engine/kumo/engine/_tracking.py b/packages/kumo-engine/kumo/engine/_tracking.py index af341fbbf..850d722ee 100644 --- a/packages/kumo-engine/kumo/engine/_tracking.py +++ b/packages/kumo-engine/kumo/engine/_tracking.py @@ -87,7 +87,6 @@ class RunSpec: def fit_inputs(manifest: Any) -> Dict[str, Dict[str, Optional[str]]]: - r"""Source tables of a training table or model manifest.""" try: table = manifest if getattr(manifest, 'kind', None) == 'model': @@ -116,8 +115,7 @@ def _parents(manifest: Any) -> Any: def logged_name(name: str) -> str: - r"""``val_auroc`` as ``val/auroc``, ``best_val_auroc`` as - ``val/best_auroc``.""" + r"""``val_auroc`` as ``val/auroc``.""" best = name.startswith('best_') split, _, metric = name.removeprefix('best_').partition('_') if split not in ('train', 'val', 'test') or not metric: @@ -131,7 +129,6 @@ def logged_names(metrics: Mapping[str, Any]) -> Dict[str, Any]: def by_step(steps: Mapping[str, Sequence[Tuple[int, float]]] ) -> List[Tuple[int, Dict[str, float]]]: - r"""Each step's metrics, in step order.""" grouped: Dict[int, Dict[str, float]] = {} for name, history in steps.items(): for step, value in history: @@ -174,7 +171,6 @@ def _user_tags(tags: Mapping[str, Any], def trial_spec(events: Sequence[Event]) -> RunSpec: - r"""One trial attempt of a fit, from its events so far.""" started = next(event for event in events if event.name == 'trial_started') trial, attempt = started.trial, int(started.params['attempt']) metrics: Dict[str, float] = {} @@ -203,25 +199,7 @@ def trial_spec(events: Sequence[Event]) -> RunSpec: class Tracking(Protocol): - r"""Records one `Trainer.fit` in an experiment tracker. - - Calls come from one background thread. `start` gets the fit run's - record without metrics, `on_trial` a trial attempt's record so far with - 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; - ``log_metrics(metrics, step)``, series of the fit run as they come (a - single trial's epochs, or a search's best validation metric after each - trial); ``tags`` (extra fit-run tags); ``log_model(bundle_dir, manifest, - 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``; ``connector``, a connector whose model - registry ``registry`` registers in when there is no ``log_model``; and - ``platform``, ``run_id`` and ``run_url``, the tracker's own run once - `start` returned, recorded as the model manifest's ``tracking``. - """ + r"""Records one `Trainer.fit` in an experiment tracker.""" def start(self, record: RunRecord) -> None: ... @@ -313,7 +291,6 @@ def __repr__(self) -> str: def as_tracking(tracking: Any) -> Optional[TrackingBase]: - r"""``tracking``, wrapped unless it subclasses `TrackingBase`.""" if tracking is None or isinstance(tracking, TrackingBase): return tracking return _Foreign(tracking) @@ -442,7 +419,6 @@ def log_model( snapshot: Any, registry: Optional['ModelRegistryLocation'] = None, ) -> Tuple[Optional['PublishedModel'], Optional[BaseException]]: - r"""The tracker's ``log_model``, run on the tracking thread.""" if self._worker is None or (self._tracker is not None and self._tracker.log_model is None): return None, None diff --git a/packages/kumo-engine/kumo/engine/connector/snowflake.py b/packages/kumo-engine/kumo/engine/connector/snowflake.py index bdfc97387..48c1b7526 100644 --- a/packages/kumo-engine/kumo/engine/connector/snowflake.py +++ b/packages/kumo-engine/kumo/engine/connector/snowflake.py @@ -101,8 +101,7 @@ def sink_backend(self) -> ArtifactSink: return SnowflakeSinkBackend(identity=self.identity) def session(self) -> Any: - r"""The Snowpark session this connector reads through: its identity, - pointed at its warehouse, database and schema, kept alive.""" + r"""The Snowpark session this connector reads through.""" with self.source_backend._session(self) as session: return session diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index 2e499a7d5..3a5e7aef5 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -272,8 +272,6 @@ def report_other_graph(message: str, before: Optional[Mapping[str, Any]], class _SnapshotEvents(EventEmitter): - r"""Snapshot progress: each event goes to the callbacks and the log, - with the seconds since the snapshot started.""" def __init__(self, callbacks: Sequence[Callable[[Event], None]]): super().__init__(callbacks) @@ -713,9 +711,7 @@ def snapshot(self, *, record_sources: bool = True, record_sources: if False, the snapshot manifest records each source as connector and table name only (no bucket, path or account) and Custom SQL as a hash only (no SQL text). - callbacks: called with each progress event, - `snapshot_started`, `table_staged` (per table) and - `snapshot_ended`; each is also logged at INFO. + callbacks: called with each progress `Event`. """ self.validate() for table in self._named_tables.values(): diff --git a/packages/kumo-engine/kumo/engine/tracking.py b/packages/kumo-engine/kumo/engine/tracking.py index af95f228f..4a7dc3341 100644 --- a/packages/kumo-engine/kumo/engine/tracking.py +++ b/packages/kumo-engine/kumo/engine/tracking.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -r"""API for writing an experiment-tracking adapter for ``kumo.MLOps(tracking=...)``.""" +r"""API for writing an experiment tracker for ``kumo.MLOps``.""" from kumo.engine._events import Event from kumo.engine._tracking import RunRecord, Tracking diff --git a/packages/kumo-engine/kumo/engine/trainer/_provenance.py b/packages/kumo-engine/kumo/engine/trainer/_provenance.py index 0df90fcf9..18c234e63 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_provenance.py +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -95,8 +95,7 @@ def read_provenance(bundle_dir: Path) -> Optional[ModelManifest]: def record_derivation(bundle_dir: Path, parent_dir: Path, parent_id: Optional[str]) -> None: - r"""Writes the manifest of ``bundle_dir``, a copy of the bundle at - ``parent_dir`` with other scoring defaults.""" + r"""Writes the manifest of a bundle derived from ``parent_dir``.""" from kumo_core.artifacts import ( PROVENANCE_FILE, add_bundle_file, diff --git a/packages/kumo-engine/kumo/engine/trainer/output_config.py b/packages/kumo-engine/kumo/engine/trainer/output_config.py index bf1f14c00..f4e0c9b63 100644 --- a/packages/kumo-engine/kumo/engine/trainer/output_config.py +++ b/packages/kumo-engine/kumo/engine/trainer/output_config.py @@ -118,8 +118,7 @@ def __post_init__(self) -> None: def check_model_location(config: ModelOutputConfig) -> None: - """Raises ``ValueError`` unless ``config.location`` holds an alias and - a model name its registry accepts.""" + """Raises ``ValueError`` unless ``config.location`` is registrable.""" location = config.location if location is None: return diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index c6ab6f5fe..dbbe49c6b 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -174,8 +174,7 @@ def snapshot_id(self) -> Optional[str]: @property def tracking(self) -> Optional[TrackingRef]: - """The experiment tracker's run of the fit: its ``platform``, - ``run_id`` and ``url``; None when the fit was not tracked.""" + """The experiment tracker's run of the fit, if tracked.""" manifest = self.manifest return None if manifest is None else manifest.tracking @@ -184,9 +183,7 @@ def predict(self, *args: Any, **kwargs: Any) -> 'PredictionResult': return self._model('predict').predict(*args, **kwargs) def export(self, config: Any) -> PublishedModel: - """Publish this model's bundle to a ``ModelOutputConfig``, or - register it at a `kumo.MLOps` ``registry`` through its tracker; a - registry that already holds the bundle returns that version.""" + """Publishes the bundle to a ``ModelOutputConfig`` or ``MLOps``.""" from kumo.engine._mlops import MLOps if isinstance(config, MLOps): @@ -204,10 +201,7 @@ def export(self, config: Any) -> PublishedModel: def with_scoring_defaults(self, defaults: ScoringDefaults) -> 'TrainingResult': - """A copy of this model whose ``predict`` applies ``defaults`` when a - call omits them; a new bundle, whose manifest names this one in - ``parents.derived_from``. To register a model with its defaults, - pass ``fit(scoring_defaults=...)`` instead.""" + """A copy of this model whose ``predict`` applies ``defaults``.""" from kumo_core.artifacts import bundle_task_type if self._trainer is None: @@ -387,7 +381,6 @@ def _store_and_register( warn: bool = True, stored: Optional[BaseException] = None, ) -> Tuple[Optional[PublishedModel], Optional[BaseException]]: - r"""Stores, logs and registers a bundle; returns it and any error.""" if stored is None and settings is not None and settings.wants('model'): try: _store_model(bundle_dir, manifest, settings) @@ -441,8 +434,7 @@ def _fit_summaries(snapshot: GraphSnapshot, train_table: TrainingTable, def _store_model(bundle_dir: Path, manifest: Optional[ModelManifest], settings: Any) -> None: - r"""Copy the bundle to ``/models//`` unless a complete - copy is there.""" + r"""Copies the bundle under the root unless it is there.""" from kumo_core.artifacts import PROVENANCE_FILE, add_bundle_file from kumo_core.object_store_sink import holds_bundle, upload_bundle from pyarrow.fs import FileSelector @@ -514,7 +506,6 @@ def _require_publisher(export_config: ModelOutputConfig) -> None: def _with_metrics(model: Optional[PublishedModel], metrics: Mapping[str, float]) -> Optional[PublishedModel]: - r"""``model`` with ``metrics``, where its registry returned none.""" if model is None or model.metrics: return model return replace(model, metrics=dict(metrics)) @@ -596,7 +587,6 @@ def _stage_model_bundle(artifacts: Mapping[str, str], dest: Path) -> Path: def _without_unactionable_warnings(method: _F) -> _F: - r"""Runs ``method`` without third-party notes users cannot act on.""" @functools.wraps(method) def wrapper(*args: Any, **kwargs: Any) -> Any: tips = logging.getLogger('pytorch_lightning.utilities.rank_zero') @@ -683,8 +673,7 @@ def _fitted(self, method: str) -> TrainingResult: return self._training_result def export(self, config: Any) -> PublishedModel: - """Publish the fitted or loaded model bundle, as - :meth:`TrainingResult.export`.""" + """Publishes the model, as :meth:`TrainingResult.export`.""" return self._fitted('export').export(config) @_without_unactionable_warnings @@ -718,9 +707,7 @@ def fit( generate_training_table(snapshot)`. snapshot: graph data and features to train against. export_config: if given, publishes the trained model bundle. - scoring_defaults: prediction arguments the bundle applies when - ``predict`` omits them, such as a classification threshold; - stored, registered and served with the model. + scoring_defaults: arguments ``predict`` applies when omitted. warm_start: a fitted trainer or one restored with ``load`` or ``load_from_connector``. Reuses its selected weights, optimizer state, fitted encoders, architecture and neighbor sampling; @@ -730,11 +717,8 @@ def fit( architecture; they support prediction but are unreliable warm-start sources. mlops: a `kumo.MLOps`; with MLOps on (this or a `kumo.mlops` - block), the bundle is also stored under the artifact root and - the fit is recorded by its ``tracking``. - callbacks: called with each progress - `kumo.engine.tracking.Event`; a callback that raises is logged - and not called again. + block), the bundle is also stored and the fit tracked. + callbacks: called with each progress `Event`. Each call returns a new `TrainingResult`; `Trainer.result` holds the latest. @@ -1422,8 +1406,7 @@ def load_from_connector( from trusted publishers; checkpoint deserialization happens at predict. ``accelerator`` selects the device for this instance (default: ``auto``). A `kumo.MLOps` ``model_registry`` reads ``path`` under its artifact - root (``models/``) and ``location`` from the registry its - tracker registers in. + root (``models/``) and ``location`` in its tracker's registry. """ from kumo.engine._mlops import MLOps, on_root diff --git a/packages/kumo-engine/test/test_fit_events.py b/packages/kumo-engine/test/test_fit_events.py index 77ecb676f..8f43b35f0 100644 --- a/packages/kumo-engine/test/test_fit_events.py +++ b/packages/kumo-engine/test/test_fit_events.py @@ -229,7 +229,7 @@ def fit(): rank_zero_info('GPU available: False') lightning = logging.getLogger('pytorch_lightning') - lightning.addHandler(caplog.handler) # it does not propagate + lightning.addHandler(caplog.handler) try: with caplog.at_level('INFO', logger='pytorch_lightning'): fit() diff --git a/packages/kumo-engine/test/test_mlops.py b/packages/kumo-engine/test/test_mlops.py index e258f3eb8..4024e0483 100644 --- a/packages/kumo-engine/test/test_mlops.py +++ b/packages/kumo-engine/test/test_mlops.py @@ -53,7 +53,7 @@ def check_model_location(connector, location): kumo.MLOps(str(tmp_path), tracking=tracker, registry=kumo.ModelRegistryLocation('USERS_CHURN', alias='LATEST')) - with pytest.raises(ValueError, match='registers through the tracker'): + with pytest.raises(ValueError, match='registry= needs tracking'): kumo.MLOps(str(tmp_path), registry=good) with pytest.raises(TypeError, match='ModelRegistryLocation'): kumo.MLOps(str(tmp_path), tracking=tracker, diff --git a/packages/kumo-engine/test/test_model_agnostic_mlops.py b/packages/kumo-engine/test/test_model_agnostic_mlops.py index 7cd604a04..ab8d70d57 100644 --- a/packages/kumo-engine/test/test_model_agnostic_mlops.py +++ b/packages/kumo-engine/test/test_model_agnostic_mlops.py @@ -1,8 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""The MLOps seams take any bundle, run and predictor, not only the -Trainer's.""" +"""The MLOps seams take any bundle, run and predictor.""" import io import json @@ -288,8 +287,7 @@ def test_a_model_serves_from_its_bundle_snapshot_and_assets_alone( class _PrefixOnlyS3: - r"""An S3 filesystem whose credentials cover only the root's prefix: - HeadBucket and parent marker objects are refused.""" + r"""S3 credentials that cover only the root's prefix.""" type_name = 's3' def __init__(self): diff --git a/packages/kumo-engine/test/test_tracking_contract.py b/packages/kumo-engine/test/test_tracking_contract.py index b48ef5fd4..2acda5a74 100644 --- a/packages/kumo-engine/test/test_tracking_contract.py +++ b/packages/kumo-engine/test/test_tracking_contract.py @@ -64,7 +64,6 @@ def _epochs(trial, *values): ] -# Trial 2's first attempt fails; epoch events carry no attempt, as in core. EVENTS = [ _event('trial_started', 1, attempt=1, **{'optimization.base_lr': 0.01}), *_epochs(1, (0.5, 0.4), (0.3, 0.35)), @@ -81,8 +80,6 @@ def _epochs(trial, *values): _event('fit_ended', metrics={'val_mae': 0.3, 'test_mae': 0.32, 'seconds': 9.0}, status='succeeded'), ] -# The fit run's own series: the best val_mae after each trial that -# succeeded and, once the fit succeeds, the selected trial 2's epochs. SEARCH = {'search/best_val_mae': ((1, 0.35), (2, 0.3))} SELECTED = {'train_loss': ((1, 0.6), (2, 0.4)), 'val_mae': ((1, 0.5), (2, 0.3))} BUNDLE = 'b' * 64 @@ -129,9 +126,7 @@ def join_tracking_threads(): def replay(tracker=None, events=EVENTS, result=RESULT, error=None, context=CONTEXT, series=None, summaries=()) -> Dict[str, RunRecord]: - r"""Feeds the recorded fit to `tracker`; returns the records expected - of it, with the fit run's ``series`` (default: the search's) and - ``summaries``.""" + r"""Feeds the recorded fit to `tracker`; returns the expected records.""" recorder = _Recorder(tracker) session = TrackingSession(recorder) session.start(lambda: fit_spec(context)) @@ -303,7 +298,7 @@ def bucket(key): for key, value in run['params'].items(): runs, name = bucket(key) - # Snowflake has no tags: they are params of the same name. + # Snowflake logs tags as params. if name.startswith('kumo.') or name in USER_TAGS: runs['tags'][name] = value else: From 6a8684b2c28de3a2cb3d6871d82949d9c918414f Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 00:50:03 +0000 Subject: [PATCH 157/278] Shorten the smoke check docstring and the mlflow test pin comment Co-Authored-By: Claude Opus 5.5 --- examples/mlops/smoke.py | 17 +---------------- setup.py | 3 +-- 2 files changed, 2 insertions(+), 18 deletions(-) diff --git a/examples/mlops/smoke.py b/examples/mlops/smoke.py index 0d4b1cfcf..3cd5c92ff 100644 --- a/examples/mlops/smoke.py +++ b/examples/mlops/smoke.py @@ -1,22 +1,7 @@ # 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:// --client-id \ - --catalog main --schema ml --warehouse --experiment /Users//kumo-smoke - python smoke.py sagemaker --arn --role-arn - python smoke.py snowflake --database ML --schema RUNS --warehouse \ - --account --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. -""" +"""Manual live check per platform; Databricks reads DATABRICKS_CLIENT_SECRET.""" import argparse import os diff --git a/setup.py b/setup.py index 8b1f65e0e..f80ad52d9 100644 --- a/setup.py +++ b/setup.py @@ -293,8 +293,7 @@ def parse_project_version(project_file_path: Path) -> str: "respx==0.22.0", "pytorch-frame==0.3.0", "pytorch-lightning==2.6.6", - # A superset of kumo-engine's `mlflow` extra (mlflow-skinny): the tests use - # MLflow's local SQLite store. Its wheel is tested with --no-deps. + # Full mlflow, for its local SQLite store. "mlflow>=3.13,<4", ] From 01d52b32328803435c670f7ac364a3fc41c3ff3b Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 01:03:14 +0000 Subject: [PATCH 158/278] Cut the Databricks sink's class docstring to a summary and its argument Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_databricks/sink.py | 42 +------------------ 1 file changed, 2 insertions(+), 40 deletions(-) diff --git a/packages/kumo-databricks/src/kumo_databricks/sink.py b/packages/kumo-databricks/src/kumo_databricks/sink.py index 80f32a6a6..e29027771 100644 --- a/packages/kumo-databricks/src/kumo_databricks/sink.py +++ b/packages/kumo-databricks/src/kumo_databricks/sink.py @@ -143,48 +143,10 @@ def write_files(self, connector, table_name, files, class DatabricksSinkBackend(FileStoreRouting, _DatabricksTables): - r"""Writes a produced artifact to a Databricks Delta table. - - Registered by default beside - :class:`~kumo_databricks.backend.DatabricksSourceBackend`, so a - run that read its graph from a catalog exports its predictions back to one - with no further setup. Installed as - ``pip install 'kumoai-studio[databricks]'``. - - A destination names its schema and its table in one string, joined by - ``__kumo__``, because that is the only field a Databricks output has -- - :func:`kumoai_studio.databricks.to_db_table_name` builds it from the - ``(schema, table)`` tuple an :class:`~kumoai_studio.artifact_export.OutputConfig` - takes. The schema must be the connector's own, which - :func:`kumo_databricks.output.within_schema` checks and the - hosted connector checks through as well. - - Every write stages through the connector's managed volume, - ``.._kumo_staging``, created on first use. That is not a - local choice: ``COPY INTO`` is how a Delta table is loaded, and hosted - states the alternatives -- S3, ADLS, an arbitrary volume, inline cloud - credentials -- are unsupported on purpose. The volume is emptied of what - this export put there once the table has landed. + """Writes artifacts to Delta tables, staged through the schema's volume. Args: - identity: A fallback identity, used only for connectors that name none - themselves. As with the reader, leaving it unset is the common - case, since a bound - :class:`~kumoai_studio.connector.DatabricksConnector` records the - identity it was created with. That identity needs ``CREATE TABLE`` - on the connector's schema and ``READ``/``WRITE VOLUME`` on its - staging volume, which a read-only role used for ingest will not - have. - - Example: - >>> import kumoai_studio as kumo # doctest: +SKIP - >>> from kumo_databricks import ( # doctest: +SKIP - ... DatabricksSinkBackend) - >>> kumo.set_backend(kumo.LocalBackend( # doctest: +SKIP - ... storage=paths, - ... sinks={kumo.DataSourceType.DATABRICKS: - ... DatabricksSinkBackend(identity='prod-dbx')}, - ... )) + identity: Fallback identity for connectors that name none. """ source_type = DataSourceType.DATABRICKS From c7588bfad534cc2cd75fdaf1f889d41eb48f533e Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 01:06:36 +0000 Subject: [PATCH 159/278] Trim engine docstrings, comments and messages Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_manifest.py | 28 ++++-------- packages/kumo-engine/kumo/engine/_mlops.py | 43 +++++-------------- .../kumo-engine/kumo/engine/_snapshot_io.py | 25 +++-------- packages/kumo-engine/kumo/engine/graph.py | 36 +++++----------- .../kumo/engine/pquery/predictive_query.py | 10 ----- .../kumo/engine/pquery/training_table.py | 11 ++--- .../kumo-engine/kumo/engine/trainer/_spec.py | 2 - .../kumo/engine/trainer/output_config.py | 3 +- .../kumo/engine/trainer/trainer.py | 21 +++------ 9 files changed, 48 insertions(+), 131 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_manifest.py b/packages/kumo-engine/kumo/engine/_manifest.py index 1e10f6e31..49451ca02 100644 --- a/packages/kumo-engine/kumo/engine/_manifest.py +++ b/packages/kumo-engine/kumo/engine/_manifest.py @@ -42,9 +42,7 @@ def omit_none() -> Any: @dataclass(frozen=True) class Artifact: - r"""An artifact a manifest names, stored at ``location``, which tracked - runs tag as ``tag``. A dataset is data a run reads, or writes when its - ``context`` is ``output``.""" + """An artifact a manifest names.""" kind: str id: str @@ -59,13 +57,12 @@ class Artifact: @dataclass(frozen=True, kw_only=True) class Record: - """A JSON object; keys this version does not know are kept in ``extra``.""" + """A JSON object that keeps unknown keys in ``extra``.""" extra: Mapping[str, Any] = field(default_factory=dict, repr=False, compare=False) def as_artifact(self, key: str) -> Optional[Artifact]: - r"""The artifact this record names as ``key``, if any.""" return None def to_dict(self) -> Dict[str, Any]: @@ -153,13 +150,12 @@ def as_artifact(self, key: str) -> Artifact: version_kind=self.version_kind) -#: The artifact kinds, as manifests and the artifact root name them. KINDS = ('snapshot', 'training_table', 'model') @dataclass(frozen=True, kw_only=True) class Manifest(Record): - """An artifact's record; a stub sets only kind, id and reconstructed.""" + """An artifact's record.""" kind: str = '' schema_version: int = SCHEMA_VERSION @@ -169,11 +165,8 @@ class Manifest(Record): reconstructed: Optional[bool] = omit_none() location: Optional[str] = omit_none() - #: Fields left out where this manifest is embedded in a child's. EMBED_OMITS = () - #: Whether this kind is data a run reads or writes. IS_DATASET: ClassVar[bool] = False - #: The tag of its location in tracked runs. URI_TAG: ClassVar[Optional[str]] = None def as_artifact(self, key: str) -> Artifact: @@ -183,7 +176,7 @@ def as_artifact(self, key: str) -> Artifact: def artifacts(manifest: Manifest) -> List[Artifact]: - r"""``manifest``'s artifact and those it names, depth first, once each.""" + """``manifest`` and the artifacts it names, once each.""" found: Dict[Any, Artifact] = {} for artifact in _walk(manifest, manifest.kind, manifest): found.setdefault((artifact.kind, artifact.id), artifact) @@ -217,7 +210,6 @@ def stamp() -> Dict[str, Any]: environment=environment()) -#: Envelope fields for a manifest rebuilt after the fact. RECONSTRUCTED: Dict[str, Any] = dict(created_at=None, environment=None, reconstructed=True) @@ -234,12 +226,12 @@ def _kinds() -> Dict[str, Type['Manifest']]: def stub(kind: str, id: str) -> 'Manifest': - """The manifest of a ``kind`` artifact known only by its ``id``.""" + """A manifest knowing only ``kind`` and ``id``.""" return _kinds()[kind](id=id, reconstructed=True) def embedded(manifest: M) -> M: - """``manifest`` without the fields a child's copy leaves out.""" + """``manifest`` minus the fields a child's copy omits.""" return replace(manifest, **{ name: None for name in manifest.EMBED_OMITS @@ -263,7 +255,7 @@ class SnapshotParents(Record): @dataclass(frozen=True, kw_only=True) class SnapshotManifest(Manifest): - """The source tables and engine build a snapshot was made from.""" + """The sources and engine build of a snapshot.""" kind: str = 'snapshot' IS_DATASET: ClassVar[bool] = True @@ -279,7 +271,7 @@ class TrainingTableParents(Record): @dataclass(frozen=True, kw_only=True) class TrainingTableManifest(Manifest): - """The snapshot, query and rows a training table was made from.""" + """The snapshot, query and rows of a training table.""" kind: str = 'training_table' IS_DATASET: ClassVar[bool] = True @@ -305,9 +297,7 @@ class FitTiming(Record): @dataclass(frozen=True, kw_only=True) class ModelManifest(Manifest): - """How a model bundle was fitted: from which training table (embedded, - with its snapshot) and warm-start model, when, and with which fit - mode. Its ``id`` is the bundle id.""" + """How a model bundle was fitted; ``id`` is the bundle id.""" kind: str = 'model' URI_TAG: ClassVar[Optional[str]] = 'kumo.bundle_uri' diff --git a/packages/kumo-engine/kumo/engine/_mlops.py b/packages/kumo-engine/kumo/engine/_mlops.py index fc4b31d97..6abfc020f 100644 --- a/packages/kumo-engine/kumo/engine/_mlops.py +++ b/packages/kumo-engine/kumo/engine/_mlops.py @@ -21,31 +21,20 @@ class _Root: _uri: str def location(self, relative: str) -> str: - r"""``relative`` as a path on :attr:`connector`.""" base = self.path.rstrip('/') return f'{base}/{relative}' if base else relative def uri(self, relative: str) -> str: - r"""``relative`` as a full URI or path, for messages and records.""" return f"{self._uri.rstrip('/')}/{relative}" @dataclass(frozen=True) class MLOps: - r"""Where artifacts go, for :func:`mlops` blocks and ``mlops=`` - arguments. + r"""Where artifacts are stored. Args: - artifact_root: an ``s3://``, ``gs://``, ``file://`` or local path - (``/Volumes/...`` on a Databricks cluster), or a ``(connector, - path)`` pair for a store a URI cannot reach, such as a volume - through a ``DatabricksConnector`` or a Snowflake stage. - store: the artifact kinds stored under the root when produced - (``snapshot``, ``training_table``, ``model``); - ``True`` for all, ``False`` for none. - - In ``load_from_connector`` calls, an ``MLOps`` stands in for the - connector, with paths relative to its artifact root. + artifact_root: a URI or path, or a ``(connector, path)`` pair. + store: artifact kinds to store; ``True`` for all, ``False`` for none. """ artifact_root: Union[str, Tuple[Any, str]] @@ -88,7 +77,6 @@ def __post_init__(self) -> None: object.__setattr__(self, 'kinds', kinds) def wants(self, kind: str) -> bool: - r"""Whether artifacts of ``kind`` are stored under the root.""" return kind in self.kinds @@ -107,9 +95,7 @@ def _require_within_root(connector: Any, path: str) -> None: name = type(connector).__name__ raise ValueError( f'artifact_root {path!r} is outside this {name}\'s {within} ' - f'({error}), so nothing could be stored there. Give the root ' - f'a connector that contains it: ({name}(root_dir={path!r}, ' - f'identity=...), {path!r}).') from error + f'({error}); give it a connector that contains it.') from error def _require_writable(root: _Root) -> None: @@ -132,7 +118,7 @@ def _require_writable(root: _Root) -> None: return try: filesystem.open_output_stream(probe).close() - except FileNotFoundError: # a store with directories, before first use + except FileNotFoundError: filesystem.create_dir(posixpath.dirname(probe), recursive=True) filesystem.open_output_stream(probe).close() delete = getattr(root.connector.sink_backend, 'delete_file', None) @@ -142,10 +128,9 @@ def _require_writable(root: _Root) -> None: delete(root.connector, probe) except Exception as error: raise PermissionError( - f'kumo.MLOps cannot write under artifact_root ' - f'{root.uri("")!r}: {type(error).__name__}: {error}. Grant this ' - f'identity write and delete rights on it, or pick another ' - f'root.') from error + f'kumo.MLOps cannot write under artifact_root {root.uri("")!r} ' + f'({type(error).__name__}: {error}); grant write and delete ' + f'rights or pick another root.') from error _CURRENT: contextvars.ContextVar[Optional[MLOps]] = contextvars.ContextVar( @@ -154,10 +139,7 @@ def _require_writable(root: _Root) -> None: @contextlib.contextmanager def mlops(settings: MLOps) -> Iterator[MLOps]: - r"""Turns MLOps on for the calls in the block (in this thread or asyncio - task): each snapshot, training table and model is stored - under ``settings.artifact_root``. An inner block replaces an outer one; - an ``mlops=`` argument replaces both for its call.""" + r"""Store artifacts of calls in this block under ``settings``.""" if not isinstance(settings, MLOps): raise TypeError('kumo.mlops() takes a kumo.MLOps') token = _CURRENT.set(settings) @@ -168,8 +150,7 @@ def mlops(settings: MLOps) -> Iterator[MLOps]: def resolve(settings: Optional[MLOps] = None) -> Optional[MLOps]: - r"""The settings a call runs with: its ``mlops=`` argument, else the - enclosing block's, else None (off).""" + r"""The ``mlops=`` argument, else the enclosing block's.""" return settings if settings is not None else _CURRENT.get() @@ -179,14 +160,12 @@ def warn_not_stored(what: str, error: Exception, stacklevel: int) -> None: def on_root(source: Any, path: str) -> Tuple[Any, str]: - r"""The connector and path a ``load_from_connector(source, path=...)`` - reads: an `MLOps` ``source`` reads ``path`` under its artifact root.""" + r"""``path`` under an `MLOps` root, else ``source`` as given.""" if isinstance(source, MLOps): return source.root.connector, source.root.location(path.lstrip('/')) return source, path def snapshot_path(snapshot: Any) -> str: - r"""Where a snapshot, or a snapshot id, is stored under a root.""" snapshot_id = snapshot if isinstance(snapshot, str) else snapshot.id return 'snapshots/' + snapshot_id.replace('@', '/') diff --git a/packages/kumo-engine/kumo/engine/_snapshot_io.py b/packages/kumo-engine/kumo/engine/_snapshot_io.py index 36b6e605e..ddeb1faaf 100644 --- a/packages/kumo-engine/kumo/engine/_snapshot_io.py +++ b/packages/kumo-engine/kumo/engine/_snapshot_io.py @@ -40,9 +40,8 @@ def _location(connector: Any, path: str) -> Tuple[FileSystem, str]: store = file_store(connector, path) if store is None: raise NotImplementedError( - f'{type(connector).__name__} cannot hold a snapshot export: export ' - 'needs a file destination (artifact root), such as an S3Connector, ' - 'GCSConnector or a FileConnector on a volume.') + f'{type(connector).__name__} cannot hold a snapshot export; use ' + 'an S3, GCS or file connector.') return store @@ -192,8 +191,7 @@ def export_location(snapshot: Any, path: str) -> str: def store(snapshot: Any, settings: Any) -> None: - r"""Export ``snapshot`` under ``/snapshots/`` unless a complete - export of it is there; failures warn.""" + r"""Export ``snapshot`` under the root unless already there.""" from kumo.engine._mlops import warn_not_stored try: @@ -230,7 +228,6 @@ def export(snapshot: Any, connector: Any, path: str) -> str: entry['rel_path']: int(entry['size_bytes']) for entry in json.loads(index_manifest)['files'] } - # One generation per export; the marker, written last, names it. generation = uuid.uuid4().hex filesystem, root = _location(connector, location) base = f'{root}/{generation}' @@ -278,7 +275,6 @@ def upload(client: FileSystem, name: str, stop: threading.Event) -> None: def _holds_copy(connector: Any, location: str, snapshot_id: str, index: Path, snapshot_manifest: bytes) -> bool: - r"""Whether ``location`` holds this index's bytes and manifest.""" from kumo_core.object_store_sink import file_sizes try: @@ -303,7 +299,6 @@ def _holds_copy(connector: Any, location: str, snapshot_id: str, def _commit(filesystem: FileSystem, root: str, generation: str, marker: bytes) -> None: - r"""Writes the marker atomically: a local rename, else one object put.""" if filesystem.type_name == 'local': staged = f'{root}/{generation}/{_MANIFEST}' _write(filesystem, staged, marker) @@ -315,7 +310,6 @@ def _commit(filesystem: FileSystem, root: str, generation: str, def _clean(filesystem: FileSystem, root: str) -> None: - r"""Removes superseded and abandoned generations, best effort.""" from pyarrow.fs import FileSelector, FileType try: @@ -355,16 +349,14 @@ def read(connector: Any, path: str) -> Tuple[str, dict]: raise _corrupt(path, f'unreadable {_MANIFEST}: {error}') from error if manifest is None: raise FileNotFoundError( - f'No complete snapshot export at {root}: it has no {_MANIFEST}. ' - f'Check the path, or export the snapshot there with ' - f'GraphSnapshot.export.') + f'No snapshot export at {root} (no {_MANIFEST}); check the path.') if not isinstance(manifest, dict): raise _corrupt(path, 'manifest is not a JSON object') if manifest.get('format') != _FORMAT: raise ValueError( f"Snapshot export at {path} has format " - f"{manifest.get('format')!r}; this kumo-engine reads format " - f"{_FORMAT}. Load it with the kumo-engine that exported it.") + f"{manifest.get('format')!r}, not {_FORMAT}; load it with the " + f"kumo-engine that exported it.") _check_key(manifest.get('graph_key'), manifest.get('data_version'), manifest.get('snapshot_id')) span, crcs = manifest.get('range_bytes'), manifest.get('crc32') @@ -379,14 +371,12 @@ def read(connector: Any, path: str) -> Tuple[str, dict]: def export_identity(manifest: dict) -> str: - r"""A hash of an export's range checksums and index manifest.""" return hashlib.sha256(json.dumps( [manifest['range_bytes'], manifest['crc32'], manifest['index_manifest']], sort_keys=True).encode()).hexdigest() def holds_export(index: Path, manifest: dict) -> bool: - r"""Whether the local ``index`` has the export's bytes.""" from diskgraph import snapshot_range_checksums identity = export_identity(manifest) @@ -529,5 +519,4 @@ def _check_readable(index: Path, manifest: dict) -> None: built = (manifest.get('environment') or {}).get('kumo_engine') raise RuntimeError( f'This DiskGraph cannot open the snapshot exported by kumo-engine ' - f'{built} ({error}). Rebuild it with this kumo-engine or load it ' - 'with the exporting one.') from error + f'{built} ({error}); load it with that kumo-engine.') from error diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index 7582f6bb3..ff2b78825 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -225,8 +225,7 @@ def _schema_tables(config: Mapping[str, Any]) -> Dict[str, Dict[str, Any]]: def schema_changes(before: Mapping[str, Any], after: Mapping[str, Any]) -> List[str]: - r"""How two graph configs differ in tables, columns, types, keys and - edges; empty when only connectors or table locations differ.""" + r"""Schema differences between two graph configs.""" old, new = _schema_tables(before), _schema_tables(after) changes = [ f"table {name} {'removed' if name in old else 'added'}" @@ -311,7 +310,6 @@ def _load_local(connector: Any, path: str, root: str, try: load_into(connector, path, root, manifest, target) except FileNotFoundError: - # A newer export of this snapshot committed meanwhile. newer_root, newer = read(connector, path) if (newer_root == root or newer['snapshot_id'] != manifest['snapshot_id']): @@ -400,26 +398,18 @@ class GraphSnapshot: @property def _graph_key(self) -> str: - """The `Graph.id` this snapshot was built from.""" return self.id.partition('@')[0] @property def _data_version(self) -> str: - """The full source-version fingerprint; `id` holds its first 16.""" return Path(self.index_path).name @cached_property def _native_graph(self) -> NativeGraph: - """`graph_config` normalized, as `Graph.native_graph`.""" return normalize_graph(self.graph_config) def export(self, output_config: 'SnapshotOutputConfig') -> str: - r"""Copy the index to `//` and - return that location. An earlier export there stays loadable until - the new one commits, then is replaced. The connector must hold files - at the path: an S3, GCS or file connector, or a volume or stage - whose sink offers a filesystem. - """ + r"""Export the snapshot via ``output_config``; returns its location.""" from kumo.engine._snapshot_io import export return export(self, output_config.connector, output_config.path) @@ -429,15 +419,13 @@ def load_from_connector(cls, connector: Connector, *, path: Optional[str] = None, id: Optional[str] = None, graph: Optional['Graph'] = None) -> 'GraphSnapshot': - r"""Copy the snapshot exported at `path` into the local store. + r"""Load a snapshot exported at `path`. Args: connector: a connector, or a `kumo.MLOps` for its artifact root. path: the location `export` returned. - id: instead of `path`, the id of a snapshot stored under the - artifact root of the `kumo.MLOps` passed as `connector`. - graph: if given, warn when its source tables changed since the - snapshot was built (see `source_changes`). + id: a snapshot id under the `kumo.MLOps` root, instead of `path`. + graph: warn if its sources changed since the snapshot. """ from kumo.engine._mlops import MLOps, on_root, snapshot_path from kumo.engine._snapshot_io import read @@ -475,8 +463,7 @@ def _warn_changed_sources(self, graph: 'Graph', path: str) -> None: UserWarning, stacklevel=3) def source_changes(self, graph: 'Graph') -> Dict[str, Dict[str, Any]]: - r"""Each of `graph`'s tables' source status since this snapshot: - `unchanged`, `changed` or `unknown`. Never raises.""" + r"""Each of `graph`'s tables as `unchanged`, `changed` or `unknown`.""" from kumo_core.sources.parallel import run_in_parallel try: @@ -498,7 +485,7 @@ def source_changes(self, graph: 'Graph') -> Dict[str, Dict[str, Any]]: @property def manifest(self) -> SnapshotManifest | None: - """The snapshot manifest stored next to the index, if any.""" + """The stored snapshot manifest, if any.""" return read_manifest(Path(self.index_path) / PROVENANCE_FILE, SnapshotManifest) @@ -677,12 +664,9 @@ def snapshot(self, *, record_sources: bool = True, it's reused as is. Args: - mlops: a `kumo.MLOps`; with MLOps on (this or a `kumo.mlops` - block), the snapshot is also exported under the artifact - root's `snapshots/`, unless it is there already. - record_sources: if False, the snapshot manifest records each - source as connector and table name only (no bucket, path or - account) and Custom SQL as a hash only (no SQL text). + mlops: also export the snapshot under this `kumo.MLOps` root. + record_sources: if False, record sources without locations or + SQL text. """ self.validate() for table in self._named_tables.values(): diff --git a/packages/kumo-engine/kumo/engine/pquery/predictive_query.py b/packages/kumo-engine/kumo/engine/pquery/predictive_query.py index fcafb625e..a6e16d5dd 100644 --- a/packages/kumo-engine/kumo/engine/pquery/predictive_query.py +++ b/packages/kumo-engine/kumo/engine/pquery/predictive_query.py @@ -16,9 +16,6 @@ `to_portable(definition)` -- the schema-level conversion, not an ingested snapshot -- that this doesn't need a `GraphSnapshot` either. -`PredictiveQuery._from_snapshot(snapshot, query)` binds a snapshot instead and -parses against `snapshot.graph_config`. - Construction itself stays I/O-free, consistent with `Table`/`Graph`: nothing parses until `_parsed` is actually touched (by `validate()`, `get_task_type()`, or a caller directly). A failed parse is *not* cached -- `functools. @@ -109,13 +106,6 @@ class PredictiveQuery: @classmethod def _from_snapshot(cls, snapshot: GraphSnapshot, query: str) -> 'PredictiveQuery': - r"""Bind `query` to `snapshot`'s graph without its `Graph` object. - - Snapshots passed to this query's methods must come from the same - graph (same `graph_key`). That only checks the snapshot belongs to - this query's graph; it says nothing about compatibility with a - trained model. - """ return cls(graph=snapshot, query=query) @cached_property diff --git a/packages/kumo-engine/kumo/engine/pquery/training_table.py b/packages/kumo-engine/kumo/engine/pquery/training_table.py index 6a772b315..c78efb581 100644 --- a/packages/kumo-engine/kumo/engine/pquery/training_table.py +++ b/packages/kumo-engine/kumo/engine/pquery/training_table.py @@ -53,20 +53,18 @@ class TrainingTable: required_stages: tuple[str, ...] #: Effective generation settings carried forward into ``Trainer.fit``. generation_plan: TrainingTableGenerationPlan - #: The artifact id: the same snapshot, query, plan, split seed and - #: engine build give the same id. None for a hand-built table. id: Optional[str] = field(default=None, compare=False) @property def tt_id(self) -> str: - """The id DiskGraph assigned to this table at generation.""" + """The DiskGraph id of this table.""" from kumo_core.training.export import read_spec return read_spec(str(self.staging_dir))['tt_id'] @property def manifest(self) -> Optional[TrainingTableManifest]: - """Lineage recorded at generation; None for older tables.""" + """The recorded manifest, if any.""" return read_manifest(self.staging_dir / PROVENANCE_FILE, TrainingTableManifest) @@ -122,10 +120,7 @@ def _export_sidecars(self, connector, output_config) -> Dict[str, Any]: def load_from_connector(cls, connector: Any, *, path: str, pquery: "PredictiveQuery", # noqa: F821 snapshot: GraphSnapshot) -> "TrainingTable": - r"""Materialize a table written by :meth:`export` (a directory with a - ``SPLIT`` column) against ``snapshot``, as :meth:`from_source_table` - with ``split_col='SPLIT'``. A `kumo.MLOps` ``connector`` reads - ``path`` under its artifact root (``training_tables/``).""" + r"""Load a table written by :meth:`export` against ``snapshot``.""" from kumo.engine._mlops import on_root connector, path = on_root(connector, path) diff --git a/packages/kumo-engine/kumo/engine/trainer/_spec.py b/packages/kumo-engine/kumo/engine/trainer/_spec.py index 6c7d4a660..8d8e04897 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_spec.py +++ b/packages/kumo-engine/kumo/engine/trainer/_spec.py @@ -21,8 +21,6 @@ SPEC_FILE = 'training-spec.json' SPEC_VERSION = 1 -#: What `kumo_core`'s `VitalArtifactPublicationPolicy` names its own -#: manifest; `training-spec.json` is listed in it like every other file. MANIFEST_FILE = 'publication-manifest.json' SPEC_ROLE = 'training_spec' PROVENANCE_ROLE = 'provenance' diff --git a/packages/kumo-engine/kumo/engine/trainer/output_config.py b/packages/kumo-engine/kumo/engine/trainer/output_config.py index 0127d2ed4..54ed5a68b 100644 --- a/packages/kumo-engine/kumo/engine/trainer/output_config.py +++ b/packages/kumo-engine/kumo/engine/trainer/output_config.py @@ -96,8 +96,7 @@ def __post_init__(self) -> None: @dataclass(frozen=True) class SnapshotOutputConfig: - """Export a graph snapshot under ``path`` via ``connector``, a file - destination (`GraphSnapshot.export`).""" + """Where `GraphSnapshot.export` writes.""" connector: Connector path: str diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 641e229e7..d35f5f28d 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -113,17 +113,17 @@ class TrainingResult: @property def id(self) -> Optional[str]: - """The model's id: its bundle id.""" + """The bundle id.""" return self.bundle_id @property def manifest(self) -> Optional[ModelManifest]: - """The bundle's ``_provenance.json``; None for older bundles.""" + """The bundle manifest, if any.""" return read_provenance(self.bundle_dir) @property def snapshot_id(self) -> Optional[str]: - """The id of the snapshot the model was trained on.""" + """The id of the training snapshot.""" if not (self.bundle_dir / SPEC_FILE).is_file(): return None return read_training_spec(self.bundle_dir).snapshot_id @@ -243,8 +243,7 @@ def _find_holdout(log_dir: str) -> tuple[Optional[Path], Optional[bool]]: def _bundle_holdout(result: TrainingResult) -> TrainingResult: - r"""Copy the holdout and its size record into the bundle, so a model - loaded from a registry or path still has them.""" + r"""Copy the holdout into the bundle.""" from kumo_core.artifacts import add_bundle_file if result.holdout_path is None: @@ -279,8 +278,7 @@ def _listed_holdout(bundle_dir: Path) -> tuple[Optional[Path], Optional[bool]]: def _store_model(result: TrainingResult, settings: Any) -> None: - r"""Copy the bundle to ``/models//`` unless a complete - copy is there; failures warn.""" + r"""Copy the bundle under the root unless already there.""" from kumo_core.object_store_sink import holds_bundle, upload_bundle from pyarrow.fs import FileSelector @@ -757,13 +755,8 @@ def predict( ) -> PredictionResult: r"""Score a snapshot and return durable file-backed results. - ``predict(snapshot)`` scores the trained query's population at - ``prediction_time`` (as :meth:`PredictiveQuery.generate_prediction_table`); - ``entities`` is a source table of ``ENTITY1`` (and ``TIMESTAMP`` - where the task needs one) rows to score instead. It warns when the - snapshot was taken of another graph than the one the model was - trained on. ``predict(prediction_table, snapshot=...)`` scores a - prediction table generated beforehand, as before. + Scores the query's population at ``prediction_time``, or the + ``entities`` given, or a ``prediction_table`` generated beforehand. The trainer must contain fitted state from :meth:`fit`, :meth:`load`, or :meth:`load_from_connector`. Without ``output_config``, From 5746dee37b6d45a2168f0a42477f85bb0b2d4816 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 01:06:36 +0000 Subject: [PATCH 160/278] Trim kumo-core and kumo-s3 docstrings and comments Co-Authored-By: Claude Opus 5.5 --- packages/kumo-core/src/kumo_core/artifacts.py | 4 ++-- packages/kumo-core/src/kumo_core/sources/identity.py | 3 +-- packages/kumo-core/src/kumo_core/sources/object_store.py | 2 -- packages/kumo-s3/src/kumo_s3/sink.py | 4 +--- 4 files changed, 4 insertions(+), 9 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/artifacts.py b/packages/kumo-core/src/kumo_core/artifacts.py index 2b9f5cdbc..f3b984386 100644 --- a/packages/kumo-core/src/kumo_core/artifacts.py +++ b/packages/kumo-core/src/kumo_core/artifacts.py @@ -345,7 +345,7 @@ def verify_model_bundle(bundle: Path) -> str: def bundle_id_from_manifest(bundle: Path) -> str: - """The bundle id from the digests a freshly published bundle lists.""" + """The bundle id from the manifest digests.""" manifest_path = Path(bundle) / VitalArtifactPublicationPolicy.MANIFEST_NAME manifest = json.loads(manifest_path.read_text(encoding='utf-8')) return _bundle_id((str(entry['relative_path']), str(entry['sha256'])) @@ -360,7 +360,7 @@ def _bundle_id(pairs: Iterable[tuple[str, str]]) -> str: def add_bundle_file(bundle_dir: Path, role: str, relative_path: str) -> None: - """List one file already written in a bundle in its publication manifest.""" + """List a written bundle file in the publication manifest.""" bundle_dir = Path(bundle_dir) path = bundle_dir / relative_path manifest_path = bundle_dir / VitalArtifactPublicationPolicy.MANIFEST_NAME diff --git a/packages/kumo-core/src/kumo_core/sources/identity.py b/packages/kumo-core/src/kumo_core/sources/identity.py index 0e3d8e664..2941f7ef1 100644 --- a/packages/kumo-core/src/kumo_core/sources/identity.py +++ b/packages/kumo-core/src/kumo_core/sources/identity.py @@ -21,8 +21,7 @@ class SourceVersion(str): - """A version string that also says what kind of evidence it is: - ``exact``, ``timestamp``, ``listing``, ``hour_bucket`` or ``unknown``.""" + """A version string tagged with its evidence ``kind``.""" kind: str dependencies: Mapping[str, str] | None diff --git a/packages/kumo-core/src/kumo_core/sources/object_store.py b/packages/kumo-core/src/kumo_core/sources/object_store.py index 657649388..312e49012 100644 --- a/packages/kumo-core/src/kumo_core/sources/object_store.py +++ b/packages/kumo-core/src/kumo_core/sources/object_store.py @@ -102,8 +102,6 @@ def _source_version(path: str, entries: Sequence[Any], identity: ObjectIdentity | None = None) -> str: if len(entries) == 1 and entries[0].path.rstrip("/") == path.rstrip("/"): if identity is not None: - # The provider's version or entity tag; a second-resolution - # modification time cannot tell same-second rewrites apart. token, version, *_ = (*identity(entries[0].path), "") return SourceVersion(version or token.strip('"'), "exact") latest = entries[0].mtime diff --git a/packages/kumo-s3/src/kumo_s3/sink.py b/packages/kumo-s3/src/kumo_s3/sink.py index cc650cb4d..0a4a7023d 100644 --- a/packages/kumo-s3/src/kumo_s3/sink.py +++ b/packages/kumo-s3/src/kumo_s3/sink.py @@ -21,9 +21,7 @@ def __repr__(self) -> str: return f'{type(self).__name__}(identity={self._identity!r})' def delete_file(self, connector: Any, path: str) -> None: - r"""Deletes the object at ``path`` (``/``). Unlike - pyarrow's delete, which recreates an emptied parent as a zero-byte - ``/`` object, this leaves no object behind.""" + r"""Delete ``path`` (``/``) leaving no ``/`` object.""" bucket, _, key = path.partition('/') session, _, connection = self._resolved_session(connector, bucket) if isinstance(session, S3FileSystem): From 4bf09897ee1b706f9208f8db72d5dd1e92d0fda2 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 01:06:36 +0000 Subject: [PATCH 161/278] Trim test comments and match the shortened messages Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/test/test_holdout_df.py | 2 +- packages/kumo-engine/test/test_mlops.py | 3 +-- packages/kumo-engine/test/test_query_from_snapshot.py | 1 - packages/kumo-engine/test/test_snapshot_export.py | 11 +++++------ packages/kumo-engine/test/test_snapshot_manifest.py | 1 - 5 files changed, 7 insertions(+), 11 deletions(-) diff --git a/packages/kumo-engine/test/test_holdout_df.py b/packages/kumo-engine/test/test_holdout_df.py index 290cce9cb..6f2435141 100644 --- a/packages/kumo-engine/test/test_holdout_df.py +++ b/packages/kumo-engine/test/test_holdout_df.py @@ -142,7 +142,7 @@ def test_a_loaded_or_fetched_model_keeps_its_holdout(graph, tmp_path, def test_a_bundle_without_a_holdout_says_so(graph, monkeypatch): monkeypatch.setattr('kumo.engine.trainer.trainer._bundle_holdout', - lambda result: result) # as bundles before the copy + lambda result: result) result = _fit(graph, QUERY) loaded = kumo.Trainer.load(result.bundle_dir).result diff --git a/packages/kumo-engine/test/test_mlops.py b/packages/kumo-engine/test/test_mlops.py index 3d3af18d9..6a740aec2 100644 --- a/packages/kumo-engine/test/test_mlops.py +++ b/packages/kumo-engine/test/test_mlops.py @@ -55,8 +55,7 @@ def test_a_connector_root_outside_its_root_dir_fails_when_built(): assert kumo.MLOps(artifact_root=(data, 's3://bucket/data/kumo'), store=False).root with pytest.raises(ValueError, match=r"'s3://bucket/kumo' is outside " - r"this S3Connector's root_dir s3://bucket/data .*" - r"S3Connector\(root_dir='s3://bucket/kumo'"): + r"this S3Connector's root_dir s3://bucket/data"): kumo.MLOps(artifact_root=(data, 's3://bucket/kumo'), store=False) diff --git a/packages/kumo-engine/test/test_query_from_snapshot.py b/packages/kumo-engine/test/test_query_from_snapshot.py index c22709da8..8afcbc7ef 100644 --- a/packages/kumo-engine/test/test_query_from_snapshot.py +++ b/packages/kumo-engine/test/test_query_from_snapshot.py @@ -40,7 +40,6 @@ def scores(**options): legacy.sort_values('ENTITY').reset_index(drop=True), expected) with pytest.raises(ValueError, match='not both'): scores(entities=rows, prediction_time=graph_table.timestamps[0]) - # Same schema under another graph id: scored, with an info log only. other = replace(snapshot, id='other@' + snapshot.id.split('@')[1]) scored = trainer.predict(other, binary_classification_threshold=0.5).frame pd.testing.assert_frame_equal( diff --git a/packages/kumo-engine/test/test_snapshot_export.py b/packages/kumo-engine/test/test_snapshot_export.py index 3d0f7c715..e6b4f18db 100644 --- a/packages/kumo-engine/test/test_snapshot_export.py +++ b/packages/kumo-engine/test/test_snapshot_export.py @@ -204,8 +204,7 @@ def racing(*args): def _rebuilt(snapshot, tmp_path, name): - r"""Another build of the same logical snapshot: same id and file sizes, - one digit different in ``col_stats.json``.""" + r"""Same id and sizes, one digit changed in ``col_stats.json``.""" index = (tmp_path / name / snapshot.index_path.parent.name / snapshot.index_path.name) shutil.copytree(snapshot.index_path, index) @@ -519,9 +518,9 @@ def test_snapshot_manifest_travels_unchanged_or_is_reconstructed( def test_warehouse_connectors_are_not_supported(scratch, graph): connector = SimpleNamespace(sink_backend=object()) - with pytest.raises(NotImplementedError, match='FileConnector'): + with pytest.raises(NotImplementedError, match='file connector'): graph.snapshot().export(SnapshotOutputConfig(connector, 'x')) - with pytest.raises(NotImplementedError, match='FileConnector'): + with pytest.raises(NotImplementedError, match='file connector'): _load(connector, 'x') @@ -753,7 +752,7 @@ def _remove(self, path, *, tree): with filesystem.open_input_file(info.path) as reader: assert (loaded.index_path / info.path[len(root) + 1:] ).read_bytes() == reader.read() - with pytest.raises(NotImplementedError, match='file destination'): + with pytest.raises(NotImplementedError, match='S3, GCS or file connector'): snapshot.export(SnapshotOutputConfig(connector, 'DB.SCHEMA')) @@ -784,5 +783,5 @@ def test_a_mlops_block_stores_each_snapshot_once(tmp_path, scratch, graph, def test_loading_a_missing_export_names_its_full_path(tmp_path, exports): missing = tmp_path / 'exports' / 'snapshots' / 'graph' / 'version' with pytest.raises(FileNotFoundError, - match=f'No complete snapshot export at {missing}: '): + match=f'No snapshot export at {missing} '): _load(exports, 'snapshots/graph/version') diff --git a/packages/kumo-engine/test/test_snapshot_manifest.py b/packages/kumo-engine/test/test_snapshot_manifest.py index 3f8acedb1..b6294a6b2 100644 --- a/packages/kumo-engine/test/test_snapshot_manifest.py +++ b/packages/kumo-engine/test/test_snapshot_manifest.py @@ -51,7 +51,6 @@ def make(events: int | None = None, **options) -> GraphSnapshot: 'EID': np.arange(events), 'ENT_ID': np.arange(events) % 5, }).to_parquet(data / 'events.parquet') - # Single-file versions are mtime seconds. stamp = 1.7e9 + events os.utime(data / 'events.parquet', (stamp, stamp)) connector = FileConnector(root_dir=str(data)) From 3bd393ab0840d11ffb567649841d0ddd0875aa35 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 01:31:08 +0000 Subject: [PATCH 162/278] Move _flatten next to its only caller in _run_spec Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_tracking.py | 9 --------- .../kumo-engine/kumo/engine/trainer/_run_spec.py | 12 ++++++++++-- 2 files changed, 10 insertions(+), 11 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_tracking.py b/packages/kumo-engine/kumo/engine/_tracking.py index 850d722ee..926d223bf 100644 --- a/packages/kumo-engine/kumo/engine/_tracking.py +++ b/packages/kumo-engine/kumo/engine/_tracking.py @@ -18,7 +18,6 @@ Any, Callable, Dict, - Iterator, List, Mapping, Optional, @@ -578,14 +577,6 @@ def _rank() -> int: return 0 -def _flatten(value: Any, prefix: str = '') -> Iterator[Tuple[str, Any]]: - if isinstance(value, Mapping): - for key, item in value.items(): - yield from _flatten(item, f'{prefix}{key}.') - else: - yield prefix[:-1], value - - def _param_value(value: Any) -> str: if isinstance(value, str): return value diff --git a/packages/kumo-engine/kumo/engine/trainer/_run_spec.py b/packages/kumo-engine/kumo/engine/trainer/_run_spec.py index 167b0877d..4be3db4be 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_run_spec.py +++ b/packages/kumo-engine/kumo/engine/trainer/_run_spec.py @@ -7,12 +7,12 @@ from dataclasses import dataclass, field, replace from datetime import datetime, timezone from pathlib import Path -from typing import Any, Mapping, Optional, Sequence +from typing import Any, Iterator, Mapping, Optional, Sequence, Tuple from kumo_core.artifacts import PROVENANCE_FILE from kumo_core.summaries import Summary -from kumo.engine._tracking import RunSpec, _flatten +from kumo.engine._tracking import RunSpec from kumo.engine._version import __version__ @@ -78,3 +78,11 @@ def _registered(model: Any, alias: Optional[str]) -> Mapping[str, Any]: return {} return {'model_name': model.name, 'model_version': model.version, 'model_alias': alias} + + +def _flatten(value: Any, prefix: str = '') -> Iterator[Tuple[str, Any]]: + if isinstance(value, Mapping): + for key, item in value.items(): + yield from _flatten(item, f'{prefix}{key}.') + else: + yield prefix[:-1], value From 1d2bb99f90c8385c40e1979d3d91d1979a9deace Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 01:31:11 +0000 Subject: [PATCH 163/278] Drop fit_spec's tags parameter that only tests passed Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/trainer/_run_spec.py | 4 +--- packages/kumo-engine/test/test_tracking_contract.py | 4 ++-- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/_run_spec.py b/packages/kumo-engine/kumo/engine/trainer/_run_spec.py index 4be3db4be..4ac3ad557 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_run_spec.py +++ b/packages/kumo-engine/kumo/engine/trainer/_run_spec.py @@ -30,8 +30,7 @@ class FitContext: model_name: Optional[str] = None -def fit_spec(context: FitContext, *, - tags: Optional[Mapping[str, Any]] = None) -> RunSpec: +def fit_spec(context: FitContext) -> RunSpec: started = context.started or datetime.now(timezone.utc) return RunSpec( kind='fit', @@ -43,7 +42,6 @@ def fit_spec(context: FitContext, *, 'tune_metric': (context.model_plan.get('training_job') or {}).get('tune_metric'), 'engine_version': __version__}, - tags=dict(tags or {}), inputs=dict(context.inputs), artifacts={ 'query.pql': context.query, diff --git a/packages/kumo-engine/test/test_tracking_contract.py b/packages/kumo-engine/test/test_tracking_contract.py index 2acda5a74..fed46a8c9 100644 --- a/packages/kumo-engine/test/test_tracking_contract.py +++ b/packages/kumo-engine/test/test_tracking_contract.py @@ -140,7 +140,7 @@ def replay(tracker=None, events=EVENTS, result=RESULT, error=None, record.name: record for hook, record in recorder.calls if hook == 'on_trial' } - spec = fit_spec(context, tags=USER_TAGS) + spec = replace(fit_spec(context), tags=USER_TAGS) if result is not None: spec = fitted_spec(spec, result, events[-1].metrics.get('seconds')) spec = replace(spec, steps=series if series is not None else { @@ -346,7 +346,7 @@ def test_records_of_the_recorded_fit(): } with pytest.warns(UserWarning, match="'kumo.bundle_id', 'training_job.tune_metric'"): - spoofed = run_record(fit_spec(CONTEXT, tags={ + spoofed = run_record(replace(fit_spec(CONTEXT), tags={ 'kumo.bundle_id': 'x', 'training_job.tune_metric': 'y', 'a': 1})) assert spoofed.tags.keys() - fit.tags.keys() == {'a'} assert fit.metrics == {'test/mae': 0.32, 'seconds': 9.0} From 27c9a5ba3c2ab65093e58076d4992a7c03e7680c Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 01:31:33 +0000 Subject: [PATCH 164/278] Reuse bundle_manifest for a result's manifest or stub Co-Authored-By: Claude Opus 5.5 --- .../kumo-engine/kumo/engine/integrations/mlflow.py | 11 ++++++----- .../kumo-engine/kumo/engine/trainer/_provenance.py | 6 +++--- packages/kumo-engine/kumo/engine/trainer/trainer.py | 11 +++++------ 3 files changed, 14 insertions(+), 14 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index ea1dd0e67..738decdf5 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -211,8 +211,8 @@ def load_model(self, location: ModelRegistryLocation, def register(self, result: Any, location: ModelRegistryLocation) -> PublishedModel: r"""Registers a `TrainingResult` at ``location``.""" - from kumo.engine._manifest import stub from kumo.engine.graph import GraphSnapshot + from kumo.engine.trainer._provenance import bundle_manifest from kumo.engine.trainer._spec import read_training_spec from kumo.engine.trainer.trainer import registry_pointers @@ -234,7 +234,8 @@ def register(self, result: Any, experiment_id, tags={'kumo.bundle_id': str(result.id)}).info.run_id info = _log_model(self.target, client, result.bundle_dir, - result.manifest or stub('model', result.id), + bundle_manifest(result.bundle_dir, + result.id), snapshot, run_id, experiment_id) with self.target.scope(): client.set_terminated(run_id, 'FINISHED') @@ -594,9 +595,9 @@ def log_model( none. target: as in `MlflowTracking`. """ - from kumo.engine._manifest import stub from kumo.engine._mlops import snapshot_uri from kumo.engine._tracking import fit_inputs + from kumo.engine.trainer._provenance import bundle_manifest mlflow = import_provider('mlflow', 'mlflow') result = trainer.result @@ -618,8 +619,8 @@ def log_model( tracking._experiment_for(client)).info.run_id experiment_id = client.get_run(run_id).info.experiment_id _log_inputs(mlflow, client, run_id, fit_inputs(result.manifest)) - info = _log_model(target, client, result.bundle_dir, result.manifest or - (None if result.id is None else stub('model', result.id)), + info = _log_model(target, client, result.bundle_dir, + bundle_manifest(result.bundle_dir, result.id), snapshot, run_id, experiment_id, name, snapshot_uri(serving_snapshot)) tracking._log_metrics(run_id, result.metrics, model_id=info.model_id) diff --git a/packages/kumo-engine/kumo/engine/trainer/_provenance.py b/packages/kumo-engine/kumo/engine/trainer/_provenance.py index 18c234e63..0001fcb30 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_provenance.py +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -79,8 +79,8 @@ def cut(model: Optional[ModelManifest]) -> Optional[ModelManifest]: derived_from=cut(parents.derived_from))) -def parent_model(bundle_dir: Path, bundle_id: Optional[str] - ) -> Optional[ModelManifest]: +def bundle_manifest(bundle_dir: Path, + bundle_id: Optional[str]) -> Optional[ModelManifest]: r"""The manifest of the bundle at ``bundle_dir``, or a stub.""" manifest = read_provenance(bundle_dir) if manifest is not None: @@ -102,7 +102,7 @@ def record_derivation(bundle_dir: Path, parent_dir: Path, bundle_id_from_manifest, ) - parent = parent_model(parent_dir, parent_id) + parent = bundle_manifest(parent_dir, parent_id) write_manifest( bundle_dir / PROVENANCE_FILE, ModelManifest( diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index a83b5e5bf..a8d38134f 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -79,7 +79,6 @@ ModelManifest, TrackingRef, artifact_summaries, - stub, ) from kumo.engine._mlops import full_error, warn_failure from kumo.engine._predictor import TRAINER_PREDICTOR @@ -97,7 +96,7 @@ from kumo.engine.pquery.training_table import TrainingTable from kumo.engine.trainer._provenance import ( FitRecord, - parent_model, + bundle_manifest, read_provenance, record_derivation, ) @@ -760,9 +759,9 @@ def fit( **result.metrics, 'seconds': seconds }, params={'status': 'succeeded'}) registered, error = _store_and_register( - result.bundle_dir, result.manifest or ( - None if result.bundle_id is None else stub( - 'model', result.bundle_id)), settings, session=session, + result.bundle_dir, + bundle_manifest(result.bundle_dir, result.bundle_id), settings, + session=session, snapshot=snapshot, metrics=result.metrics, warn=warn, stored=(deps.errors or [None])[0]) if registered is not None: @@ -1001,7 +1000,7 @@ def _fit( model_plan=effective_plan, provenance=FitRecord( train_table=train_table, - warm_start=None if warm_start is None else parent_model( + warm_start=None if warm_start is None else bundle_manifest( warm_start._training_result.bundle_dir, warm_start_bundle_id), started_at=started_at, fit_mode=fit_mode, From 955e37bee957861029efa7073a9d62609fdee0ff Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 01:31:44 +0000 Subject: [PATCH 165/278] Verify a fetched asset with loaded(), dropping a fallback no record reaches Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_deps.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_deps.py b/packages/kumo-engine/kumo/engine/_deps.py index 6c11abd07..5cc78627f 100644 --- a/packages/kumo-engine/kumo/engine/_deps.py +++ b/packages/kumo-engine/kumo/engine/_deps.py @@ -267,8 +267,7 @@ def _verified(directory: Path, record: Dependency) -> bool: stored = read_manifest(directory / PROVENANCE_FILE, AssetManifest) if stored is None or stored.id != record.sha256: return False - loader = imported(record.entry_point or entry_point(FileAsset)) - return bool(loader.load(directory, record).fingerprint() == record.sha256) + return bool(loaded(directory, record).fingerprint() == record.sha256) def loaded(directory: Path, record: Dependency) -> Any: From 192cdd71c850f5bb9fe2c3ef70facfb335a8d192 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 01:31:44 +0000 Subject: [PATCH 166/278] Read GloVeVectors.path directly in TextEncoderAsset Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_embeddings.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/kumo-engine/kumo/engine/_embeddings.py b/packages/kumo-engine/kumo/engine/_embeddings.py index e06675647..0591588b5 100644 --- a/packages/kumo-engine/kumo/engine/_embeddings.py +++ b/packages/kumo-engine/kumo/engine/_embeddings.py @@ -63,7 +63,7 @@ class TextEncoderAsset: def __init__(self, name: str, model: Any) -> None: self.name, self.model = name, model self.version = model.label - self.path = getattr(model, 'path', None) + self.path = model.path def fingerprint(self) -> str: return str(self.model.fingerprint()) From 67651fbdb59e3a0a58a1427e33b6400a506e4b0c Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 01:31:45 +0000 Subject: [PATCH 167/278] Drop _new_run_name's unreachable None branch Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/integrations/snowflake.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/integrations/snowflake.py b/packages/kumo-engine/kumo/engine/integrations/snowflake.py index 8dc4aa931..04b2ce151 100644 --- a/packages/kumo-engine/kumo/engine/integrations/snowflake.py +++ b/packages/kumo-engine/kumo/engine/integrations/snowflake.py @@ -206,9 +206,7 @@ def log_metrics(self, metrics: Mapping[str, float], step: int) -> None: self._log_metrics(metrics, _EPOCH_METRICS, step=step) def _new_run_name(self, tracking: Any, experiment: Any, - run_name: Optional[str]) -> Optional[str]: - if run_name is None: - return None + run_name: str) -> Optional[str]: # start_run matches name with LIKE. existing = [ str(row['name']) From 8d63cbc89f167c8a2b93c4b58f8ad43ab8dc7c5b Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 01:32:09 +0000 Subject: [PATCH 168/278] Keep kumoapi's registry key prefix private Co-Authored-By: Claude Opus 5.5 --- packages/kumo-api/kumoapi/plugin.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/packages/kumo-api/kumoapi/plugin.py b/packages/kumo-api/kumoapi/plugin.py index 633effb05..709e4a59d 100644 --- a/packages/kumo-api/kumoapi/plugin.py +++ b/packages/kumo-api/kumoapi/plugin.py @@ -46,7 +46,6 @@ __all__ = [ 'BUNDLE_ID_KEY', 'ENGINE_MARKER', - 'REGISTRY_KEY_PREFIX', 'ArtifactLocation', 'ArtifactSink', 'FileStoreSink', @@ -66,7 +65,7 @@ 'registry_lineage', ] -REGISTRY_KEY_PREFIX = 'kumo_' +_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}') @@ -761,7 +760,7 @@ def registry_lineage(recorded: Mapping[str, Any]) -> Dict[str, str]: return { key: str(value) for key, value in recorded.items() - if key.startswith(REGISTRY_KEY_PREFIX) and key != ENGINE_MARKER + if key.startswith(_REGISTRY_KEY_PREFIX) and key != ENGINE_MARKER } From a985faf3e9cc3379e3d012658ae2057948679268 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 01:32:09 +0000 Subject: [PATCH 169/278] Drop the searcher test's identity check on the fallback sampler Co-Authored-By: Claude Opus 5.5 --- packages/kumo-core/tests/test_searched_params.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/packages/kumo-core/tests/test_searched_params.py b/packages/kumo-core/tests/test_searched_params.py index c59c8390a..af255f322 100644 --- a/packages/kumo-core/tests/test_searched_params.py +++ b/packages/kumo-core/tests/test_searched_params.py @@ -34,5 +34,3 @@ def chosen(cfg): searcher.max_attempts = 0 cfg = searcher.suggest() assert searcher.sampled_params == chosen(cfg) - assert (searcher.sampled_params - is searcher.fallback_sampler.sampled_params) From b81eb5c915143183197735c6ca8c7e5ea4a81da4 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 01:32:09 +0000 Subject: [PATCH 170/278] Shorten the kumo-snowflake sdk extra comment Co-Authored-By: Claude Opus 5.5 --- packages/kumo-snowflake/pyproject.toml | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/packages/kumo-snowflake/pyproject.toml b/packages/kumo-snowflake/pyproject.toml index 47547c8b9..2f7591943 100644 --- a/packages/kumo-snowflake/pyproject.toml +++ b/packages/kumo-snowflake/pyproject.toml @@ -27,10 +27,9 @@ dependencies = [ [project.optional-dependencies] # What the SDK plugin half needs beyond the platform reader: Snowpark for the -# session it opens, the model registry for what a sink publishes (and, from -# 2.0, the experiment tracking kumo-engine records fits in; capped below 3 -# because that adapter reads ExperimentTracking internals), and pyyaml for the -# manifest that carries. The platform installs this distribution +# session it opens, the model registry for what a sink publishes (and +# experiment tracking, >=2.0; <3 because kumo-engine reads its internals), and +# pyyaml for the manifest that carries. The platform installs this distribution # without this extra and imports `SnowflakeSourceReader` alone, so none of it # is loaded there -- `__init__` exports the plugin lazily to keep that true. sdk = [ From 3fbf44b318e487f31c78b4f762e33e7de55244c0 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 01:35:14 +0000 Subject: [PATCH 171/278] Keep the model download inside ObjectStoreSinkBackend._download_file Co-Authored-By: Claude Opus 5.5 --- packages/kumo-core/src/kumo_core/object_store_sink.py | 11 +++-------- 1 file changed, 3 insertions(+), 8 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/object_store_sink.py b/packages/kumo-core/src/kumo_core/object_store_sink.py index bb4ec832d..b097f726b 100644 --- a/packages/kumo-core/src/kumo_core/object_store_sink.py +++ b/packages/kumo-core/src/kumo_core/object_store_sink.py @@ -47,13 +47,6 @@ def upload_file(filesystem: FileSystem, source: str, shutil.copyfileobj(reader, writer, _CHUNK) -def download_file(filesystem: FileSystem, source: str, - destination: str) -> None: - with filesystem.open_input_stream(source) as reader, \ - open(destination, 'wb') as writer: - shutil.copyfileobj(reader, writer, _CHUNK) - - def upload_bundle(filesystem: FileSystem, destination: str, bundle_dir: str) -> None: """Copy ``bundle_dir`` under ``destination``, provenance last.""" @@ -174,7 +167,9 @@ def publish_model( def _download_file(self, filesystem, source, destination): """Providers may transfer a model file without the streaming fallback.""" - download_file(filesystem, source, destination) + with filesystem.open_input_stream(source) as reader, \ + open(destination, 'wb') as writer: + shutil.copyfileobj(reader, writer, _CHUNK) def fetch_model( self, From 80445ea40a1e2bcacd1d6aaac469901411731812 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 01:35:24 +0000 Subject: [PATCH 172/278] Import tempfile at module level in object_store_sink Co-Authored-By: Claude Opus 5.5 --- packages/kumo-core/src/kumo_core/object_store_sink.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/object_store_sink.py b/packages/kumo-core/src/kumo_core/object_store_sink.py index b097f726b..30b9decb4 100644 --- a/packages/kumo-core/src/kumo_core/object_store_sink.py +++ b/packages/kumo-core/src/kumo_core/object_store_sink.py @@ -5,6 +5,7 @@ import os import posixpath import shutil +import tempfile from abc import ABC, abstractmethod from pathlib import Path from typing import Any, Mapping, Optional, Tuple @@ -289,9 +290,8 @@ def _write_one(filesystem: FileSystem, frame: pd.DataFrame, path: str, if not _chunked(filesystem): pq.write_table(table, path, filesystem=filesystem) return - import tempfile - from kumo_core.chunked_fs import spool_root + # Spool to learn the size. with tempfile.TemporaryDirectory( dir=spool_root(filesystem.handler)) as spool: From 91368050aa185a48bb902f8620b6cb5b9dec23f9 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 01:35:49 +0000 Subject: [PATCH 173/278] Type the file-sink helpers with FileStoreSink and Connector Co-Authored-By: Claude Opus 5.5 --- packages/kumo-core/src/kumo_core/object_store_sink.py | 3 ++- packages/kumo-engine/kumo/engine/connector/__init__.py | 10 +++++----- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/object_store_sink.py b/packages/kumo-core/src/kumo_core/object_store_sink.py index 30b9decb4..d09484b4c 100644 --- a/packages/kumo-core/src/kumo_core/object_store_sink.py +++ b/packages/kumo-core/src/kumo_core/object_store_sink.py @@ -13,6 +13,7 @@ import pandas as pd from kumoapi.jobs import WriteMode from kumoapi.plugin import ( + FileStoreSink, ModelRegistryLocation, PublishedModel, SourceConnector, @@ -394,7 +395,7 @@ def convert_csv(self, source, destination): class FilesystemSink(ObjectStoreSinkBackend): """The object-store file writers over any sink's ``filesystem()``.""" - def __init__(self, sink: Any) -> None: + def __init__(self, sink: FileStoreSink) -> None: self._sink = sink self.provider = type(sink).__name__ diff --git a/packages/kumo-engine/kumo/engine/connector/__init__.py b/packages/kumo-engine/kumo/engine/connector/__init__.py index e5ac09e6f..c2bdeea22 100644 --- a/packages/kumo-engine/kumo/engine/connector/__init__.py +++ b/packages/kumo-engine/kumo/engine/connector/__init__.py @@ -19,6 +19,8 @@ from typing import Any, Optional, Tuple +from kumoapi.plugin import FileStoreSink + from kumo.engine.connector.adls import ADLSConnector from kumo.engine.connector.base import Connector from kumo.engine.connector.bigquery import BigQueryConnector @@ -39,9 +41,7 @@ from kumo.engine.connector.source_table import SourceTable, SQLSourceTable -def _file_sink(connector: Any, path: str) -> Any: - from kumoapi.plugin import FileStoreSink - +def _file_sink(connector: Connector, path: str) -> Optional[FileStoreSink]: try: sink = connector.sink_backend except NotImplementedError: @@ -50,12 +50,12 @@ def _file_sink(connector: Any, path: str) -> Any: connector, path) else None -def is_file_destination(connector: Any, path: str) -> bool: +def is_file_destination(connector: Connector, path: str) -> bool: r"""Whether ``path`` is a file location rather than a table.""" return _file_sink(connector, path) is not None -def file_store(connector: Any, path: str) -> Optional[Tuple[Any, str]]: +def file_store(connector: Connector, path: str) -> Optional[Tuple[Any, str]]: r"""The ``(filesystem, path)`` of a file destination; None for a table.""" sink = _file_sink(connector, path) return None if sink is None else sink.filesystem(connector, path) From faad76d08df6239a247a63113d821aa3f4087524 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 01:35:52 +0000 Subject: [PATCH 174/278] Reuse the sink local in the file-store sink tests Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/test/test_filestore_sink.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/kumo-engine/test/test_filestore_sink.py b/packages/kumo-engine/test/test_filestore_sink.py index 8a7e36f52..da12f3bb0 100644 --- a/packages/kumo-engine/test/test_filestore_sink.py +++ b/packages/kumo-engine/test/test_filestore_sink.py @@ -163,7 +163,7 @@ def test_parquet_writes_stay_plain_up_to_one_part(connector): sink.write_table(connector, '@stage/small.parquet', frame) sink.write_table(connector, '@stage/big.parquet', pd.DataFrame({'x': [os.urandom(50).hex() for _ in range(100)]})) - store = connector.sink_backend.store + store = sink.store assert '@stage/small.parquet' in store.objects assert any( name.startswith('@stage/big.parquet.chunks/') @@ -178,7 +178,7 @@ def test_an_occupied_directory_is_refused_before_scoring(connector): sink = connector.sink_backend sink.preflight_write_files(connector, '@stage/out/') - connector.sink_backend.store.objects['@stage/out/old.parquet'] = b'x' + sink.store.objects['@stage/out/old.parquet'] = b'x' for mode in (None, 'OVERWRITE', 'APPEND'): with pytest.raises(ArtifactDestinationConflict, match='already'): sink.preflight_write_files(connector, '@stage/out/', mode) From dda58f571ea64e51df4b2a529057e9c570529719 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 01:49:00 +0000 Subject: [PATCH 175/278] Drop duplicate dockerignore entries the base already has Co-Authored-By: Claude Opus 5.5 --- tests/public/Dockerfile.dockerignore | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/tests/public/Dockerfile.dockerignore b/tests/public/Dockerfile.dockerignore index 14a52306a..c6d21c61a 100644 --- a/tests/public/Dockerfile.dockerignore +++ b/tests/public/Dockerfile.dockerignore @@ -14,15 +14,12 @@ !packages/kumo-core/src/kumo_core/artifact_staging.py # Model registries validate bundles through the portable artifact contract. !packages/kumo-core/src/kumo_core/artifacts.py -# Volume and stage sinks hold files through the shared chunked filesystem. -!packages/kumo-core/src/kumo_core/chunked_fs.py -!packages/kumo-core/src/kumo_core/object_store_sink.py -!packages/kumo-core/src/kumo_core/artifact_export.py !packages/kumo-core/src/kumo_core/config.py !packages/kumo-core/src/kumo_core/manifest.py !packages/kumo-core/src/kumo_core/io/__init__.py !packages/kumo-core/src/kumo_core/io/filesystem.py !packages/kumo-core/src/kumo_core/io/json.py +!packages/kumo-core/src/kumo_core/chunked_fs.py !packages/kumo-core/src/kumo_core/object_store_sink.py !packages/kumo-core/src/kumo_core/artifact_export.py !packages/kumo-core/src/kumo_core/stype_inference.py From c10be5a43efc4fd8779c67cd06efac66503a3af0 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 01:49:10 +0000 Subject: [PATCH 176/278] Share one upload-finish path between write and close Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_databricks/volume.py | 16 ++++++---------- 1 file changed, 6 insertions(+), 10 deletions(-) diff --git a/packages/kumo-databricks/src/kumo_databricks/volume.py b/packages/kumo-databricks/src/kumo_databricks/volume.py index 2e14c31d8..1f0657de1 100644 --- a/packages/kumo-databricks/src/kumo_databricks/volume.py +++ b/packages/kumo-databricks/src/kumo_databricks/volume.py @@ -203,7 +203,9 @@ def writable(self) -> bool: def write(self, data) -> int: piece = bytes(data) - self._put(piece) + if not self._offer(piece): + self._finish() + raise OSError(f'{self._path}: the upload stopped reading') return len(piece) def _offer(self, piece: Any) -> bool: @@ -215,14 +217,10 @@ def _offer(self, piece: Any) -> bool: continue return False - def _put(self, piece: Optional[bytes]) -> None: - if self._offer(piece): - return + def _finish(self) -> None: self._thread.join() if self._error is not None: raise self._error - if piece is not None: - raise OSError(f'{self._path}: the upload stopped reading') def abort(self) -> None: if self.closed: @@ -235,9 +233,7 @@ def close(self) -> None: if self.closed: return try: - self._put(None) - self._thread.join() - if self._error is not None: - raise self._error + self._offer(None) + self._finish() finally: super().close() From df21902fec64f4554f2e019ccaf710514b6520f5 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 01:49:13 +0000 Subject: [PATCH 177/278] Share the fake volume sink setup across volume tests Co-Authored-By: Claude Opus 5.5 --- packages/kumo-databricks/tests/test_volume.py | 23 +++++++++---------- 1 file changed, 11 insertions(+), 12 deletions(-) diff --git a/packages/kumo-databricks/tests/test_volume.py b/packages/kumo-databricks/tests/test_volume.py index 15cef7d32..6bcdf8bcc 100644 --- a/packages/kumo-databricks/tests/test_volume.py +++ b/packages/kumo-databricks/tests/test_volume.py @@ -25,6 +25,13 @@ def filesystem(objects=None, part_bytes=None): return PyFileSystem(handler), volume +def volume_sink(): + sink = DatabricksSinkBackend() + volume = FakeVolume() + sink._client = lambda connector: FakeWorkspaceClient(volume) + return sink, volume + + def write(fs, path, body, piece=700): with fs.open_output_stream(path) as writer: for start in range(0, len(body), piece): @@ -108,9 +115,7 @@ def refuse(path, contents, *, overwrite): def test_sink_filesystem_only_takes_volume_paths(): - sink = DatabricksSinkBackend() - volume = FakeVolume() - sink._client = lambda connector: FakeWorkspaceClient(volume) + sink, volume = volume_sink() connector = SimpleNamespace() fs, path = sink.filesystem(connector, f'{ROOT}/exports/') assert path == f'{ROOT}/exports' @@ -125,9 +130,7 @@ def test_sink_filesystem_only_takes_volume_paths(): def test_bundle_publish_and_fetch_through_a_volume_path(tmp_path): - sink = DatabricksSinkBackend() - volume = FakeVolume() - sink._client = lambda connector: FakeWorkspaceClient(volume) + sink, volume = volume_sink() bundle = tmp_path / 'bundle' (bundle / 'model').mkdir(parents=True) (bundle / 'training-spec.json').write_text('{}') @@ -145,9 +148,7 @@ def test_bundle_publish_and_fetch_through_a_volume_path(tmp_path): def test_predictions_and_frames_written_to_a_volume_path(tmp_path): - sink = DatabricksSinkBackend() - volume = FakeVolume() - sink._client = lambda connector: FakeWorkspaceClient(volume) + sink, volume = volume_sink() part = tmp_path / 'part-0.parquet' pq.write_table(pa.table({'users': [1, 2], 'score': [0.1, 0.9]}), part) written = sink.write_files(SimpleNamespace(), f'{ROOT}/predictions/', @@ -164,9 +165,7 @@ def test_an_occupied_volume_directory_is_refused_before_scoring(tmp_path): from kumo_core.artifact_export import ArtifactDestinationConflict from kumoapi.jobs import WriteMode - sink = DatabricksSinkBackend() - volume = FakeVolume() - sink._client = lambda connector: FakeWorkspaceClient(volume) + sink, volume = volume_sink() target = f'{ROOT}/predictions/' sink.preflight_write_files(SimpleNamespace(), target) part = tmp_path / 'part-0.parquet' From eb15c6d92ed41e76b7aa6ace2ff8b94f60b63d11 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 01:49:17 +0000 Subject: [PATCH 178/278] Reflow the volume README paragraph Co-Authored-By: Claude Opus 5.5 --- packages/kumo-databricks/README.md | 20 +++++++++----------- 1 file changed, 9 insertions(+), 11 deletions(-) diff --git a/packages/kumo-databricks/README.md b/packages/kumo-databricks/README.md index 2c0a1a53f..2874ff510 100644 --- a/packages/kumo-databricks/README.md +++ b/packages/kumo-databricks/README.md @@ -86,17 +86,15 @@ trainer.predict(prediction_table, snapshot=snapshot, output_config=PredictionOut For an MLOps artifact root on a volume, see `packages/kumo-engine/docs/mlops.md`. -A file whose size is known when it is written (bundle files, exported -Parquet parts) is stored as one object up to 4 GiB, so Spark -and other tools read it directly. A larger one streams into -`.chunks//00000`, `00001`, ... (4 GiB each, under the Files -API's 5 GiB upload cap) with an empty `_SUCCESS` written last, and Kumo reads -the newest committed generation back as one file. A file written without a -known size stays one object up to -256 MiB, staged on local disk under the engine scratch directory, and streams -into parts above that. The identity needs `USE CATALOG`, `USE SCHEMA`, and -`READ VOLUME` (loads) and `WRITE VOLUME` (exports) on the volume; no -warehouse is used. +A file whose size is known when it is written (bundle files, exported Parquet +parts) is stored as one object up to 4 GiB, so Spark and other tools read it +directly. A larger one streams into `.chunks//00000`, +`00001`, ... (4 GiB each, under the Files API's 5 GiB upload cap) with an empty +`_SUCCESS` written last, and Kumo reads the newest committed generation back as +one file. A file written without a known size stays one object up to 256 MiB, +staged on local disk under the engine scratch directory, and streams into parts +above that. The identity needs `USE CATALOG`, `USE SCHEMA`, and `READ VOLUME` +(loads) and `WRITE VOLUME` (exports) on the volume; no warehouse is used. ## Model registry From 6d4ee0ec91e4f27e3c06153f691b407486b82af3 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 01:50:11 +0000 Subject: [PATCH 179/278] Cut a warm-start parent to a stub with stub() Co-Authored-By: Claude Opus 5.5 --- .../kumo-engine/kumo/engine/trainer/_provenance.py | 11 +++-------- 1 file changed, 3 insertions(+), 8 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/_provenance.py b/packages/kumo-engine/kumo/engine/trainer/_provenance.py index 965ddb05d..b63bb3ac8 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_provenance.py +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -61,15 +61,10 @@ def embedded_model(manifest: ModelManifest) -> ModelManifest: r"""``manifest`` for embedding, its own model parents cut to stubs.""" copy = embedded(manifest) parents = copy.parents - if parents is None: + if parents is None or parents.warm_start is None: return copy - - def cut(model: Optional[ModelManifest]) -> Optional[ModelManifest]: - return None if model is None else ModelManifest(id=model.id, - reconstructed=True) - - return replace(copy, - parents=replace(parents, warm_start=cut(parents.warm_start))) + return replace(copy, parents=replace( + parents, warm_start=stub('model', parents.warm_start.id))) def parent_model(bundle_dir: Path, bundle_id: Optional[str] From 0999522cc5317cfdf3cac80ce9f08d64ca3c5a1c Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 01:50:11 +0000 Subject: [PATCH 180/278] Import the manifest helpers at the top of training_table Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/pquery/training_table.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/pquery/training_table.py b/packages/kumo-engine/kumo/engine/pquery/training_table.py index c78efb581..b84b26d9a 100644 --- a/packages/kumo-engine/kumo/engine/pquery/training_table.py +++ b/packages/kumo-engine/kumo/engine/pquery/training_table.py @@ -21,8 +21,10 @@ SourceTableEntry, TrainingTableManifest, TrainingTableParents, + embedded, read_manifest, stamp, + stub, write_manifest, ) from kumo.engine.connector.source_table import SourceTable @@ -331,12 +333,10 @@ def _table_id(table: TrainingTable, snapshot: GraphSnapshot, def _record_manifest(table: TrainingTable, snapshot: GraphSnapshot, request: TrainTableRequest, seed: Optional[int], *, source_table: Optional[SourceTableEntry] = None) -> None: - from kumo.engine._manifest import embedded, stub + from kumo_core.io.json import to_json_dict path = table.staging_dir / PROVENANCE_FILE try: - from kumo_core.io.json import to_json_dict - parent = snapshot.manifest write_manifest( path, From c31373cf51cae5e37001022e4616916d665bc1e0 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 01:50:20 +0000 Subject: [PATCH 181/278] Copy the holdout into the bundle without returning the result unchanged Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/trainer/trainer.py | 7 +++---- packages/kumo-engine/test/test_holdout_df.py | 2 +- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index d35f5f28d..5046109cb 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -242,12 +242,12 @@ def _find_holdout(log_dir: str) -> tuple[Optional[Path], Optional[bool]]: HOLDOUT_META_FILE = 'holdout-meta.json' -def _bundle_holdout(result: TrainingResult) -> TrainingResult: +def _bundle_holdout(result: TrainingResult) -> None: r"""Copy the holdout into the bundle.""" from kumo_core.artifacts import add_bundle_file if result.holdout_path is None: - return result + return try: shutil.copyfile(result.holdout_path, result.bundle_dir / HOLDOUT_FILE) add_bundle_file(result.bundle_dir, HOLDOUT_ROLE, HOLDOUT_FILE) @@ -261,7 +261,6 @@ def _bundle_holdout(result: TrainingResult) -> TrainingResult: f"Could not add the holdout to the model bundle: {error}; a " f"loaded copy of this model has no holdout_df().", UserWarning, stacklevel=4) - return result def _listed_holdout(bundle_dir: Path) -> tuple[Optional[Path], Optional[bool]]: @@ -694,7 +693,7 @@ def _finish_fit( training_table_generation=training_table_generation, graph_config=graph_config, ) - training_result = _bundle_holdout(training_result) + _bundle_holdout(training_result) try: training_result = replace( training_result, diff --git a/packages/kumo-engine/test/test_holdout_df.py b/packages/kumo-engine/test/test_holdout_df.py index 6f2435141..406d2a00b 100644 --- a/packages/kumo-engine/test/test_holdout_df.py +++ b/packages/kumo-engine/test/test_holdout_df.py @@ -142,7 +142,7 @@ def test_a_loaded_or_fetched_model_keeps_its_holdout(graph, tmp_path, def test_a_bundle_without_a_holdout_says_so(graph, monkeypatch): monkeypatch.setattr('kumo.engine.trainer.trainer._bundle_holdout', - lambda result: result) + lambda result: None) result = _fit(graph, QUERY) loaded = kumo.Trainer.load(result.bundle_dir).result From 3fe6442360d41ca1f7453efba6bf353c79dfa9c6 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 01:50:26 +0000 Subject: [PATCH 182/278] Drop a stray blank line before ObjectStoreSinkBackend Co-Authored-By: Claude Opus 5.5 --- packages/kumo-core/src/kumo_core/object_store_sink.py | 1 - 1 file changed, 1 deletion(-) diff --git a/packages/kumo-core/src/kumo_core/object_store_sink.py b/packages/kumo-core/src/kumo_core/object_store_sink.py index 41c4a861a..9770a031d 100644 --- a/packages/kumo-core/src/kumo_core/object_store_sink.py +++ b/packages/kumo-core/src/kumo_core/object_store_sink.py @@ -100,7 +100,6 @@ def file_sizes(filesystem: FileSystem, directory: str) -> dict[str, int]: } - class ObjectStoreSinkBackend(ABC): scheme: str provider: str From a7426456d523cb253e04d387e91f9a15df685437 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 01:50:33 +0000 Subject: [PATCH 183/278] Read a version's kind only from a SourceVersion Co-Authored-By: Claude Opus 5.5 --- packages/kumo-core/src/kumo_core/sources/identity.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/kumo-core/src/kumo_core/sources/identity.py b/packages/kumo-core/src/kumo_core/sources/identity.py index 2941f7ef1..4097a1e1a 100644 --- a/packages/kumo-core/src/kumo_core/sources/identity.py +++ b/packages/kumo-core/src/kumo_core/sources/identity.py @@ -36,7 +36,7 @@ def __new__(cls, value: str, kind: str = "unknown", def version_kind(version: object) -> str: - return getattr(version, "kind", "unknown") + return version.kind if isinstance(version, SourceVersion) else "unknown" def epoch_second(moment: datetime) -> str: From 55316b24b8b9a050a916b3c8f56bb992ddd21113 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 01:50:41 +0000 Subject: [PATCH 184/278] Check a stub once, by equality, without a no-op embedded() case Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/test/test_mlops.py | 23 +++-------------------- 1 file changed, 3 insertions(+), 20 deletions(-) diff --git a/packages/kumo-engine/test/test_mlops.py b/packages/kumo-engine/test/test_mlops.py index 6a740aec2..b3ec3827b 100644 --- a/packages/kumo-engine/test/test_mlops.py +++ b/packages/kumo-engine/test/test_mlops.py @@ -10,12 +10,7 @@ import pytest import kumo.engine as kumo -from kumo.engine._manifest import ( - SnapshotManifest, - SnapshotParents, - embedded, - stub, -) +from kumo.engine._manifest import SnapshotManifest, stub from kumo.engine._mlops import resolve @@ -99,17 +94,5 @@ async def task(): def test_stubs_stand_in_for_parents_without_a_manifest(): - snapshot = stub('snapshot', 'abc@0123') - assert isinstance(snapshot, SnapshotManifest) - assert snapshot.to_dict() == { - 'kind': 'snapshot', - 'schema_version': 1, - 'id': 'abc@0123', - 'created_at': None, - 'environment': None, - 'reconstructed': True, - 'parents': None, - } - full = SnapshotManifest(id='abc@0123', - parents=SnapshotParents(source_tables={})) - assert embedded(full) == full + assert stub('snapshot', 'abc@0123') == SnapshotManifest(id='abc@0123', + reconstructed=True) From 895fa00ad6a0f53ae030b15967289b87956b4eea Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 02:25:15 +0000 Subject: [PATCH 185/278] Return a PublishedModel from the warm-start test's fake sink, as real sinks do Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/test/test_warm_start.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/packages/kumo-engine/test/test_warm_start.py b/packages/kumo-engine/test/test_warm_start.py index 86d274770..6d37621f2 100644 --- a/packages/kumo-engine/test/test_warm_start.py +++ b/packages/kumo-engine/test/test_warm_start.py @@ -21,7 +21,7 @@ from kumo_core.train.trainer import get_best_ckpt_path from kumo_core.types import TaskType from kumoapi.model_plan import ModelPlan -from kumoapi.plugin import ModelFetchingSink +from kumoapi.plugin import ModelFetchingSink, PublishedModel from kumo.engine import FileConnector, Graph, PredictiveQuery, Table, Trainer @@ -111,8 +111,9 @@ def test_bundle_warm_start_fit_reload_and_predict(tmp_path, monkeypatch, refit, shutil.rmtree(initial.scratch_dir) if refit == 'full': sink = Mock(spec=ModelFetchingSink) - sink.fetch_model.side_effect = lambda connector, target, **kwargs: shutil.copytree( - durable, target) + sink.fetch_model.side_effect = lambda connector, target, **kwargs: ( + shutil.copytree(durable, target) and PublishedModel( + name='m', version='1', uri=str(durable))) source = Trainer.load_from_connector(SimpleNamespace(sink_backend=sink), path='retained-model') else: From a3592c041dfb56e46f049374d3f1decd7ac62128 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 02:44:11 +0000 Subject: [PATCH 186/278] Log a registered result's source tables from its manifest's Datasets Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/integrations/mlflow.py | 53 +++++++------------ .../kumo-engine/test/test_mlflow_model.py | 4 +- 2 files changed, 22 insertions(+), 35 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index 738decdf5..4efcfdb04 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -26,6 +26,7 @@ Literal, Mapping, Optional, + Sequence, Set, Tuple, ) @@ -275,9 +276,6 @@ def start(self, record: RunRecord) -> None: _SystemMetrics(self.log_metrics)) self._log_params(self._parent, record.params) self._log_artifacts(record.artifacts) - if record.tags.get('kumo.kind') != 'fit': - with self.target.scope(): - _log_inputs(self._mlflow, client, self._parent, record.inputs) def on_trial(self, record: RunRecord, event: Event) -> None: key = (event.trial or 0, int(event.params['attempt'])) @@ -388,9 +386,18 @@ def _log_table(self, summary: Table) -> None: json.dumps([*logged, entry])) def _log_dataset(self, summary: Dataset) -> None: + self._log_inputs(self._parent, [summary]) + + def _log_inputs(self, run_id: str, summaries: Sequence[Dataset]) -> None: + datasets = [d for d in map(self._dataset_input, summaries) if d] + if datasets: + with self.target.scope(): + self._client.log_inputs(run_id, datasets=datasets) + + def _dataset_input(self, summary: Dataset) -> Any: source = _dataset_source(summary.source) if source is None: - return + return None entities = self._mlflow.entities version = summary.source.get('version') or '' dataset = entities.Dataset( @@ -404,10 +411,8 @@ def _log_dataset(self, summary: Dataset) -> None: 'string')} for name, column in summary.schema.items()]}), profile=json.dumps(summary.profile, default=str)) - with self.target.scope(): - self._client.log_inputs(self._parent, datasets=[ - entities.DatasetInput(dataset, tags=[entities.InputTag( - 'mlflow.data.context', summary.context)])]) + return entities.DatasetInput(dataset, tags=[entities.InputTag( + 'mlflow.data.context', summary.context)]) def _log_artifacts(self, artifacts: Mapping[str, str]) -> None: for name, text in artifacts.items(): @@ -595,8 +600,8 @@ def log_model( none. target: as in `MlflowTracking`. """ + from kumo.engine._manifest import artifact_summaries from kumo.engine._mlops import snapshot_uri - from kumo.engine._tracking import fit_inputs from kumo.engine.trainer._provenance import bundle_manifest mlflow = import_provider('mlflow', 'mlflow') @@ -618,7 +623,11 @@ def log_model( run_id = client.create_run( tracking._experiment_for(client)).info.run_id experiment_id = client.get_run(run_id).info.experiment_id - _log_inputs(mlflow, client, run_id, fit_inputs(result.manifest)) + if result.manifest is not None: + tracking._log_inputs(run_id, [ + summary for summary in artifact_summaries(result.manifest, + 'training') + if isinstance(summary, Dataset)]) info = _log_model(target, client, result.bundle_dir, bundle_manifest(result.bundle_dir, result.id), snapshot, run_id, experiment_id, name, @@ -773,30 +782,6 @@ def _requirement(config: Any) -> str: return f"kumo-engine[{','.join(extras)}]=={__version__}" -def _log_inputs(mlflow: Any, client: Any, run_id: str, - inputs: Mapping[str, Mapping[str, Optional[str]]]) -> None: - try: - meta = import_provider('mlflow.data.meta_dataset', 'mlflow') - entities = mlflow.entities - datasets = [] - for name, entry in inputs.items(): - source = _dataset_source(entry) - if source is None: - continue - digest = hashlib.sha256( - json.dumps(dict(entry), sort_keys=True).encode()).hexdigest() - dataset = meta.MetaDataset(source, name=name, digest=digest[:8]) - datasets.append( - entities.DatasetInput( - dataset._to_mlflow_entity(), - tags=[entities.InputTag('mlflow.data.context', 'train')])) - if datasets: - client.log_inputs(run_id, datasets=datasets) - except Exception as error: - logger.warning('Could not log the fit\'s source tables to MLflow: %s', - format_error(error), exc_info=True) - - def _register(target: Any, client: Any, info: Any, location: ModelRegistryLocation, metadata: Mapping[str, str]) -> PublishedModel: diff --git a/packages/kumo-engine/test/test_mlflow_model.py b/packages/kumo-engine/test/test_mlflow_model.py index 64a926126..a093bf5d3 100644 --- a/packages/kumo-engine/test/test_mlflow_model.py +++ b/packages/kumo-engine/test/test_mlflow_model.py @@ -134,7 +134,9 @@ def test_logged_model_scores_the_query_population_and_timed_rows( info = log_model(trainer, snapshot=snapshot, serving_snapshot=str(tmp_path / 'exports' / location)) - assert mlflow.get_run(info.run_id).info.status == 'FINISHED' + run = mlflow.get_run(info.run_id) + assert run.info.status == 'FINISHED' + assert 'ents' in {item.dataset.name for item in run.inputs.dataset_inputs} requirements = mlflow.pyfunc.get_model_dependencies(info.model_uri) assert f'kumo-engine[gcs,s3]=={kumo.__version__}' in open( requirements).read().splitlines() From b47562e441c75c77f2b7ae646be952745b2acd33 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 02:44:11 +0000 Subject: [PATCH 187/278] Drop RunRecord.inputs: trackers get source tables as Dataset summaries Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/docs/mlops.md | 2 +- packages/kumo-engine/kumo/engine/_tracking.py | 36 ------------------ .../kumo/engine/trainer/_run_spec.py | 5 +-- .../kumo/engine/trainer/trainer.py | 8 +--- .../kumo-engine/test/test_mlflow_tracking.py | 37 ------------------- .../test/test_model_agnostic_mlops.py | 3 +- .../test/test_tracking_contract.py | 3 +- 7 files changed, 5 insertions(+), 89 deletions(-) diff --git a/packages/kumo-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index e209eeaed..6ce8efefb 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -291,7 +291,7 @@ A text column `infer_metadata()` reads as categorical (short values from a small ### Custom trackers -Other systems plug in through `kumo.engine.tracking.Tracking`, 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 attempt after each of its events, and `finish(record, error)` with the final record (`None` when tracking stopped early). Optional members: `prepare()`, `tags`, `check(registry)` (when `kumo.MLOps` is built), `log_model(bundle_dir, manifest, snapshot, registry)` (logs and registers, returning the version) or `connector` (a connector whose model registry `registry` registers in), `load_model(location)` and `register(result, location)` (for loading and `result.export(cfg)`), `platform`/`run_id`/`run_url` (recorded as `result.tracking`). A tracker that raises is not called again for that fit and its run ends failed. `RunRecord.inputs` maps each source table to its `connector`, `location`, `version` and `version_kind`; a tracker's `tags` are added except `kumo.*` and param names. +Other systems plug in through `kumo.engine.tracking.Tracking`, 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 attempt after each of its events, and `finish(record, error)` with the final record (`None` when tracking stopped early). Optional members: `prepare()`, `tags`, `check(registry)` (when `kumo.MLOps` is built), `log_model(bundle_dir, manifest, snapshot, registry)` (logs and registers, returning the version) or `connector` (a connector whose model registry `registry` registers in), `load_model(location)` and `register(result, location)` (for loading and `result.export(cfg)`), `platform`/`run_id`/`run_url` (recorded as `result.tracking`). A tracker that raises is not called again for that fit and its run ends failed. Source tables arrive in `finish` as `Dataset` summaries in `record.summaries`, each with its `name`, `source` (`connector`, `location`, `version`, `version_kind`), `schema`, `profile` and `context`; a tracker's `tags` are added except `kumo.*` and param names. ### Progress events diff --git a/packages/kumo-engine/kumo/engine/_tracking.py b/packages/kumo-engine/kumo/engine/_tracking.py index 926d223bf..654f297d0 100644 --- a/packages/kumo-engine/kumo/engine/_tracking.py +++ b/packages/kumo-engine/kumo/engine/_tracking.py @@ -57,8 +57,6 @@ class RunRecord: ...]] = field(default_factory=dict) tags: Mapping[str, str] = field(default_factory=dict) artifacts: Mapping[str, str] = field(default_factory=dict) - inputs: Mapping[str, Mapping[str, Optional[str]]] = field( - default_factory=dict) summaries: Tuple[Summary, ...] = () @@ -73,8 +71,6 @@ class RunSpec: steps: Mapping[str, Sequence[Tuple[int, float]]] = field( default_factory=dict) tags: Mapping[str, Any] = field(default_factory=dict) - inputs: Mapping[str, Mapping[str, Optional[str]]] = field( - default_factory=dict) parents: Mapping[str, Optional[str]] = field(default_factory=dict) artifacts: Mapping[str, str] = field(default_factory=dict) summaries: Sequence[Summary] = () @@ -82,37 +78,6 @@ class RunSpec: error: Optional[str] = None -INPUT_FIELDS = ('connector', 'location', 'version', 'version_kind') - - -def fit_inputs(manifest: Any) -> Dict[str, Dict[str, Optional[str]]]: - try: - table = manifest - if getattr(manifest, 'kind', None) == 'model': - table = getattr(_parents(manifest), 'training_table', None) - parents = _parents(table) - snapshot = _parents(getattr(parents, 'snapshot', None)) - tables = dict(getattr(snapshot, 'source_tables', None) or {}) - source = getattr(parents, 'source_table', None) - if source is not None: - tables['training_table'] = source - return { - str(name): {key: None if getattr(entry, key) is None else str( - getattr(entry, key)) for key in INPUT_FIELDS} - for name, entry in tables.items() - } - except Exception as error: - logger.warning('Could not read the fit\'s source tables: %s', - format_error(error), exc_info=_debug()) - return {} - - -def _parents(manifest: Any) -> Any: - if manifest is None or getattr(manifest, 'reconstructed', False): - return None - return getattr(manifest, 'parents', None) - - def logged_name(name: str) -> str: r"""``val_auroc`` as ``val/auroc``.""" best = name.startswith('best_') @@ -152,7 +117,6 @@ def run_record(spec: RunSpec) -> RunRecord: **({} if spec.error is None else {'kumo.error': spec.error}), }, artifacts=dict(spec.artifacts), - inputs=dict(spec.inputs), summaries=tuple(spec.summaries), ) diff --git a/packages/kumo-engine/kumo/engine/trainer/_run_spec.py b/packages/kumo-engine/kumo/engine/trainer/_run_spec.py index 4ac3ad557..e33bb8212 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_run_spec.py +++ b/packages/kumo-engine/kumo/engine/trainer/_run_spec.py @@ -4,7 +4,7 @@ r"""A `Trainer.fit` as a tracked run's `RunSpec`.""" import json -from dataclasses import dataclass, field, replace +from dataclasses import dataclass, replace from datetime import datetime, timezone from pathlib import Path from typing import Any, Iterator, Mapping, Optional, Sequence, Tuple @@ -22,8 +22,6 @@ class FitContext: query: str model_plan: Mapping[str, Any] - inputs: Mapping[str, Mapping[str, Optional[str]]] = field( - default_factory=dict) started: Optional[datetime] = None snapshot_id: Optional[str] = None task_type: Optional[str] = None @@ -42,7 +40,6 @@ def fit_spec(context: FitContext) -> RunSpec: 'tune_metric': (context.model_plan.get('training_job') or {}).get('tune_metric'), 'engine_version': __version__}, - inputs=dict(context.inputs), artifacts={ 'query.pql': context.query, 'model_plan.json': json.dumps(dict(context.model_plan), indent=2, diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 8f3f8a9c8..625437941 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -83,12 +83,7 @@ from kumo.engine._mlops import full_error, warn_failure from kumo.engine._predictor import TRAINER_PREDICTOR from kumo.engine._scratch import new_run_id, trainer_scratch_dir -from kumo.engine._tracking import ( - RunSpec, - TrackingSession, - as_tracking, - fit_inputs, -) +from kumo.engine._tracking import RunSpec, TrackingSession, as_tracking from kumo.engine.connector.base import Connector from kumo.engine.connector.source_table import SourceTable from kumo.engine.graph import GraphSnapshot, report_other_graph @@ -931,7 +926,6 @@ def _fit( tracking.start(lambda: fit_spec(FitContext( query=train_table.query, model_plan=to_json_dict(materialized.model_plan), - inputs=fit_inputs(train_table.manifest), started=started_at, snapshot_id=snapshot.id, task_type=materialized.task.task_type.value, diff --git a/packages/kumo-engine/test/test_mlflow_tracking.py b/packages/kumo-engine/test/test_mlflow_tracking.py index 54beec57f..b804557b9 100644 --- a/packages/kumo-engine/test/test_mlflow_tracking.py +++ b/packages/kumo-engine/test/test_mlflow_tracking.py @@ -395,43 +395,6 @@ def test_source_tables_are_logged_as_native_inputs( assert mlflow.data.get_source(orders).delta_table_id == 'table-id' -def test_fit_inputs_come_from_the_embedded_parents(): - from kumo.engine._manifest import ( - ModelManifest, - ModelParents, - SnapshotManifest, - SnapshotParents, - SourceTableEntry, - TrainingTableManifest, - TrainingTableParents, - stub, - ) - from kumo.engine._tracking import fit_inputs - - entry = {'connector': 'LOCAL_FILE', 'location': '/data/ents.parquet', - 'version': '1790000000', 'version_kind': 'timestamp'} - snapshot = SnapshotManifest(id='s', parents=SnapshotParents( - source_tables={'ents': SourceTableEntry(**entry, sql='SELECT 1')})) - custom = SourceTableEntry(**{**entry, 'version': None}) - table = TrainingTableManifest(id='t', parents=TrainingTableParents( - snapshot=snapshot, source_table=custom)) - model = ModelManifest(id='m', parents=ModelParents( - training_table=table, warm_start=ModelManifest( - id='w', parents=ModelParents(training_table=TrainingTableManifest( - id='u', parents=TrainingTableParents(snapshot=SnapshotManifest( - id='v', parents=SnapshotParents(source_tables={ - 'old': SourceTableEntry(**entry)})))))))) - - expected = {'ents': entry, 'training_table': {**entry, 'version': None}} - assert fit_inputs(table) == expected - assert fit_inputs(model) == expected - assert fit_inputs(TrainingTableManifest(id='t', parents=TrainingTableParents( - snapshot=stub('snapshot', 's'), source_table=custom))) == { - 'training_table': {**entry, 'version': None}} - assert fit_inputs(stub('model', 'm')) == {} - assert fit_inputs(None) == {} - - def test_link_prediction_metric_names_are_logged(tracking_uri): from kumo.engine._tracking import RunRecord from kumo.engine.integrations.mlflow import MlflowTracking diff --git a/packages/kumo-engine/test/test_model_agnostic_mlops.py b/packages/kumo-engine/test/test_model_agnostic_mlops.py index ab8d70d57..ef94de3b8 100644 --- a/packages/kumo-engine/test/test_model_agnostic_mlops.py +++ b/packages/kumo-engine/test/test_model_agnostic_mlops.py @@ -86,8 +86,7 @@ def finish(self, record, error): def test_any_run_is_tracked_from_its_spec(): tracker = _Tracker() session = TrackingSession(tracker) - spec = RunSpec(kind='custom', params={'lr': 0.1}, - inputs={'users': {'connector': 'S3'}}) + spec = RunSpec(kind='custom', params={'lr': 0.1}) session.start(lambda: spec) session.on_event(Event('trial_started', 0.0, trial=1, params={'attempt': 1})) diff --git a/packages/kumo-engine/test/test_tracking_contract.py b/packages/kumo-engine/test/test_tracking_contract.py index fed46a8c9..0e46bfc25 100644 --- a/packages/kumo-engine/test/test_tracking_contract.py +++ b/packages/kumo-engine/test/test_tracking_contract.py @@ -40,7 +40,7 @@ } SNAPSHOT = f"{'s' * 64}@{'v' * 16}" CONTEXT = FitContext( - query=QUERY, inputs=INPUTS, snapshot_id=SNAPSHOT, task_type='regression', + query=QUERY, snapshot_id=SNAPSHOT, task_type='regression', started=datetime(2026, 10, 29, 6, 18, tzinfo=timezone.utc), model_plan={ 'optimization': {'base_lr': [0.01, 0.001], 'max_epochs': 2}, @@ -361,7 +361,6 @@ def test_records_of_the_recorded_fit(): '3', 'kumo.model_alias': 'candidate'}.items() assert fit.artifacts['query.pql'] == QUERY assert json.loads(fit.artifacts['model_plan.json']) == CONTEXT.model_plan - assert fit.inputs == INPUTS oom = records['trial-2-attempt-1'] assert oom == RunRecord( From 9cf57c6703024c6a04f94412f2cbec2ea106c3f9 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 02:44:55 +0000 Subject: [PATCH 188/278] Keep a registered result's source-table logging best-effort Co-Authored-By: Claude Opus 5.5 --- .../kumo-engine/kumo/engine/integrations/mlflow.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index 4efcfdb04..a5bc50620 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -623,11 +623,14 @@ def log_model( run_id = client.create_run( tracking._experiment_for(client)).info.run_id experiment_id = client.get_run(run_id).info.experiment_id - if result.manifest is not None: - tracking._log_inputs(run_id, [ - summary for summary in artifact_summaries(result.manifest, - 'training') - if isinstance(summary, Dataset)]) + summaries = [] if result.manifest is None else artifact_summaries( + result.manifest, 'training') + try: + tracking._log_inputs( + run_id, [s for s in summaries if isinstance(s, Dataset)]) + except Exception as error: + logger.warning('Could not log the fit\'s source tables to MLflow: %s', + format_error(error), exc_info=True) info = _log_model(target, client, result.bundle_dir, bundle_manifest(result.bundle_dir, result.id), snapshot, run_id, experiment_id, name, From da0f66fa27802b26f36bcf717bd5eaf14c1b5f92 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 02:51:24 +0000 Subject: [PATCH 189/278] Name the published model in the holdout test parameters The refit_full case now also explains, as in the upper PRs. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/test/test_holdout_df.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/packages/kumo-engine/test/test_holdout_df.py b/packages/kumo-engine/test/test_holdout_df.py index 70bfe82c6..b207a0f86 100644 --- a/packages/kumo-engine/test/test_holdout_df.py +++ b/packages/kumo-engine/test/test_holdout_df.py @@ -95,18 +95,19 @@ def test_holdout_df_after_fit(graph): assert not (result.bundle_dir / 'test_df.parquet').exists() -@pytest.mark.parametrize('refit_trainval,refit_full,explain,suffix', [ - (False, False, False, ''), - (True, False, True, '_refit_trainval'), - (True, True, False, None), +@pytest.mark.parametrize('refit_trainval,refit_full,explain,model', [ + (False, False, False, 'fit'), + (True, False, True, 'refit_trainval'), + (True, True, True, None), ]) def test_holdout_is_the_published_models(graph, refit_trainval, refit_full, - explain, suffix): + explain, model): result = _fit(graph, QUERY, refit_trainval, refit_full, explain) - if suffix is None: + if model is None: assert result.holdout_path is None else: + suffix = '' if model == 'fit' else '_refit_trainval' assert re.fullmatch(rf'version_\d+{suffix}', result.holdout_path.parent.name) assert len(result.holdout_df()) > 0 From 9d82a62a7aa1b175dac0e6d75312fa31267a4e53 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 02:51:24 +0000 Subject: [PATCH 190/278] Wrap the carry-forward fake training result like the formatter does Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/test/test_semantic_carry_forward.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/packages/kumo-engine/test/test_semantic_carry_forward.py b/packages/kumo-engine/test/test_semantic_carry_forward.py index a1533b131..9481d14d1 100644 --- a/packages/kumo-engine/test/test_semantic_carry_forward.py +++ b/packages/kumo-engine/test/test_semantic_carry_forward.py @@ -874,10 +874,9 @@ def fake_train(training_input, **kwargs): both_workers.wait(timeout=5) time.sleep(.05) run_name = Path(training_input.home_dir).name - return SimpleNamespace( - metrics={'run_name': - run_name}, artifact_paths={}, explain_disabled=True, - best_trainer=_best_trainer(tmp_path)) + return SimpleNamespace(metrics={'run_name': run_name}, + artifact_paths={}, explain_disabled=True, + best_trainer=_best_trainer(tmp_path)) monkeypatch.setattr( 'kumo.engine.trainer.trainer.build_generated_training_task_dataset', From b375ee5cad5ce9c665f8c546cee25e901452c5e7 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 02:51:33 +0000 Subject: [PATCH 191/278] Move metric_key into a shared _text module with the query's entity Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_text.py | 17 +++++++++++++++++ .../kumo/engine/integrations/mlflow.py | 6 +----- 2 files changed, 18 insertions(+), 5 deletions(-) create mode 100644 packages/kumo-engine/kumo/engine/_text.py diff --git a/packages/kumo-engine/kumo/engine/_text.py b/packages/kumo-engine/kumo/engine/_text.py new file mode 100644 index 000000000..a660aebe2 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/_text.py @@ -0,0 +1,17 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""Text helpers shared by summaries and cards.""" + +import re +from typing import Optional + + +def entity(query: Optional[str]) -> Optional[str]: + r"""The entity table a predictive query predicts for.""" + found = re.search(r'FOR\s+EACH\s+(\w+)', query or '', re.IGNORECASE) + return found and found.group(1) + + +def metric_key(name: str) -> str: + return name.replace('@', '_at_') diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index a5bc50620..24b145f3e 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -38,6 +38,7 @@ from kumo.engine._events import Event from kumo.engine._mlops import full_error +from kumo.engine._text import metric_key from kumo.engine._tracking import ( DROPPED_PARAMS, RegistrationError, @@ -578,11 +579,6 @@ def _sagemaker_ui_hint(tracking_uri: str) -> Optional[str]: f'from SageMaker Studio or at the URL `{command}` prints.') -def metric_key(name: str) -> str: - r"""``map@10`` as ``map_at_10``: MLflow rejects ``@``.""" - return name.replace('@', '_at_') - - def log_model( trainer: 'Trainer', *, From 3e22e0374feb4481e475c03b3ee07d242b4df319 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 02:55:58 +0000 Subject: [PATCH 192/278] Refuse occupied output directories on every object store, by full location The preflight moves from FilesystemSink to ObjectStoreSinkBackend, so S3, GCS and ABFS outputs are refused before scoring too, and a conflict found while writing names the same full location. Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_core/object_store_sink.py | 49 ++++++++++++------- .../kumo-engine/test/test_filestore_sink.py | 22 +++++++-- 2 files changed, 48 insertions(+), 23 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/object_store_sink.py b/packages/kumo-core/src/kumo_core/object_store_sink.py index d09484b4c..d7433cf3d 100644 --- a/packages/kumo-core/src/kumo_core/object_store_sink.py +++ b/packages/kumo-core/src/kumo_core/object_store_sink.py @@ -123,17 +123,42 @@ def write_files(self, connector, table_name, files, if write_mode is not WriteMode.OVERWRITE: raise NotImplementedError( 'Object-store artifact output does not support APPEND') - from kumo_core.artifact_export import export_selected_files + from kumo_core.artifact_export import ( + ArtifactDestinationConflict, + export_selected_files, + ) from kumo_core.sources import layout filesystem, path = self._destination(connector, table_name) kind = layout.file_type(path.rstrip('/')) or layout.PARQUET storage = _ArrowExportFiles(filesystem, extra_columns) - _, output = export_selected_files(files, path, fs=storage, - identity=export_id or str(files), - file_type=kind, - convert_csv=storage.convert_csv) + try: + _, output = export_selected_files( + files, path, fs=storage, identity=export_id or str(files), + file_type=kind, convert_csv=storage.convert_csv) + except ArtifactDestinationConflict: + raise self._conflict(path) from None return self._uri(output) + def preflight_write_files(self, connector, table_name, write_mode=None): + """Refuse an occupied output directory before scoring.""" + del write_mode + from pyarrow.fs import FileSelector + + if not table_name.endswith('/'): + return + filesystem, path = self._destination(connector, table_name) + if filesystem.get_file_info( + FileSelector(path.rstrip('/'), allow_not_found=True)): + raise self._conflict(path) + + def _conflict(self, path: str) -> ValueError: + from kumo_core.artifact_export import ArtifactDestinationConflict + + return ArtifactDestinationConflict( + f"Output directory '{self._uri(path.rstrip('/'))}/' already " + f"contains data from another export. Choose an empty destination " + "path.") + def publish_model( self, connector: SourceConnector, @@ -406,20 +431,6 @@ def _destination(self, connector, coordinate): def _uri(self, path: str) -> str: return path - def preflight_write_files(self, connector, table_name, write_mode=None): - """Refuse an occupied output directory before scoring.""" - del write_mode - from kumo_core.artifact_export import ArtifactDestinationConflict - from pyarrow.fs import FileSelector - - if not table_name.endswith('/'): - return - filesystem, path = self._sink.filesystem(connector, table_name) - if filesystem.get_file_info(FileSelector(path, allow_not_found=True)): - raise ArtifactDestinationConflict( - f"Output directory '{path}/' already contains data from " - f"another export. Choose an empty destination path.") - class FileStoreRouting: """Sends ``holds_files`` paths to :class:`FilesystemSink`, the rest on.""" diff --git a/packages/kumo-engine/test/test_filestore_sink.py b/packages/kumo-engine/test/test_filestore_sink.py index da12f3bb0..f8dd18e95 100644 --- a/packages/kumo-engine/test/test_filestore_sink.py +++ b/packages/kumo-engine/test/test_filestore_sink.py @@ -8,9 +8,10 @@ import pandas as pd import pytest +from kumo_core.artifact_export import ArtifactDestinationConflict from kumo_core.chunked_fs import ChunkedFileHandler from kumo_core.object_store_sink import FileStoreRouting -from kumoapi.jobs import PredictionArtifactType +from kumoapi.jobs import PredictionArtifactType, WriteMode from kumoapi.plugin import FileStoreSink from pyarrow.fs import PyFileSystem @@ -173,15 +174,28 @@ def test_parquet_writes_stay_plain_up_to_one_part(connector): filesystem=PyFileSystem(store)).shape == (100, 1) -def test_an_occupied_directory_is_refused_before_scoring(connector): - from kumo_core.artifact_export import ArtifactDestinationConflict +def test_an_occupied_directory_is_refused_by_its_full_location( + tmp_path, connector): + files = connector.sink_backend + part = tmp_path / 'part-0.parquet' + pd.DataFrame({'x': range(3)}).to_parquet(part) + files.write_files(connector, '@stage/out/', [str(part)]) + + files.preflight_write_files(connector, '@stage/new/') + with pytest.raises(ArtifactDestinationConflict, match="'@stage/out/'"): + files.preflight_write_files(connector, '@stage/out/') + +def test_an_occupied_directory_is_refused_before_scoring(connector): sink = connector.sink_backend sink.preflight_write_files(connector, '@stage/out/') sink.store.objects['@stage/out/old.parquet'] = b'x' - for mode in (None, 'OVERWRITE', 'APPEND'): + for mode in (None, WriteMode.APPEND): with pytest.raises(ArtifactDestinationConflict, match='already'): sink.preflight_write_files(connector, '@stage/out/', mode) + with pytest.raises(ArtifactDestinationConflict, match='already'): + sink.preflight_write_files(connector, '@stage/out/', + WriteMode.OVERWRITE) def test_a_file_only_sink_refuses_a_table(connector, tmp_path): From d5f46e4adf7db660a7775346b3f9ba79f75be040 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 02:55:58 +0000 Subject: [PATCH 193/278] Extract the model publisher check and the output preflight from fit and predict Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/trainer/trainer.py | 31 ++++++++++++------- 1 file changed, 19 insertions(+), 12 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 8c3a604a4..a208a3677 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -213,6 +213,22 @@ def _to_core_model_plan(model_plan: ModelPlan) -> CoreModelPlan: return internal_model_plan_from_public(model_plan) +def _require_publisher(export_config: ModelOutputConfig) -> None: + if not isinstance(export_config.connector.sink_backend, + ModelPublishingSink): + raise NotImplementedError( + f'{type(export_config.connector).__name__} cannot publish ' + f'models') + + +def _preflight(output_config: PredictionOutputConfig) -> None: + preflight = getattr(output_config.writer(), 'preflight_write_files', None) + if callable(preflight): + for destination in output_config.destinations().values(): + preflight(output_config.connector, destination, + write_mode=output_config.write_mode) + + def _stage_model_bundle(artifacts: Mapping[str, str], dest: Path) -> Path: from kumo_core.artifacts import stage_model_bundle return stage_model_bundle(artifacts, dest) @@ -354,12 +370,8 @@ def fit( from kumo.engine.trainer._runtime import local_training_bindings - if export_config is not None and not isinstance( - export_config.connector.sink_backend, - ModelPublishingSink): - raise NotImplementedError( - f'{type(export_config.connector).__name__} cannot publish ' - f'models') + if export_config is not None: + _require_publisher(export_config) run_id = f'{self._id}-fit-{new_run_id()}' run_dir = self.scratch_dir / 'runs' / run_id @@ -649,12 +661,7 @@ def predict( ) if output_config is not None: - preflight = getattr(output_config.writer(), - 'preflight_write_files', None) - if callable(preflight): - for destination in destinations.values(): - preflight(output_config.connector, destination, - write_mode=output_config.write_mode) + _preflight(output_config) from kumo_core.encoders.config import EncodersConfig from kumo_core.io.json import load_json From edbb28d411a33b89acaf8e5211531e0afad36985 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 03:00:18 +0000 Subject: [PATCH 194/278] Name parent_model bundle_manifest, as it reads any bundle's manifest Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/trainer/_provenance.py | 4 ++-- packages/kumo-engine/kumo/engine/trainer/trainer.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/_provenance.py b/packages/kumo-engine/kumo/engine/trainer/_provenance.py index b63bb3ac8..df2c8ba18 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_provenance.py +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -67,8 +67,8 @@ def embedded_model(manifest: ModelManifest) -> ModelManifest: parents, warm_start=stub('model', parents.warm_start.id))) -def parent_model(bundle_dir: Path, bundle_id: Optional[str] - ) -> Optional[ModelManifest]: +def bundle_manifest(bundle_dir: Path, + bundle_id: Optional[str]) -> Optional[ModelManifest]: r"""The manifest of the bundle at ``bundle_dir``, or a stub.""" manifest = read_provenance(bundle_dir) if manifest is not None: diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index a3f4d2230..b90bfaa30 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -65,7 +65,7 @@ from kumo.engine.pquery.training_table import TrainingTable from kumo.engine.trainer._provenance import ( FitRecord, - parent_model, + bundle_manifest, read_provenance, ) from kumo.engine.trainer._spec import ( @@ -673,7 +673,7 @@ def fit( model_plan=effective_plan, provenance=FitRecord( train_table=train_table, - warm_start=None if warm_start is None else parent_model( + warm_start=None if warm_start is None else bundle_manifest( warm_start._training_result.bundle_dir, warm_start_bundle_id), started_at=started_at, fit_mode=fit_mode), From 1f4f81f5bd7426a1e0f5477348bd79a288d53b5d Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 03:00:34 +0000 Subject: [PATCH 195/278] Cut an embedded model's warm-start parent through a small helper Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/trainer/_provenance.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/_provenance.py b/packages/kumo-engine/kumo/engine/trainer/_provenance.py index df2c8ba18..5d8b0fbdb 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_provenance.py +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -61,10 +61,14 @@ def embedded_model(manifest: ModelManifest) -> ModelManifest: r"""``manifest`` for embedding, its own model parents cut to stubs.""" copy = embedded(manifest) parents = copy.parents - if parents is None or parents.warm_start is None: + if parents is None: return copy + + def cut(model: Optional[ModelManifest]) -> Optional[ModelManifest]: + return None if model is None else stub('model', model.id) + return replace(copy, parents=replace( - parents, warm_start=stub('model', parents.warm_start.id))) + parents, warm_start=cut(parents.warm_start))) def bundle_manifest(bundle_dir: Path, From 2a4bf2545ec59d827985bdc0832cdabac028c96b Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 03:00:34 +0000 Subject: [PATCH 196/278] List the roles kept out of the artifact paths one per line Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/trainer/_spec.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/_spec.py b/packages/kumo-engine/kumo/engine/trainer/_spec.py index 8d8e04897..d81bfb2e6 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_spec.py +++ b/packages/kumo-engine/kumo/engine/trainer/_spec.py @@ -144,6 +144,10 @@ def read_bundle_artifact_paths( return { entry['role']: str(bundle_dir / entry['relative_path']) for entry in artifacts - if entry['role'] not in (SPEC_ROLE, PROVENANCE_ROLE, HOLDOUT_ROLE, - HOLDOUT_META_ROLE) + if entry['role'] not in ( + SPEC_ROLE, + PROVENANCE_ROLE, + HOLDOUT_ROLE, + HOLDOUT_META_ROLE, + ) }, bundle_id From d5997cafbf6623d337cb316b023496775ba72f73 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 03:01:48 +0000 Subject: [PATCH 197/278] Name runs the same on every tracker through _run_names Fits are fit--, trials trial--, a user run_name names the fit fit-, and Snowflake numbers a taken name _2, _3. Trackers report the experiment and run name on TrackingRef. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_manifest.py | 2 + .../kumo-engine/kumo/engine/_run_names.py | 77 +++++++++++++++++++ packages/kumo-engine/kumo/engine/_tracking.py | 35 +++++++-- .../kumo/engine/integrations/mlflow.py | 36 ++++++--- .../kumo/engine/integrations/snowflake.py | 63 ++++++--------- .../kumo/engine/trainer/_run_spec.py | 7 +- .../test/_snowflake_experiment_fake.py | 11 ++- .../kumo-engine/test/test_mlflow_tracking.py | 22 ++++-- .../test/test_snowflake_tracking.py | 46 +++++++---- .../test/test_tracking_contract.py | 26 ++++--- 10 files changed, 229 insertions(+), 96 deletions(-) create mode 100644 packages/kumo-engine/kumo/engine/_run_names.py diff --git a/packages/kumo-engine/kumo/engine/_manifest.py b/packages/kumo-engine/kumo/engine/_manifest.py index 6b615ed77..4e513829e 100644 --- a/packages/kumo-engine/kumo/engine/_manifest.py +++ b/packages/kumo-engine/kumo/engine/_manifest.py @@ -343,6 +343,8 @@ class TrackingRef(Record): platform: str run_id: str url: Optional[str] = omit_none() + experiment: Optional[str] = omit_none() + run_name: Optional[str] = omit_none() @dataclass(frozen=True, kw_only=True) diff --git a/packages/kumo-engine/kumo/engine/_run_names.py b/packages/kumo-engine/kumo/engine/_run_names.py new file mode 100644 index 000000000..875a36d5c --- /dev/null +++ b/packages/kumo-engine/kumo/engine/_run_names.py @@ -0,0 +1,77 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""Run names, the same on every tracker.""" + +import re +from datetime import datetime, timezone +from typing import TYPE_CHECKING, Optional + +from kumo.engine._text import entity + +if TYPE_CHECKING: + from kumo.engine._tracking import RunRecord + +_KIND_WORDS = {'prediction': 'predict'} + + +def model_label(registered: Optional[str], query: Optional[str], + task_type: Optional[str]) -> str: + r"""The registered name's last part, else -.""" + if registered: + return registered.rsplit('.', 1)[-1] + return '-'.join(filter(None, (entity(query), task_type))) or 'model' + + +def run_name(kind: str, core: str, suffix: Optional[str] = None) -> str: + r"""-, the kind not repeated.""" + word = _KIND_WORDS.get(kind, kind) + named = re.match(rf'{word}([-_]|$)', core, re.IGNORECASE) + return (core if named else f'{word}-{core}') + (suffix or '') + + +def versioned(model: str, version: str) -> str: + r"""-v, or - for a named version.""" + return f'{model}-v{version}' if version.isdigit() else f'{model}-{version}' + + +def scored_model(model: str, version: Optional[str], bundle_id: str) -> str: + r"""-v, else -.""" + return versioned(model, version) if version else f'{model}-{bundle_id[:8]}' + + +def fit_name(model: str, started: datetime) -> str: + started = started.astimezone(timezone.utc) + return run_name('fit', f'{model}-{started:%Y%m%d-%H%M%S}') + + +def is_dated_fit(name: str) -> bool: + r"""Whether a fit run is still named by its start, not a version.""" + return re.search(r'-\d{8}-\d{6}$', name) is not None + + +def registered_fit_name(model_name: str, version: str) -> str: + return run_name('fit', versioned(model_label(model_name, None, None), + version)) + + +def trial_name(trial: int, attempt: int, fit: str) -> str: + return f'trial-{trial}-{fit}' + (f'-attempt-{attempt}' if attempt > 1 + else '') + + +def record_name(record: 'RunRecord', named: Optional[str] = None) -> str: + r"""The run name of ``record``; ``named`` replaces -v.""" + tags = record.tags + kind = tags.get('kumo.kind') + model = record.model_label or 'model' + if kind == 'fit': + return run_name('fit', named) if named else fit_name( + model, record.started or datetime.now(timezone.utc)) + suffix = record.name_suffix or ( + f"-{tags['kumo.as_of']}" if tags.get('kumo.as_of') else '') + word = (kind if kind != 'monitoring' else + 'eval' if tags.get('kumo.monitor') == 'evaluate' else 'drift') + core = named or scored_model(model, tags.get('kumo.model_version'), + tags.get('kumo.bundle_id', '')) + return run_name(word, core, suffix) diff --git a/packages/kumo-engine/kumo/engine/_tracking.py b/packages/kumo-engine/kumo/engine/_tracking.py index 654f297d0..14bc81ee5 100644 --- a/packages/kumo-engine/kumo/engine/_tracking.py +++ b/packages/kumo-engine/kumo/engine/_tracking.py @@ -12,6 +12,7 @@ import time import warnings from dataclasses import dataclass, field +from datetime import datetime from pathlib import Path from typing import ( TYPE_CHECKING, @@ -58,6 +59,9 @@ class RunRecord: tags: Mapping[str, str] = field(default_factory=dict) artifacts: Mapping[str, str] = field(default_factory=dict) summaries: Tuple[Summary, ...] = () + model_label: Optional[str] = None + started: Optional[datetime] = None + name_suffix: Optional[str] = None @dataclass(frozen=True) @@ -76,6 +80,9 @@ class RunSpec: summaries: Sequence[Summary] = () status: Optional[str] = None error: Optional[str] = None + model_label: Optional[str] = None + started: Optional[datetime] = None + name_suffix: Optional[str] = None def logged_name(name: str) -> str: @@ -118,6 +125,9 @@ def run_record(spec: RunSpec) -> RunRecord: }, artifacts=dict(spec.artifacts), summaries=tuple(spec.summaries), + model_label=spec.model_label, + started=spec.started, + name_suffix=spec.name_suffix, ) @@ -184,6 +194,8 @@ class TrackingBase(Tracking): log_model: Optional[Callable[..., Optional['PublishedModel']]] = None register: Optional[Callable[..., 'PublishedModel']] = None load_model: Optional[Callable[..., Any]] = None + experiment: Optional[str] = None + logged_run_name: Optional[str] = None @property def platform(self) -> str: @@ -342,16 +354,10 @@ def begin(tracker: TrackingBase) -> None: self._started.set() def pointer(self) -> Optional['TrackingRef']: - r"""The tracker's own run, once it started.""" - from kumo.engine._manifest import TrackingRef - + r"""The tracker's own run, once it started (`tracking_ref`).""" if self.spec is None or not self._started.wait(FINISH_SECONDS): return None - tracker = self._tracker - if tracker is None or not tracker.run_id: - return None - return TrackingRef(platform=str(tracker.platform), - run_id=str(tracker.run_id), url=tracker.run_url) + return tracking_ref(self._tracker) def on_event(self, event: Event) -> None: if self.spec is None or event.trial is None: @@ -520,6 +526,19 @@ def _fail(self, error: BaseException) -> None: self.error = error +def tracking_ref(tracking: Any) -> Optional['TrackingRef']: + r"""The tracker's latest run.""" + from kumo.engine._manifest import TrackingRef + + tracker = as_tracking(tracking) + if tracker is None or not tracker.run_id: + return None + return TrackingRef( + platform=str(tracker.platform), run_id=str(tracker.run_id), + url=tracker.run_url, experiment=tracker.experiment, + run_name=tracker.logged_run_name) + + def _debug() -> bool: return logger.isEnabledFor(logging.DEBUG) diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index 24b145f3e..b379e6c6e 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -29,6 +29,7 @@ Sequence, Set, Tuple, + cast, ) from urllib.parse import quote @@ -38,6 +39,7 @@ from kumo.engine._events import Event from kumo.engine._mlops import full_error +from kumo.engine._run_names import record_name, trial_name from kumo.engine._text import metric_key from kumo.engine._tracking import ( DROPPED_PARAMS, @@ -111,12 +113,14 @@ def __init__( self.tags = dict(tags or {}) self.logs_model = log_model self._outer: Optional[str] = None + self.logged_run_name: Optional[str] = None + self.experiment: Optional[str] = None + self._experiment_id: Optional[str] = None self._reset() def _reset(self) -> None: self._client: Any = None self._parent: Optional[str] = None - self._experiment_id: Optional[str] = None self._children: Dict[Tuple[int, int], str] = {} self._artifacts: Set[str] = set() self._model_id: Optional[str] = None @@ -261,16 +265,27 @@ def start(self, record: RunRecord) -> None: context = import_provider('mlflow.tracking.context.registry', 'mlflow') with self.target.scope(): - parent = {} if self._outer is None else { - 'mlflow.parentRunId': self._outer - } + experiment = (client.get_experiment_by_name(self.experiment_name) + if self.experiment_name is not None else + client.get_experiment(self._default_experiment + or '0')) + if experiment is None: + experiment = client.get_experiment( + client.create_experiment(self.experiment_name)) + experiment_id, self.experiment = (experiment.experiment_id, + experiment.name) + tags = dict(record.tags) + if self._outer is not None: + tags['mlflow.parentRunId'] = self._outer + kind = record.tags.get('kumo.kind') + name = (record_name(record, self.run_name) + if self.run_name and kind in ('fit', 'prediction') + else record.name) run = client.create_run( - self._experiment_for(client), - run_name=(self.run_name if self.run_name - and record.tags.get('kumo.kind') == 'fit' else - record.name), - tags=context.resolve_tags({**record.tags, **parent})) + experiment_id, run_name=name, + tags=context.resolve_tags(tags)) self._parent = self.run_id = run.info.run_id + self.logged_run_name = run.info.run_name self._experiment_id = run.info.experiment_id if self.log_system_metrics and record.tags.get('kumo.kind') == 'fit': self._guard.enter_context( @@ -283,7 +298,8 @@ def on_trial(self, record: RunRecord, event: Event) -> None: if event.name == 'trial_started': with self.target.scope(): run = self._client.create_run( - self._experiment_id, run_name=record.name, + self._experiment_id, run_name=trial_name( + *key, cast(str, self.logged_run_name)), tags={**record.tags, 'mlflow.parentRunId': self._parent}) self._children[key] = run.info.run_id self._log_params(run.info.run_id, record.params) diff --git a/packages/kumo-engine/kumo/engine/integrations/snowflake.py b/packages/kumo-engine/kumo/engine/integrations/snowflake.py index 04b2ce151..b02a7f4fe 100644 --- a/packages/kumo-engine/kumo/engine/integrations/snowflake.py +++ b/packages/kumo-engine/kumo/engine/integrations/snowflake.py @@ -7,7 +7,6 @@ import posixpath import re import tempfile -import uuid from contextlib import ExitStack from pathlib import Path from typing import Any, Dict, List, Mapping, Optional, Set, Tuple @@ -17,6 +16,7 @@ from kumo_core.summaries import Scalar, Table, Tag from kumo.engine._events import Event +from kumo.engine._run_names import record_name from kumo.engine._tracking import ( DROPPED_METRICS, DROPPED_PARAMS, @@ -72,6 +72,7 @@ def __init__( self.run_name = run_name self.tags = dict(tags or {}) self.native_model_lineage = native_model_lineage + self.logged_run_name: Optional[str] = None self._reset() def _reset(self) -> None: @@ -123,7 +124,7 @@ def check(self, registry: Any = None) -> None: def start(self, record: RunRecord) -> None: self._reset() - self.run_id = None + self.run_id = self.logged_run_name = None cls = self._module.ExperimentTracking tracking = getattr(cls, '_instance', None) if tracking is None: @@ -139,11 +140,14 @@ def start(self, record: RunRecord) -> None: experiment = tracking.set_experiment(self.experiment_name, database_name=database, schema_name=schema) - self._run = tracking.start_run( - self._new_run_name( - tracking, experiment.name, - self.run_name or _run_identifier(record.name or 'fit'))) - self.run_id = str(self._run.name) + existing = [ + str(row['name']) + for row in tracking._sql_client.show_runs_in_experiment( + experiment_name=experiment.name) + ] + self._run = tracking.start_run(self._new_run_name( + existing, _identifier(record_name(record, self.run_name)))) + self.run_id = self.logged_run_name = str(self._run.name) if self.native_model_lineage: self._run_context.enter_context(self._run) self._start_tags = set(record.tags) @@ -205,26 +209,16 @@ def finish(self, record: Optional[RunRecord], def log_metrics(self, metrics: Mapping[str, float], step: int) -> None: self._log_metrics(metrics, _EPOCH_METRICS, step=step) - def _new_run_name(self, tracking: Any, experiment: Any, - run_name: str) -> Optional[str]: + def _new_run_name(self, existing: List[str], run_name: str) -> str: # start_run matches name with LIKE. - existing = [ - str(row['name']) - for row in tracking._sql_client.show_runs_in_experiment( - experiment_name=experiment) - ] - short = uuid.uuid4().hex - for name in (run_name, _suffixed(run_name, short[:8]), - _suffixed(run_name, short)): - if not _taken(name, existing): - break - else: - name = None - if name != run_name: + name, number = run_name, 1 + while _taken(name, existing): + number += 1 + name = f'{run_name}_{number}' + if number > 1: logger.warning( 'Snowflake run %s already exists in experiment %s; recording ' - 'this fit as %s', run_name, self.experiment_name, - name or 'a generated run name') + 'this run as %s', run_name, self.experiment_name, name) return name def _log_scalar(self, summary: Scalar) -> None: @@ -276,22 +270,11 @@ def _log_metrics(self, metrics: Mapping[str, float], limit: int, self._tracking.log_metrics(kept, step=step) -def _run_identifier(name: str) -> str: - name = re.sub(r'[^A-Za-z0-9_]', '_', name) - return name if re.match(r'[A-Za-z_]', name) else f'_{name}' - - -def _suffixed(name: str, suffix: str) -> str: - if len(name) > 1 and name[0] == name[-1] == '"': - return f'{name[:-1]}_{suffix}"' - return f'{name}_{suffix}' +def _identifier(name: str) -> str: + return re.sub(r'[^A-Za-z0-9_]+', '_', name).upper() def _taken(name: str, existing: List[str]) -> bool: - resolved = resolve_snowflake_identifier(name) - regex = ''.join('.*' if c == '%' else '.' if c == '_' else re.escape(c) - for c in name) - return any( - other == resolved - or re.fullmatch(regex, other, re.IGNORECASE | re.DOTALL) is not None - for other in existing) + regex = ''.join('.' if c == '_' else re.escape(c) for c in name) + return any(re.fullmatch(regex, other, re.IGNORECASE | re.DOTALL) + for other in existing) diff --git a/packages/kumo-engine/kumo/engine/trainer/_run_spec.py b/packages/kumo-engine/kumo/engine/trainer/_run_spec.py index e33bb8212..a4a3b95ce 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_run_spec.py +++ b/packages/kumo-engine/kumo/engine/trainer/_run_spec.py @@ -12,6 +12,7 @@ from kumo_core.artifacts import PROVENANCE_FILE from kumo_core.summaries import Summary +from kumo.engine._run_names import fit_name, model_label from kumo.engine._tracking import RunSpec from kumo.engine._version import __version__ @@ -30,10 +31,12 @@ class FitContext: def fit_spec(context: FitContext) -> RunSpec: started = context.started or datetime.now(timezone.utc) + model = model_label(context.model_name, context.query, context.task_type) return RunSpec( kind='fit', - name='-'.join(filter(None, ('fit', context.model_name, - started.strftime('%Y%m%d-%H%M')))), + name=fit_name(model, started), + model_label=model, + started=started, params=dict(_flatten(context.model_plan)), parents={'snapshot_id': context.snapshot_id, 'task_type': context.task_type, diff --git a/packages/kumo-engine/test/_snowflake_experiment_fake.py b/packages/kumo-engine/test/_snowflake_experiment_fake.py index 3c914cedf..ff150009c 100644 --- a/packages/kumo-engine/test/_snowflake_experiment_fake.py +++ b/packages/kumo-engine/test/_snowflake_experiment_fake.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""A recording fake of `snowflake.ml.experiment.ExperimentTracking`.""" +"""A recording stand-in for Snowflake's ExperimentTracking.""" import os import re @@ -157,9 +157,14 @@ def log_metrics(self, metrics: Dict[str, float], step: int = 0) -> None: def log_artifact(self, local_path: str, artifact_path: Optional[str] = None) -> None: self._record('log_artifact', local_path, artifact_path) + # Artifact URIs do not quote identifiers. + if not re.fullmatch(r'[A-Z_][A-Z0-9_$]*', self._run.name): + raise RuntimeError(f'Run "{self._run.name}" does not exist in ' + f'Experiment {self.experiment}.') with open(local_path, encoding='utf-8') as file: - self._active()['artifacts'][os.path.basename(local_path)] = ( - file.read()) + self._active()['artifacts'][os.path.join( + artifact_path or '', os.path.basename(local_path))] = ( + file.read()) class FakeSession: diff --git a/packages/kumo-engine/test/test_mlflow_tracking.py b/packages/kumo-engine/test/test_mlflow_tracking.py index b804557b9..2bb4918d5 100644 --- a/packages/kumo-engine/test/test_mlflow_tracking.py +++ b/packages/kumo-engine/test/test_mlflow_tracking.py @@ -26,6 +26,7 @@ ) import kumo.engine as kumo +from kumo.engine._manifest import TrackingRef from kumo.engine._scratch import get_scratch_dir from kumo.engine._tracking import logged_names, run_record from kumo.engine.trainer._run_spec import fit_spec @@ -113,7 +114,8 @@ def test_fit_logs_a_run_per_fit_and_per_trial(table_and_snapshot, tmp_path, if run.data.tags.get('mlflow.parentRunId') == outer.info.run_id ] assert tracking.run_id == parent.info.run_id - assert re.fullmatch(r'fit-\d{8}-\d{4}', parent.info.run_name) + assert re.fullmatch(r'fit-ents-regression-\d{8}-\d{6}', + parent.info.run_name) children = sorted( (run for run in runs if run.data.tags.get('mlflow.parentRunId') == parent.info.run_id), @@ -133,8 +135,11 @@ def test_fit_logs_a_run_per_fit_and_per_trial(table_and_snapshot, tmp_path, assert parent.data.tags['kumo.bundle_id'] == result.id stored = tmp_path / 'root' / 'models' / result.id / PROVENANCE_FILE assert json.loads(stored.read_text())['tracking'] == { - 'platform': 'mlflow', 'run_id': parent.info.run_id} - assert result.tracking.run_id == parent.info.run_id + 'platform': 'mlflow', 'run_id': parent.info.run_id, + 'experiment': 'ents', 'run_name': parent.info.run_name} + assert result.tracking == TrackingRef( + platform='mlflow', run_id=parent.info.run_id, experiment='ents', + run_name=parent.info.run_name) logged = client.download_artifacts(parent.info.run_id, PROVENANCE_FILE, str(tmp_path)) assert json.loads(Path(logged).read_text())['id'] == result.id @@ -161,8 +166,8 @@ def test_fit_logs_a_run_per_fit_and_per_trial(table_and_snapshot, tmp_path, 'uri': str((tmp_path / 'data' / 'ents.parquet').resolve()) } - assert [run.info.run_name - for run in children] == ['trial-1-attempt-1', 'trial-2-attempt-1'] + assert [run.info.run_name for run in children] == [ + f'trial-{i}-{parent.info.run_name}' for i in (1, 2)] for child in children: assert child.info.status == 'FINISHED' assert child.data.tags['kumo.status'] == 'succeeded' @@ -224,15 +229,16 @@ def down(self, *args, metrics=(), **kwargs): assert statuses == {'FAILED'} -def test_a_user_run_name_wins_over_the_default(tracking_uri): +def test_a_user_run_name_names_every_run(tracking_uri): from kumo.engine.integrations.mlflow import MlflowTracking - replay(MlflowTracking(run_name='nightly', log_system_metrics=False)) + tracking = MlflowTracking(run_name='nightly', log_system_metrics=False) + replay(tracking) join_tracking_threads() names = {run.info.run_name for run in _runs(mlflow.MlflowClient()) if 'mlflow.parentRunId' not in run.data.tags} - assert names == {'nightly'} + assert names == {'fit-nightly'} SAGEMAKER = ('arn:aws:sagemaker:us-west-2:account:' diff --git a/packages/kumo-engine/test/test_snowflake_tracking.py b/packages/kumo-engine/test/test_snowflake_tracking.py index b6de6819e..be6b766e2 100644 --- a/packages/kumo-engine/test/test_snowflake_tracking.py +++ b/packages/kumo-engine/test/test_snowflake_tracking.py @@ -3,8 +3,8 @@ # SPDX-License-Identifier: Apache-2.0 import inspect -import re import threading +from datetime import datetime, timezone from unittest import mock import _snowflake_experiment_fake @@ -14,7 +14,7 @@ from test_fit_events import table_and_snapshot # noqa: F401 from test_tracking_contract import CONTEXT, EVENTS, RESULT -from kumo.engine._tracking import TrackingSession +from kumo.engine._tracking import RunRecord, TrackingSession from kumo.engine.trainer._run_spec import fit_spec, fitted_spec FAKED = ('__init__', 'set_experiment', 'start_run', 'end_run', 'log_params', @@ -153,8 +153,9 @@ def test_location_defaults_to_the_session_not_the_last_fit(fake): assert not any(call[0].startswith('run.') for call in fake._instance.calls) -@pytest.mark.parametrize('taken', ['FIT', 'MYXRUN']) -def test_an_existing_run_is_never_resumed(fake, taken, caplog): +@pytest.mark.parametrize('taken,run_name,new', [ + ('FIT', 'fit', 'FIT_2'), ('FIT_MYXRUN', 'my_run', 'FIT_MY_RUN_2')]) +def test_an_existing_run_is_never_resumed(fake, taken, run_name, new, caplog): session = FakeSession() user = fake(session) user.set_experiment('ents') @@ -168,12 +169,11 @@ def test_an_existing_run_is_never_resumed(fake, taken, caplog): tracking = _tracking(session, run_name=run_name) _replay(tracking) - new = tracking.run_id - assert re.fullmatch(f'{run_name}_[0-9a-f]{{8}}'.upper(), new) + assert tracking.run_id == tracking.logged_run_name == new assert set(user.runs) == {taken, new} assert user.runs[taken]['params'] == {} assert user.runs[new]['status'] == 'FINISHED' - assert 'recording this fit as' in caplog.text + assert 'recording this run as' in caplog.text def test_metric_and_param_limits(fake, monkeypatch): @@ -197,14 +197,6 @@ def test_metric_and_param_limits(fake, monkeypatch): assert run['params']['level'] == '2' -def test_quoted_run_names_are_taken_only_by_their_exact_name(): - from kumo.engine.integrations.snowflake import _taken - - assert _taken('"MyRun"', ['MyRun']) - assert not _taken('"MyRun"', ['MYRUN']) - assert _taken('my_run', ['MYXRUN']) - - def test_the_users_session_is_not_altered(fake): session = FakeSession() @@ -272,3 +264,27 @@ def test_mlops_loads_a_registered_model_through_the_tracker( })] with pytest.raises(ValueError, match='SnowflakeTracking\\(connector'): _tracking().load_model(location) + + +@pytest.mark.parametrize('record, name', [ + (RunRecord(tags={'kumo.kind': 'fit'}, model_label='kumo-demo-churn', + started=datetime(2026, 10, 29, 6, 18, 5, tzinfo=timezone.utc)), + 'FIT_KUMO_DEMO_CHURN_20261029_061805'), +]) +def test_runs_are_named_by_kind_model_version_and_date(fake, record, name): + tracking = _tracking() + tracking.start(record) + tracking.finish(record, None) + + assert list(fake._instance.runs) == [name] + + +@pytest.mark.parametrize('record, name', [ + (RunRecord(tags={'kumo.kind': 'fit'}), 'FIT_NIGHTLY'), +]) +def test_a_user_run_name_names_every_run(fake, record, name): + tracking = _tracking(run_name='nightly') + tracking.start(record) + tracking.finish(record, None) + + assert list(fake._instance.runs) == [name] diff --git a/packages/kumo-engine/test/test_tracking_contract.py b/packages/kumo-engine/test/test_tracking_contract.py index 0e46bfc25..79df5a3b2 100644 --- a/packages/kumo-engine/test/test_tracking_contract.py +++ b/packages/kumo-engine/test/test_tracking_contract.py @@ -172,7 +172,7 @@ class Adapter: nested_runs: bool = False step_metrics: bool = False artifacts: bool = False - run_name: Callable[[str], str] = str + run_name: Callable[[str], Optional[str]] = str def gaps(expected: Dict[str, RunRecord], observed: Dict[str, Observed], @@ -221,6 +221,12 @@ def differ(what, have, want): return found +def _trial(run): + r"""An MLflow trial run's record name, ``trial--attempt-``.""" + trial = run.info.run_name.split('-')[1] + return f"trial-{trial}-attempt-{run.data.params['attempt']}" + + @pytest.fixture def mlflow_adapter(tmp_path, monkeypatch): mlflow = pytest.importorskip('mlflow') @@ -236,7 +242,7 @@ def observe(): [e.experiment_id for e in client.search_experiments()]) kumo_ids = { r.info.run_id: - FIT if r.data.tags['kumo.kind'] == 'fit' else r.info.run_name + FIT if r.data.tags['kumo.kind'] == 'fit' else _trial(r) for r in runs } @@ -247,7 +253,8 @@ def history(run_id, key): return { kumo_ids[run.info.run_id]: Observed( - name=run.info.run_name, params=run.data.params, + name=(run.info.run_name if run.data.tags['kumo.kind'] != 'trial' + else _trial(run)), params=run.data.params, metrics=run.data.metrics, tags=run.data.tags, steps={k: history(run.info.run_id, k) for k in run.data.metrics}, @@ -278,10 +285,7 @@ def snowflake_adapter(monkeypatch): import _snowflake_experiment_fake fake = _snowflake_experiment_fake.install(monkeypatch) - from kumo.engine.integrations.snowflake import ( - SnowflakeTracking, - _run_identifier, - ) + from kumo.engine.integrations.snowflake import SnowflakeTracking def observe(): observed = {} @@ -319,7 +323,7 @@ def bucket(key): make=lambda: SnowflakeTracking(_snowflake_experiment_fake.FakeSession(), experiment_name='ents', tags=USER_TAGS), observe=observe, step_metrics=True, artifacts=True, - run_name=lambda name: _run_identifier(name).upper()) + run_name=lambda name: None) @pytest.fixture(params=['mlflow', 'snowflake']) @@ -331,7 +335,9 @@ def test_records_of_the_recorded_fit(): records = replay() fit = records[FIT] - assert fit.name == 'fit-20261029-0618' + assert fit.name == 'fit-ents-regression-20261029-061800' + assert (fit.model_label, fit.started) == ('ents-regression', + CONTEXT.started) assert fit.tags == {'team': 'ents', 'kumo.kind': 'fit', 'kumo.bundle_id': BUNDLE, 'kumo.snapshot_id': SNAPSHOT, 'kumo.task_type': 'regression', @@ -415,7 +421,7 @@ def test_adapter_logs_each_summary_type(adapter, mlflow_adapter): tables = {'col_stats.json': {'columns': ['table', 'rows'], 'data': [['ents', 400]]}} if adapter is not mlflow_adapter: - assert {name: json.loads(fit.artifacts[name]) + assert {name: json.loads(fit.artifacts[f'graph/{name}']) for name in tables} == tables return import mlflow From 1f5bb865bf7bd3cc3e5d01dc21f6144e5676d8a5 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 03:02:02 +0000 Subject: [PATCH 198/278] Store the model under the artifact root after the fit, from Trainer.fit fit resolves the MLOps settings, runs the new _fit, then stores the bundle; _finish_fit no longer takes settings. Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/trainer/trainer.py | 49 ++++++++++++------- .../test/test_bundle_provenance.py | 10 ++-- 2 files changed, 38 insertions(+), 21 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index b90bfaa30..ceac4ae39 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -276,28 +276,29 @@ def _listed_holdout(bundle_dir: Path) -> tuple[Optional[Path], Optional[bool]]: return path, truncated -def _store_model(result: TrainingResult, settings: Any) -> None: - r"""Copy the bundle under the root unless already there.""" +def _store_model(bundle_dir: Path, manifest: Optional[ModelManifest], + settings: Any) -> None: + r"""Copies the bundle under the root unless it is there.""" from kumo_core.object_store_sink import holds_bundle, upload_bundle from pyarrow.fs import FileSelector from kumo.engine._mlops import warn_not_stored from kumo.engine.connector import file_store - if result.bundle_id is None: + if manifest is None: return try: filesystem, root = file_store( settings.root.connector, - settings.root.location(f'models/{result.bundle_id}')) - if holds_bundle(filesystem, root, str(result.bundle_dir)): + settings.root.location(f'models/{manifest.id}')) + if holds_bundle(filesystem, root, str(bundle_dir)): return if filesystem.get_file_info( FileSelector(root, allow_not_found=True)): filesystem.delete_dir_contents(root, missing_dir_ok=True) - upload_bundle(filesystem, root, str(result.bundle_dir)) + upload_bundle(filesystem, root, str(bundle_dir)) except Exception as error: - warn_not_stored(f'model {result.bundle_id}', error, stacklevel=4) + warn_not_stored(f'model {manifest.id}', error, stacklevel=3) def _published_holdout(trainer: Any) -> tuple[Optional[Path], Optional[bool]]: @@ -474,6 +475,30 @@ def fit( Each call returns a new `TrainingResult`; `Trainer.result` holds the latest. """ + if export_config is not None: + _require_publisher(export_config) + from kumo.engine._mlops import resolve + + settings = resolve(mlops) + run_id = f'{self._id}-fit-{new_run_id()}' + result = self._fit(train_table, snapshot=snapshot, + export_config=export_config, + warm_start=warm_start, run_id=run_id) + if settings is not None and settings.wants('model'): + _store_model(result.bundle_dir, + bundle_manifest(result.bundle_dir, result.bundle_id), + settings) + return result + + def _fit( + self, + train_table: TrainingTable, + *, + snapshot: GraphSnapshot, + export_config: Optional[ModelOutputConfig], + warm_start: Optional[Self], + run_id: str, + ) -> TrainingResult: from kumo_core.automl.runtime import ( inherit_warm_start_model_plan, validate_warm_start_task_type, @@ -495,13 +520,7 @@ def fit( from kumo.engine.trainer._runtime import local_training_bindings - if export_config is not None: - _require_publisher(export_config) - from kumo.engine._mlops import resolve - - settings = resolve(mlops) started_at = datetime.now(timezone.utc) - run_id = f'{self._id}-fit-{new_run_id()}' run_dir = self.scratch_dir / 'runs' / run_id model_plan = _to_core_model_plan(self._model_plan) @@ -677,7 +696,6 @@ def fit( warm_start._training_result.bundle_dir, warm_start_bundle_id), started_at=started_at, fit_mode=fit_mode), - settings=settings, ) def _finish_fit( @@ -691,7 +709,6 @@ def _finish_fit( graph_config: Optional[Dict[str, Any]] = None, model_plan: Optional[ModelPlan] = None, provenance: Optional[FitRecord] = None, - settings: Optional[Any] = None, ) -> TrainingResult: """Commit durable fit outputs before making this trainer usable.""" from kumo_core.artifacts import bundle_id_from_manifest @@ -737,8 +754,6 @@ def _finish_fit( ) training_result = replace(training_result, published_model=published) - if settings is not None and settings.wants('model'): - _store_model(training_result, settings) training_result = replace(training_result, _trainer=self) self._training_result = training_result diff --git a/packages/kumo-engine/test/test_bundle_provenance.py b/packages/kumo-engine/test/test_bundle_provenance.py index 1a1954d30..7b0f6bd28 100644 --- a/packages/kumo-engine/test/test_bundle_provenance.py +++ b/packages/kumo-engine/test/test_bundle_provenance.py @@ -22,7 +22,7 @@ ) from kumo.engine.trainer._provenance import FitRecord, embedded_model from kumo.engine.trainer._spec import MANIFEST_FILE -from kumo.engine.trainer.trainer import TrainingResult +from kumo.engine.trainer.trainer import TrainingResult, _store_model def test_fit_records_a_self_contained_manifest(fitted): @@ -109,13 +109,15 @@ def _finish_stub_fit(bundle, trainer, warm_start=None, settings=None): (bundle / MANIFEST_FILE).write_text(json.dumps({'artifacts': []})) add_bundle_file(bundle, 'trainer_config', 'model/config.json') table = SimpleNamespace(tt_id='tt-raw', id='tt-1', manifest=None) - return trainer._finish_fit( + result = trainer._finish_fit( TrainingResult({}, {}, bundle), predictive_query='PREDICT x', snapshot_id='snap', export_config=None, training_table_generation=TrainingTableGenerationPlan(), provenance=FitRecord(table, warm_start, datetime.now(timezone.utc), - None), - settings=settings) + None)) + if settings is not None and settings.wants('model'): + _store_model(result.bundle_dir, result.manifest, settings) + return result def test_missing_parents_are_stubs_and_retrain_chains_stay_flat(tmp_path): From b6fe35d3f10bd7a35b041f958f0d08207ef98de6 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 03:04:19 +0000 Subject: [PATCH 199/278] Report MLOps store failures as MLOps.on_error says A failed snapshot, training table or model store warns by default, as before, and raises with on_error='raise'. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_mlops.py | 46 +++++++++++++++++-- .../kumo-engine/kumo/engine/_snapshot_io.py | 7 ++- packages/kumo-engine/kumo/engine/graph.py | 6 ++- .../kumo/engine/pquery/training_table.py | 13 ++++-- .../kumo/engine/trainer/trainer.py | 17 ++++--- packages/kumo-engine/test/test_mlops.py | 2 + .../kumo-engine/test/test_snapshot_export.py | 14 ++++++ 7 files changed, 87 insertions(+), 18 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_mlops.py b/packages/kumo-engine/kumo/engine/_mlops.py index 6abfc020f..e1f140169 100644 --- a/packages/kumo-engine/kumo/engine/_mlops.py +++ b/packages/kumo-engine/kumo/engine/_mlops.py @@ -9,7 +9,16 @@ import posixpath import warnings from dataclasses import dataclass, field -from typing import Any, FrozenSet, Iterable, Iterator, Optional, Tuple, Union +from typing import ( + Any, + FrozenSet, + Iterable, + Iterator, + Literal, + Optional, + Tuple, + Union, +) from kumo.engine._manifest import KINDS @@ -35,10 +44,12 @@ class MLOps: Args: artifact_root: a URI or path, or a ``(connector, path)`` pair. store: artifact kinds to store; ``True`` for all, ``False`` for none. + on_error: ``'warn'`` or ``'raise'`` when MLOps fails. """ artifact_root: Union[str, Tuple[Any, str]] store: Union[bool, Iterable[str]] = field(default=True, kw_only=True) + on_error: Literal['warn', 'raise'] = field(default='warn', kw_only=True) root: _Root = field(init=False, repr=False, compare=False) kinds: FrozenSet[str] = field(init=False, repr=False, compare=False) @@ -50,6 +61,9 @@ def __post_init__(self) -> None: if not self.artifact_root: raise ValueError('MLOps needs an artifact_root.') + if self.on_error not in ('warn', 'raise'): + raise ValueError(f"on_error must be 'warn' or 'raise', not " + f"{self.on_error!r}.") if isinstance(self.artifact_root, str): if self.artifact_root.startswith('@'): raise ValueError( @@ -154,9 +168,33 @@ def resolve(settings: Optional[MLOps] = None) -> Optional[MLOps]: return settings if settings is not None else _CURRENT.get() -def warn_not_stored(what: str, error: Exception, stacklevel: int) -> None: - warnings.warn(f'Could not store {what} under the artifact root: {error}', - UserWarning, stacklevel=stacklevel + 1) +def raise_on_error(settings: Optional[MLOps], error: Optional[BaseException], + call: str, stacklevel: int = 2) -> None: + r"""Raises or warns about ``error``, as ``settings.on_error`` says.""" + if error is None or settings is None: + return + message = f'{call} finished, but MLOps failed: {full_error(error)}' + if settings.on_error == 'raise': + raise RuntimeError(message) from error + warn_failure(message, error, stacklevel + 1) + + +def full_error(error: BaseException) -> str: + r"""``Type: message``, not shortened.""" + return f'{type(error).__name__}: {error}' + + +def warn_failure(message: str, error: BaseException, stacklevel: int = 2, + also: Optional[BaseException] = None) -> None: + r"""Warns ``message`` once per ``error`` (and ``also``).""" + if getattr(error, '_kumo_warned', False): + return + warnings.warn(message, UserWarning, stacklevel=stacklevel + 1) + for reported in (error, also): + try: + setattr(reported, '_kumo_warned', True) + except AttributeError: + pass def on_root(source: Any, path: str) -> Tuple[Any, str]: diff --git a/packages/kumo-engine/kumo/engine/_snapshot_io.py b/packages/kumo-engine/kumo/engine/_snapshot_io.py index ddeb1faaf..b0f70d88d 100644 --- a/packages/kumo-engine/kumo/engine/_snapshot_io.py +++ b/packages/kumo-engine/kumo/engine/_snapshot_io.py @@ -191,9 +191,7 @@ def export_location(snapshot: Any, path: str) -> str: def store(snapshot: Any, settings: Any) -> None: - r"""Export ``snapshot`` under the root unless already there.""" - from kumo.engine._mlops import warn_not_stored - + r"""Exports ``snapshot`` under the root unless it is there.""" try: connector = settings.root.connector path = settings.root.location('snapshots') @@ -207,7 +205,8 @@ def store(snapshot: Any, settings: Any) -> None: and marker.get('snapshot_id') == snapshot.id): export(snapshot, connector, path) except Exception as error: - warn_not_stored(f'snapshot {snapshot.id}', error, stacklevel=3) + raise RuntimeError(f"Could not store snapshot {snapshot.id} under the " + f"artifact root: {error}") from error def export(snapshot: Any, connector: Any, path: str) -> str: diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index 0dcb13165..2b4a99ece 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -757,9 +757,13 @@ def snapshot(self, *, record_sources: bool = True, snapshot = self._snapshot_no_lock(graph_key, persist_dir, record_sources) if settings is not None and settings.wants('snapshot'): + from kumo.engine._mlops import raise_on_error from kumo.engine._snapshot_io import store - store(snapshot, settings) + try: + store(snapshot, settings) + except Exception as error: + raise_on_error(settings, error, 'Graph.snapshot') return snapshot def _snapshot_no_lock(self, graph_key: str, persist_dir: Path, diff --git a/packages/kumo-engine/kumo/engine/pquery/training_table.py b/packages/kumo-engine/kumo/engine/pquery/training_table.py index b84b26d9a..c805684c7 100644 --- a/packages/kumo-engine/kumo/engine/pquery/training_table.py +++ b/packages/kumo-engine/kumo/engine/pquery/training_table.py @@ -133,7 +133,6 @@ def load_from_connector(cls, connector: Any, *, path: str, def _store(self, settings: Any) -> None: from pyarrow.fs import FileType - from kumo.engine._mlops import warn_not_stored from kumo.engine.connector import file_store from kumo.engine.trainer.output_config import TrainingTableOutputConfig @@ -145,7 +144,9 @@ def _store(self, settings: Any) -> None: self.export( TrainingTableOutputConfig(settings.root.connector, path)) except Exception as error: - warn_not_stored(f'training table {self.id}', error, stacklevel=3) + raise RuntimeError( + f"Could not store training table {self.id} under the " + f"artifact root: {error}") from error @classmethod def from_source_table( @@ -307,7 +308,13 @@ def finish_generated(table: TrainingTable, snapshot: GraphSnapshot, _record_manifest(table, snapshot, request, seed, source_table=source_table) settings = resolve(mlops) if settings is not None and settings.wants('training_table'): - table._store(settings) + from kumo.engine._mlops import raise_on_error + + try: + table._store(settings) + except Exception as error: + raise_on_error(settings, error, 'Training table generation', + stacklevel=3) return table diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index ceac4ae39..913bb9a5c 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -282,7 +282,6 @@ def _store_model(bundle_dir: Path, manifest: Optional[ModelManifest], from kumo_core.object_store_sink import holds_bundle, upload_bundle from pyarrow.fs import FileSelector - from kumo.engine._mlops import warn_not_stored from kumo.engine.connector import file_store if manifest is None: @@ -298,7 +297,8 @@ def _store_model(bundle_dir: Path, manifest: Optional[ModelManifest], filesystem.delete_dir_contents(root, missing_dir_ok=True) upload_bundle(filesystem, root, str(bundle_dir)) except Exception as error: - warn_not_stored(f'model {manifest.id}', error, stacklevel=3) + raise RuntimeError(f"Could not store model {manifest.id} under the " + f"artifact root: {error}") from error def _published_holdout(trainer: Any) -> tuple[Optional[Path], Optional[bool]]: @@ -477,7 +477,7 @@ def fit( """ if export_config is not None: _require_publisher(export_config) - from kumo.engine._mlops import resolve + from kumo.engine._mlops import raise_on_error, resolve settings = resolve(mlops) run_id = f'{self._id}-fit-{new_run_id()}' @@ -485,9 +485,14 @@ def fit( export_config=export_config, warm_start=warm_start, run_id=run_id) if settings is not None and settings.wants('model'): - _store_model(result.bundle_dir, - bundle_manifest(result.bundle_dir, result.bundle_id), - settings) + try: + _store_model( + result.bundle_dir, + bundle_manifest(result.bundle_dir, result.bundle_id), + settings) + except Exception as error: + raise_on_error(settings, error, + 'Trainer.fit (its model is in trainer.result)') return result def _fit( diff --git a/packages/kumo-engine/test/test_mlops.py b/packages/kumo-engine/test/test_mlops.py index b3ec3827b..3289c8467 100644 --- a/packages/kumo-engine/test/test_mlops.py +++ b/packages/kumo-engine/test/test_mlops.py @@ -27,6 +27,8 @@ def test_mlops_needs_an_artifact_root_it_can_write_files_to(tmp_path): kumo.MLOps(artifact_root='@DB.SCHEMA.STAGE/kumo') with pytest.raises(ValueError, match='Unknown store kinds'): kumo.MLOps(artifact_root=str(tmp_path), store={'models'}) + with pytest.raises(ValueError, match='on_error'): + kumo.MLOps(artifact_root=str(tmp_path), on_error='ignore') def test_a_root_is_a_uri_or_a_connector_and_path(tmp_path): diff --git a/packages/kumo-engine/test/test_snapshot_export.py b/packages/kumo-engine/test/test_snapshot_export.py index e6b4f18db..29aab0f49 100644 --- a/packages/kumo-engine/test/test_snapshot_export.py +++ b/packages/kumo-engine/test/test_snapshot_export.py @@ -785,3 +785,17 @@ def test_loading_a_missing_export_names_its_full_path(tmp_path, exports): with pytest.raises(FileNotFoundError, match=f'No snapshot export at {missing} '): _load(exports, 'snapshots/graph/version') + + +@pytest.mark.parametrize('on_error', ['warn', 'raise']) +def test_a_failed_store_follows_on_error(tmp_path, scratch, graph, + monkeypatch, on_error): + monkeypatch.setattr(snapshot_io, 'export', + Mock(side_effect=OSError('ACCESS_DENIED'))) + cfg = kumo.MLOps(str(tmp_path / 'root'), on_error=on_error) + if on_error == 'raise': + with pytest.raises(RuntimeError, match='Could not store snapshot'): + graph.snapshot(mlops=cfg) + else: + with pytest.warns(UserWarning, match='ACCESS_DENIED'): + graph.snapshot(mlops=cfg) From 689435f594e28b84b67b003cc40da35ea8e91885 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 03:05:02 +0000 Subject: [PATCH 200/278] Route summaries through _summary_tables instead of TrackingBase dispatch Adds Note, Table.brief and Scalar.role; MLflow logs a run's tables with one loggedArtifacts update and digests datasets by name, source and schema. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-core/src/kumo_core/summaries.py | 18 ++- .../kumo/engine/_summary_tables.py | 59 ++++++++++ packages/kumo-engine/kumo/engine/_tracking.py | 30 +---- .../kumo/engine/integrations/mlflow.py | 109 ++++++++++-------- .../kumo/engine/integrations/snowflake.py | 23 ++-- .../kumo-engine/test/test_summary_tables.py | 19 +++ .../test/test_tracking_contract.py | 5 +- 7 files changed, 174 insertions(+), 89 deletions(-) create mode 100644 packages/kumo-engine/kumo/engine/_summary_tables.py create mode 100644 packages/kumo-engine/test/test_summary_tables.py diff --git a/packages/kumo-core/src/kumo_core/summaries.py b/packages/kumo-core/src/kumo_core/summaries.py index 55064547e..25c366db5 100644 --- a/packages/kumo-core/src/kumo_core/summaries.py +++ b/packages/kumo-core/src/kumo_core/summaries.py @@ -4,13 +4,16 @@ r"""Typed, tracker-neutral results, named ``
/``.""" from dataclasses import dataclass -from typing import Any, Mapping, Optional, Tuple, Union +from typing import Any, Literal, Mapping, Optional, Tuple, Union @dataclass(frozen=True) class Scalar: + r"""A metric; ``role='model'`` also logs it on the model.""" + name: str value: float + role: Optional[Literal['model']] = None @dataclass(frozen=True) @@ -21,9 +24,20 @@ class Tag: @dataclass(frozen=True) class Table: + r"""A table; a ``brief`` one (at most 12 x 5) goes on the run card.""" + name: str columns: Tuple[str, ...] rows: Tuple[Tuple[Any, ...], ...] + brief: bool = False + + +@dataclass(frozen=True) +class Note: + r"""A line of text for the run card, under ``section``.""" + + section: str + text: str @dataclass(frozen=True) @@ -37,4 +51,4 @@ class Dataset: context: str -Summary = Union[Scalar, Tag, Table, Dataset] +Summary = Union[Scalar, Tag, Table, Note, Dataset] diff --git a/packages/kumo-engine/kumo/engine/_summary_tables.py b/packages/kumo-engine/kumo/engine/_summary_tables.py new file mode 100644 index 000000000..91f26a06d --- /dev/null +++ b/packages/kumo-engine/kumo/engine/_summary_tables.py @@ -0,0 +1,59 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""Summaries routed for trackers, and tables as tracker JSON.""" + +import json +import math +from dataclasses import dataclass, field +from typing import Any, Dict, List, Sequence + +from kumo_core.summaries import Dataset, Note, Scalar, Summary, Table, Tag + +_DEFINITIONS = frozenset({'from', 'to', 'threshold'}) + + +@dataclass +class Routed: + r"""Summaries sorted by how a tracker writes them.""" + + scalars: List[Scalar] = field(default_factory=list) + tags: Dict[str, str] = field(default_factory=dict) + tables: List[Table] = field(default_factory=list) + card: List[Summary] = field(default_factory=list) + datasets: List[Dataset] = field(default_factory=list) + + +def route(summaries: Sequence[Summary]) -> Routed: + r"""Sorts summaries by how a tracker writes them.""" + routed, tables = Routed(), [] + for summary in summaries: + if isinstance(summary, Scalar): + routed.scalars.append(summary) + elif isinstance(summary, Tag): + routed.tags[summary.name] = summary.value + elif isinstance(summary, Dataset): + routed.datasets.append(summary) + elif isinstance(summary, Note) or (isinstance(summary, Table) + and summary.brief): + routed.card.append(summary) + elif isinstance(summary, Table): + tables.append(summary) + else: + raise TypeError(f'No route for {type(summary).__name__}') + routed.tables = tables + return routed + + +def table_json(table: Table) -> str: + r"""Table JSON for tracker table views.""" + digits = [10 if column in _DEFINITIONS else 6 for column in table.columns] + return json.dumps({'columns': list(table.columns), 'data': [ + [_rounded(value, places) for value, places in zip(row, digits)] + for row in table.rows]}, default=str) + + +def _rounded(value: Any, digits: int) -> Any: + if not isinstance(value, float) or not math.isfinite(value) or not value: + return value + return round(value, digits - 1 - math.floor(math.log10(abs(value)))) diff --git a/packages/kumo-engine/kumo/engine/_tracking.py b/packages/kumo-engine/kumo/engine/_tracking.py index 14bc81ee5..3ca520c8f 100644 --- a/packages/kumo-engine/kumo/engine/_tracking.py +++ b/packages/kumo-engine/kumo/engine/_tracking.py @@ -28,7 +28,7 @@ ) from kumo_core.instrumentation import format_error -from kumo_core.summaries import Dataset, Scalar, Summary, Table, Tag +from kumo_core.summaries import Summary from kumo.engine._events import Event, _finite_floats from kumo.engine._mlops import full_error, warn_failure @@ -211,34 +211,6 @@ def check(self, registry: Optional['ModelRegistryLocation'] = None def log_metrics(self, metrics: Mapping[str, float], step: int) -> None: pass - def log_summaries(self, summaries: Sequence[Summary]) -> None: - for summary in summaries: - if isinstance(summary, Scalar): - self._log_scalar(summary) - elif isinstance(summary, Tag): - self._log_tag(summary) - elif isinstance(summary, Dataset): - self._log_dataset(summary) - else: - self._log_table(summary) - - def _log_scalar(self, summary: Scalar) -> None: - raise NotImplementedError - - def _log_tag(self, summary: Tag) -> None: - raise NotImplementedError - - def _log_table(self, summary: Table) -> None: - raise NotImplementedError - - def _log_dataset(self, summary: Dataset) -> None: - pass - - -def table_json(table: Table) -> str: - return json.dumps({'columns': list(table.columns), - 'data': [list(row) for row in table.rows]}, default=str) - class _Foreign(TrackingBase): r"""A tracker that does not subclass `TrackingBase`, with its defaults.""" diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index b379e6c6e..98969d532 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -34,12 +34,13 @@ from urllib.parse import quote from kumo_core.instrumentation import format_error -from kumo_core.summaries import Dataset, Scalar, Table, Tag +from kumo_core.summaries import Dataset, Summary, Table from kumoapi.plugin import BUNDLE_ID_KEY, ModelRegistryLocation, PublishedModel from kumo.engine._events import Event from kumo.engine._mlops import full_error from kumo.engine._run_names import record_name, trial_name +from kumo.engine._summary_tables import Routed, route, table_json from kumo.engine._text import metric_key from kumo.engine._tracking import ( DROPPED_PARAMS, @@ -47,7 +48,6 @@ RunRecord, TrackingBase, by_step, - table_json, ) from kumo.engine._version import __version__ from kumo.engine.connector._extras import import_provider @@ -360,11 +360,11 @@ def finish(self, record: Optional[RunRecord], if record is not None and self._parent is not None: for epoch, metrics in by_step(record.steps): self._log_metrics(self._parent, metrics, epoch) + self._log_summaries(record.summaries) self._log_metrics(self._parent, record.metrics, model_id=self._model_id) self._log_tags(self._parent, record.tags) self._log_artifacts(record.artifacts) - self.log_summaries(record.summaries) if error is None: status = 'FINISHED' finally: @@ -383,27 +383,20 @@ def finish(self, record: Optional[RunRecord], def log_metrics(self, metrics: Mapping[str, float], step: int) -> None: self._log_metrics(self._parent, metrics, step) - def _log_scalar(self, summary: Scalar) -> None: - model_id = self._model_id if summary.name.startswith('test/') else None - self._log_metrics(self._parent, {summary.name: summary.value}, - model_id=model_id) - - def _log_tag(self, summary: Tag) -> None: - self._log_tags(self._parent, {summary.name: summary.value}) - - def _log_table(self, summary: Table) -> None: - path = f'{summary.name}.json' - self._log_artifacts({path: table_json(summary)}) - with self.target.scope(): - tags = self._client.get_run(self._parent).data.tags - logged = json.loads(tags.get('mlflow.loggedArtifacts', '[]')) - entry = {'path': path, 'type': 'table'} - if entry not in logged: - self._client.set_tag(self._parent, 'mlflow.loggedArtifacts', - json.dumps([*logged, entry])) - - def _log_dataset(self, summary: Dataset) -> None: - self._log_inputs(self._parent, [summary]) + def _log_summaries(self, summaries: Sequence[Summary]) -> None: + self._log_routed(route(summaries), self._parent) + + def _log_routed(self, routed: Routed, run_id: str) -> None: + self._log_inputs(run_id, routed.datasets) + scalars = routed.scalars + self._log_metrics(run_id, {s.name: s.value for s in scalars + if s.role == 'model'}, + model_id=self._model_id) + self._log_metrics(run_id, {s.name: s.value for s in scalars + if s.role != 'model'}) + if routed.tags: + self._log_tags(run_id, routed.tags) + self._write_tables(run_id, routed.tables) def _log_inputs(self, run_id: str, summaries: Sequence[Dataset]) -> None: datasets = [d for d in map(self._dataset_input, summaries) if d] @@ -416,38 +409,57 @@ def _dataset_input(self, summary: Dataset) -> Any: if source is None: return None entities = self._mlflow.entities - version = summary.source.get('version') or '' dataset = entities.Dataset( name=summary.name, - digest=version[:36] or hashlib.sha256(json.dumps( - dict(summary.source), sort_keys=True).encode()).hexdigest()[:8], + digest=_digest(summary.name, summary.source, summary.schema), source_type=source._get_source_type(), source=source.to_json(), schema=json.dumps({'mlflow_colspec': [ {'name': name, 'type': _COLUMN_TYPES.get(column['dtype'], 'string')} for name, column in summary.schema.items()]}), - profile=json.dumps(summary.profile, default=str)) - return entities.DatasetInput(dataset, tags=[entities.InputTag( - 'mlflow.data.context', summary.context)]) + profile=json.dumps({ + key: value for key, value in summary.profile.items() + if not isinstance(value, (Mapping, list, tuple)) + }, default=str)) + return entities.DatasetInput( + dataset, tags=[entities.InputTag('mlflow.data.context', + summary.context)]) + + def _write_tables(self, run_id: str, tables: Sequence[Table]) -> None: + if not tables: + return + for table in tables: + self._write_text(run_id, f'{table.name}.json', table_json(table)) + with self.target.scope(): + tags = self._client.get_run(run_id).data.tags + logged = json.loads(tags.get('mlflow.loggedArtifacts', '[]')) + entries = [{'path': f'{table.name}.json', 'type': 'table'} + for table in tables] + added = [entry for entry in entries if entry not in logged] + if added: + self._client.set_tag(run_id, 'mlflow.loggedArtifacts', + json.dumps([*logged, *added])) def _log_artifacts(self, artifacts: Mapping[str, str]) -> None: for name, text in artifacts.items(): - if name in self._artifacts: - continue - with self.target.scope(): - repository = self.target.artifact_repository( - self._client.get_run(self._parent).info.artifact_uri - ) if self.target.explicit else None - if repository is None: - self._client.log_text(self._parent, text, name) - else: - with tempfile.TemporaryDirectory() as tmp: - path = Path(tmp) / posixpath.basename(name) - path.write_text(text) - repository.log_artifact( - str(path), posixpath.dirname(name) or None) - self._artifacts.add(name) + if name not in self._artifacts: + self._write_text(self._parent, name, text) + self._artifacts.add(name) + + def _write_text(self, run_id: str, name: str, text: str) -> None: + with self.target.scope(): + repository = self.target.artifact_repository( + self._client.get_run(run_id).info.artifact_uri + ) if self.target.explicit else None + if repository is None: + self._client.log_text(run_id, text, name) + else: + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / posixpath.basename(name) + path.write_text(text) + repository.log_artifact(str(path), + posixpath.dirname(name) or None) def _log_tags(self, run_id: str, tags: Mapping[str, str]) -> None: entities = self._mlflow.entities @@ -488,6 +500,13 @@ def _log_metrics(self, run_id: Optional[str], run_id, metrics=batch[start:start + _BATCH_METRICS]) +def _digest(name: str, source: Mapping[str, Optional[str]], + schema: Mapping[str, Any]) -> str: + return hashlib.sha256(json.dumps( + [name, dict(source), dict(schema)], sort_keys=True, + default=str).encode()).hexdigest()[:8] + + _BATCH_METRICS = 1000 _SYSTEM_METRICS_FAILURES = 6 _COLUMN_TYPES = {'int64': 'long', 'int32': 'integer', 'float64': 'double', diff --git a/packages/kumo-engine/kumo/engine/integrations/snowflake.py b/packages/kumo-engine/kumo/engine/integrations/snowflake.py index b02a7f4fe..3f1107411 100644 --- a/packages/kumo-engine/kumo/engine/integrations/snowflake.py +++ b/packages/kumo-engine/kumo/engine/integrations/snowflake.py @@ -9,21 +9,21 @@ import tempfile from contextlib import ExitStack from pathlib import Path -from typing import Any, Dict, List, Mapping, Optional, Set, Tuple +from typing import Any, Dict, List, Mapping, Optional, Sequence, Set, Tuple from kumo_core.instrumentation import format_error from kumo_core.sources.sql import resolve_snowflake_identifier -from kumo_core.summaries import Scalar, Table, Tag +from kumo_core.summaries import Summary from kumo.engine._events import Event from kumo.engine._run_names import record_name +from kumo.engine._summary_tables import route, table_json from kumo.engine._tracking import ( DROPPED_METRICS, DROPPED_PARAMS, RunRecord, TrackingBase, by_step, - table_json, ) from kumo.engine.connector._extras import import_provider @@ -187,7 +187,7 @@ def finish(self, record: Optional[RunRecord], for k, v in record.tags.items() if k not in self._start_tags}, _MAX_PARAMS) self._log_artifacts(record.artifacts) - self.log_summaries(record.summaries) + self._log_summaries(record.summaries) dropped = {DROPPED_METRICS: len(self._dropped_metrics), DROPPED_PARAMS: self._dropped_params} if any(dropped.values()): @@ -221,14 +221,13 @@ def _new_run_name(self, existing: List[str], run_name: str) -> str: 'this run as %s', run_name, self.experiment_name, name) return name - def _log_scalar(self, summary: Scalar) -> None: - self._log_metrics({summary.name: summary.value}, _MAX_METRICS) - - def _log_tag(self, summary: Tag) -> None: - self._log_params({summary.name: summary.value}) - - def _log_table(self, summary: Table) -> None: - self._log_artifacts({f'{summary.name}.json': table_json(summary)}) + def _log_summaries(self, summaries: Sequence[Summary]) -> None: + routed = route(summaries) + self._log_metrics({s.name: s.value for s in routed.scalars}, + _MAX_METRICS) + self._log_params(routed.tags) + self._log_artifacts({f'{table.name}.json': table_json(table) + for table in routed.tables}) def _log_artifacts(self, artifacts: Mapping[str, str]) -> None: with tempfile.TemporaryDirectory() as directory: diff --git a/packages/kumo-engine/test/test_summary_tables.py b/packages/kumo-engine/test/test_summary_tables.py new file mode 100644 index 000000000..afe37e5f5 --- /dev/null +++ b/packages/kumo-engine/test/test_summary_tables.py @@ -0,0 +1,19 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json + +from kumo_core.summaries import Table + +from kumo.engine._summary_tables import table_json + + +def test_tables_round_measurements_and_keep_definitions_and_counts(): + table = Table('holdout/by_threshold', ('threshold', 'tp', 'precision'), ( + (0.07 + 1e-17 * 3, 12, 0.48774387193112), + (1 / 3, 7, 27.610799999999998), + )) + + assert json.loads(table_json(table))['data'] == [ + [0.07, 12, 0.487744], [0.3333333333, 7, 27.6108]] diff --git a/packages/kumo-engine/test/test_tracking_contract.py b/packages/kumo-engine/test/test_tracking_contract.py index 79df5a3b2..68daed1e8 100644 --- a/packages/kumo-engine/test/test_tracking_contract.py +++ b/packages/kumo-engine/test/test_tracking_contract.py @@ -3,6 +3,7 @@ # SPDX-License-Identifier: Apache-2.0 """Every tracking adapter records the same run records as a recorded fit.""" +import hashlib import json import threading import time @@ -437,7 +438,9 @@ def test_adapter_logs_each_summary_type(adapter, mlflow_adapter): {'path': 'graph/col_stats.json', 'type': 'table'}] dataset, = run.inputs.dataset_inputs assert (dataset.dataset.name, dataset.dataset.digest, - dataset.tags[0].value) == ('ents', '1790000000', 'training') + dataset.tags[0].value) == ('ents', hashlib.sha256(json.dumps( + ['ents', INPUTS['ents'], SUMMARIES[-1].schema], + sort_keys=True).encode()).hexdigest()[:8], 'training') assert json.loads(dataset.dataset.schema) == { 'mlflow_colspec': [{'name': 'ID', 'type': 'long'}]} assert json.loads(dataset.dataset.profile) == {'rows': 400} From 2e664c36eed71f277eabb48434be49056e11cea0 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 03:05:33 +0000 Subject: [PATCH 201/278] Publish a fitted bundle through a _publish helper Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/trainer/trainer.py | 35 +++++++++++-------- 1 file changed, 20 insertions(+), 15 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 913bb9a5c..ecc88c4c8 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -343,6 +343,22 @@ def _preflight(output_config: PredictionOutputConfig) -> None: write_mode=output_config.write_mode) +def _publish(bundle_dir: Path, metrics: Mapping[str, float], + export_config: ModelOutputConfig) -> PublishedModel: + spec = read_training_spec(bundle_dir) + return export_config.connector.sink_backend.publish_model( + export_config.connector, + str(bundle_dir), + path=export_config.path, + location=export_config.location, + metadata={ + name: str(getattr(spec, name)) + for name in ('spec_version', 'predictive_query', 'snapshot_id') + }, + metrics=metrics, + ) + + def _stage_model_bundle(artifacts: Mapping[str, str], dest: Path) -> Path: from kumo_core.artifacts import stage_model_bundle return stage_model_bundle(artifacts, dest) @@ -744,21 +760,10 @@ def _finish_fit( f"Publishing the model bundle without _provenance.json: " f"{error}", UserWarning, stacklevel=3) if export_config is not None: - spec = read_training_spec(training_result.bundle_dir) - published = export_config.connector.sink_backend.publish_model( - export_config.connector, - str(training_result.bundle_dir), - path=export_config.path, - location=export_config.location, - metadata={ - name: str(getattr(spec, name)) - for name in ('spec_version', 'predictive_query', - 'snapshot_id') - }, - metrics=training_result.metrics, - ) - training_result = replace(training_result, - published_model=published) + training_result = replace( + training_result, published_model=_publish( + training_result.bundle_dir, training_result.metrics, + export_config)) training_result = replace(training_result, _trainer=self) self._training_result = training_result From c1d7bd030d66325a5e21fba9b61b65babc33e34b Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 03:07:59 +0000 Subject: [PATCH 202/278] Split Trainer._fit into warm-start, validation, materialization, training and explain steps A pure extraction; two comments the step names now say are dropped. Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/trainer/trainer.py | 243 ++++++++++-------- 1 file changed, 133 insertions(+), 110 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index ecc88c4c8..27ff90ba6 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -21,6 +21,7 @@ Literal, Mapping, Optional, + Tuple, TypeVar, Union, cast, @@ -359,6 +360,87 @@ def _publish(bundle_dir: Path, metrics: Mapping[str, float], ) +def _warm_start_plan( + public_plan: ModelPlan, warm_start: Optional['Trainer'] +) -> Tuple[CoreModelPlan, ModelPlan, Optional[Dict[str, str]], Optional[str]]: + r"""The model plan with the warm start's inherited, and its artifacts.""" + from kumo_core.automl.runtime import inherit_warm_start_model_plan + from kumo_core.io.json import load_json + + model_plan = _to_core_model_plan(public_plan) + if warm_start is None: + return model_plan, public_plan, None, None + if warm_start._training_result is None: + raise ValueError("warm_start requires a fitted or loaded Trainer") + artifacts, bundle_id = read_bundle_artifact_paths( + warm_start._training_result.bundle_dir) + inherit_warm_start_model_plan( + model_plan, + load_json(artifacts['resolved_model_plan'], CoreModelPlan), + ) + return (model_plan, public_model_plan_from_internal(model_plan), + artifacts, bundle_id) + + +def _validate_model_plan(model_plan: CoreModelPlan, + workflow: TrainingWorkflowRequest, + train_table: TrainingTable) -> None: + from kumo_core.pquery import ( + query_is_temporal, + target_to_entity_paths, + to_plan_graph, + ) + from kumo_core.training.custom_table import contract + from kumo_core.training.export import read_spec, stage_parts + from kumo_core.training.workflow import resolve_model_plan + from kumo_core.types import TaskType + from kumoapi.pquery import QueryType + from pyarrow import parquet + + parts = stage_parts(read_spec(str(train_table.staging_dir))) + seed_part = next(part for paths in parts.values() for part in paths) + seed_schema = parquet.read_schema(train_table.staging_dir / seed_part) + entity_tables = [workflow.query.entity_table] + if (workflow.query.rhs_node_type is not None + and workflow.query.rhs_node_type not in entity_tables): + entity_tables.append(workflow.query.rhs_node_type) + resolve_model_plan( + model_plan, + workflow.search, + training_table_generation_plan=train_table.generation_plan, + task_type=TaskType(workflow.query.task_type), + query_type=QueryType(workflow.query.query_type), + has_train_table_time_col=query_is_temporal(workflow.query, + workflow.graph), + entity_table_names=entity_tables, + target_to_entity_paths=target_to_entity_paths( + workflow.query, workflow.graph), + has_train_table_weight_col=contract()['weight_output_column'] + in seed_schema.names, + graph=to_plan_graph(workflow.graph), + ) + + +def _explain(best_trainer: Any, run_id: str) -> None: + try: + best_trainer.explain( + explain_methods=[ + ExplainMethod.COLUMN_ANALYSIS, + ExplainMethod.GRADIENT, + ExplainMethod.SUBGRAPH_COLLECTOR, + ], + stage='test', + ) + aggregate_subgraph_features( + str(Path(best_trainer.subgraph_collector_dir()) / 'test'), + TrainingStage.TEST) + except Exception: + logger.exception( + 'Training succeeded, but post-training explanations ' + 'failed for run %s. The trained model remains available.', + run_id) + + def _stage_model_bundle(artifacts: Mapping[str, str], dest: Path) -> Path: from kumo_core.artifacts import stage_model_bundle return stage_model_bundle(artifacts, dest) @@ -520,47 +602,15 @@ def _fit( warm_start: Optional[Self], run_id: str, ) -> TrainingResult: - from kumo_core.automl.runtime import ( - inherit_warm_start_model_plan, - validate_warm_start_task_type, - ) - from kumo_core.io.json import load_json - from kumo_core.training.orchestration import ( - CallableTrainingOptions, - ResolvedTrainingInput, - run_callable_training, - ) + from kumo_core.automl.runtime import validate_warm_start_task_type from kumo_core.training.search import AutoMLSearchOptions - from kumo_core.training.table import GeneratedTrainTableResult - from kumo_core.training.workflow import ( - PreparedGraph, - materialize_generated_training_task, - resolve_model_plan, - ) + from kumo_core.training.workflow import PreparedGraph from kumo_core.types import TaskType - from kumo.engine.trainer._runtime import local_training_bindings - started_at = datetime.now(timezone.utc) run_dir = self.scratch_dir / 'runs' / run_id - - model_plan = _to_core_model_plan(self._model_plan) - warm_start_artifacts = None - warm_start_bundle_id = None - effective_plan = self._model_plan - if warm_start is not None: - if warm_start._training_result is None: - raise ValueError( - "warm_start requires a fitted or loaded Trainer") - warm_start_artifacts, warm_start_bundle_id = ( - read_bundle_artifact_paths( - warm_start._training_result.bundle_dir)) - inherit_warm_start_model_plan( - model_plan, - load_json(warm_start_artifacts['resolved_model_plan'], - CoreModelPlan), - ) - effective_plan = public_model_plan_from_internal(model_plan) + model_plan, effective_plan, warm_start_artifacts, warm_start_bundle_id = ( + _warm_start_plan(self._model_plan, warm_start)) workflow = TrainingWorkflowRequest.from_spec( graph=snapshot.graph_config, predictive_query=train_table.query, @@ -582,41 +632,37 @@ def _fit( TaskType(workflow.query.task_type), TaskType(source_schema['task_type']), ) - from kumo_core.pquery import ( - query_is_temporal, - target_to_entity_paths, - to_plan_graph, + _validate_model_plan(model_plan, workflow, train_table) + materialized = self._materialize_training(workflow, train_table, + run_dir, effective_plan) + training_result, fit_mode = self._train(materialized, run_dir, run_id, + warm_start_artifacts) + _collect_garbage('fit') + return self._finish_fit( + training_result, + predictive_query=train_table.query, + snapshot_id=snapshot.id, + graph_config=snapshot.graph_config, + training_table_generation=train_table.generation_plan, + export_config=export_config, + model_plan=effective_plan, + provenance=FitRecord( + train_table=train_table, + warm_start=None if warm_start is None else bundle_manifest( + warm_start._training_result.bundle_dir, + warm_start_bundle_id), + started_at=started_at, fit_mode=fit_mode), ) - from kumo_core.training.custom_table import contract - from kumo_core.training.export import read_spec, stage_parts - from kumoapi.pquery import QueryType - from pyarrow import parquet - - parts = stage_parts(read_spec(str(train_table.staging_dir))) - seed_part = next(part for paths in parts.values() for part in paths) - seed_schema = parquet.read_schema(train_table.staging_dir / seed_part) - entity_tables = [workflow.query.entity_table] - if (workflow.query.rhs_node_type is not None - and workflow.query.rhs_node_type not in entity_tables): - entity_tables.append(workflow.query.rhs_node_type) - # Validate a copy with query/schema facts before extending the index. - # Materialization still resolves defaults and data-derived fanouts. - resolve_model_plan( - model_plan, - workflow.search, - training_table_generation_plan=train_table.generation_plan, - task_type=TaskType(workflow.query.task_type), - query_type=QueryType(workflow.query.query_type), - has_train_table_time_col=query_is_temporal(workflow.query, - workflow.graph), - entity_table_names=entity_tables, - target_to_entity_paths=target_to_entity_paths( - workflow.query, workflow.graph), - has_train_table_weight_col=contract()['weight_output_column'] - in seed_schema.names, - graph=to_plan_graph(workflow.graph), + + def _materialize_training(self, workflow: TrainingWorkflowRequest, + train_table: TrainingTable, run_dir: Path, + effective_plan: ModelPlan) -> Any: + from kumo_core.training.table import GeneratedTrainTableResult + from kumo_core.training.workflow import ( + materialize_generated_training_task, ) - ingest_tt_id = f'tt-{run_id}' + + ingest_tt_id = f'tt-{workflow.run_id}' lag_timesteps = inferred_or_missing_type_to_none( train_table.generation_plan.lag_timesteps) year_over_year = inferred_or_missing_type_to_none( @@ -625,7 +671,7 @@ def _fit( query=workflow.query, graph=workflow.graph, graph_index_path=workflow.prepared_graph.index_path, - run_id=run_id, + run_id=workflow.run_id, artifact_dir=str(run_dir), tt_id=ingest_tt_id, required_stages=train_table.required_stages, @@ -647,11 +693,25 @@ def _fit( ), tt_id=ingest_tt_id, ) - del task_dataset if (materialized.task.split.stage_counts.get('test') == 0 and not materialized.task.task_type.is_link_pred): raise ValueError('The test split is empty: shorten the query ' 'window or add data.') + return materialized + + def _train( + self, materialized: Any, run_dir: Path, run_id: str, + warm_start_artifacts: Optional[Mapping[str, str]] + ) -> Tuple[TrainingResult, Optional[str]]: + r"""The trained model and its fit mode.""" + from kumo_core.training.orchestration import ( + CallableTrainingOptions, + ResolvedTrainingInput, + run_callable_training, + ) + + from kumo.engine.trainer._runtime import local_training_bindings + training_input = ResolvedTrainingInput( task=materialized.task, model_plan=materialized.model_plan, @@ -667,29 +727,8 @@ def _fit( accelerator=self._accelerator), options=CallableTrainingOptions(), ) - if not result.explain_disabled: - try: - # The shared trainer owns task support, checkpoint selection - # (refit_trainval when available), and artifact generation. - result.best_trainer.explain( - explain_methods=[ - ExplainMethod.COLUMN_ANALYSIS, - ExplainMethod.GRADIENT, - ExplainMethod.SUBGRAPH_COLLECTOR, - ], - stage='test', - ) - aggregate_subgraph_features( - str( - Path(result.best_trainer.subgraph_collector_dir()) / - 'test'), TrainingStage.TEST) - except Exception: - logger.exception( - 'Training succeeded, but post-training explanations ' - 'failed for run %s. The trained model remains available.', - run_id) - + _explain(result.best_trainer, run_id) holdout_path, holdout_truncated = _published_holdout( result.best_trainer) fit_mode = result.best_trainer.published_fit_mode() @@ -701,23 +740,7 @@ def _fit( holdout_path=holdout_path, holdout_truncated=holdout_truncated, ) - del result - _collect_garbage('fit') - return self._finish_fit( - training_result, - predictive_query=train_table.query, - snapshot_id=snapshot.id, - graph_config=snapshot.graph_config, - training_table_generation=train_table.generation_plan, - export_config=export_config, - model_plan=effective_plan, - provenance=FitRecord( - train_table=train_table, - warm_start=None if warm_start is None else bundle_manifest( - warm_start._training_result.bundle_dir, - warm_start_bundle_id), - started_at=started_at, fit_mode=fit_mode), - ) + return training_result, fit_mode def _finish_fit( self, From 4c92ca44022a8e04a7a935170e418068c82ca9d8 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 03:08:08 +0000 Subject: [PATCH 203/278] Word the MLOps docstrings as the upper PRs do Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_mlops.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_mlops.py b/packages/kumo-engine/kumo/engine/_mlops.py index e1f140169..d1e1f14cc 100644 --- a/packages/kumo-engine/kumo/engine/_mlops.py +++ b/packages/kumo-engine/kumo/engine/_mlops.py @@ -39,11 +39,11 @@ def uri(self, relative: str) -> str: @dataclass(frozen=True) class MLOps: - r"""Where artifacts are stored. + r"""MLOps settings. Args: - artifact_root: a URI or path, or a ``(connector, path)`` pair. - store: artifact kinds to store; ``True`` for all, ``False`` for none. + artifact_root: a URI, local path or ``(connector, path)`` pair. + store: artifact kinds to store; ``True`` for all, ``False`` none. on_error: ``'warn'`` or ``'raise'`` when MLOps fails. """ @@ -153,7 +153,7 @@ def _require_writable(root: _Root) -> None: @contextlib.contextmanager def mlops(settings: MLOps) -> Iterator[MLOps]: - r"""Store artifacts of calls in this block under ``settings``.""" + r"""Turns MLOps on, as ``settings`` says, for the calls in the block.""" if not isinstance(settings, MLOps): raise TypeError('kumo.mlops() takes a kumo.MLOps') token = _CURRENT.set(settings) From be6c8f3cbd6bedba7cb1ea37a751232c47bc0947 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 03:08:45 +0000 Subject: [PATCH 204/278] Tag only models with their location when naming manifest artifacts Snapshots and training tables no longer get a made-up kumo._uri tag; a model keeps kumo.bundle_uri. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_manifest.py | 3 +-- packages/kumo-engine/test/test_artifacts.py | 8 ++++---- 2 files changed, 5 insertions(+), 6 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_manifest.py b/packages/kumo-engine/kumo/engine/_manifest.py index 49451ca02..2fb4d7d05 100644 --- a/packages/kumo-engine/kumo/engine/_manifest.py +++ b/packages/kumo-engine/kumo/engine/_manifest.py @@ -171,8 +171,7 @@ class Manifest(Record): def as_artifact(self, key: str) -> Artifact: return Artifact(self.kind, self.id, self.location, self.IS_DATASET, - self.URI_TAG or f'kumo.{self.kind}_uri', - version=self.id, version_kind='exact') + self.URI_TAG, version=self.id, version_kind='exact') def artifacts(manifest: Manifest) -> List[Artifact]: diff --git a/packages/kumo-engine/test/test_artifacts.py b/packages/kumo-engine/test/test_artifacts.py index 0d51cc262..d6cbb3cb2 100644 --- a/packages/kumo-engine/test/test_artifacts.py +++ b/packages/kumo-engine/test/test_artifacts.py @@ -50,13 +50,13 @@ def test_a_new_kind_names_its_artifacts_without_other_changes(): assert artifacts(report) == [ Artifact('report', 'r', 's3://root/reports/r', True, - 'kumo.report_uri', version='r', version_kind='exact'), + version='r', version_kind='exact'), Artifact('snapshot', 'g@d', 's3://root/snapshots/g/d', True, - 'kumo.snapshot_uri', version='g@d', version_kind='exact'), + version='g@d', version_kind='exact'), Artifact('source_table', 'orders', 's3://data/orders', True, connector='S3', version='7', version_kind='exact'), Artifact('model', 'm', 's3://root/models/m', False, 'kumo.bundle_uri', version='m', version_kind='exact'), - Artifact('snapshot', 'p', None, True, 'kumo.snapshot_uri', - version='p', version_kind='exact'), + Artifact('snapshot', 'p', None, True, version='p', + version_kind='exact'), ] From 10569d5765384018f948c1a8d0aa57f840cfa4eb Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 03:09:38 +0000 Subject: [PATCH 205/278] Name the changed columns in a schema change Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/graph.py | 13 ++++++++++++- packages/kumo-engine/test/test_graph_mismatch.py | 2 +- 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index 2b4a99ece..3f1b40c3c 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -234,7 +234,8 @@ def schema_changes(before: Mapping[str, Any], ] for name in sorted(old.keys() & new.keys()): parts = sorted( - _SCHEMA_PARTS.get(key, key) + _SCHEMA_PARTS.get(key, key) + (_changed_columns( + old[name][key], new[name][key]) if key == 'cols' else '') for key in old[name].keys() | new[name].keys() if old[name].get(key) != new[name].get(key)) if parts: @@ -242,6 +243,16 @@ def schema_changes(before: Mapping[str, Any], return changes +def _changed_columns(before: List[str], after: List[str]) -> str: + old, new = ({ + column['name']: column + for column in map(json.loads, side or []) + } for side in (before, after)) + names = sorted(name for name in old.keys() | new.keys() + if old.get(name) != new.get(name)) + return f" ({', '.join(names)})" if names else '' + + def report_other_graph(message: str, before: Optional[Mapping[str, Any]], after: Mapping[str, Any], stacklevel: int) -> None: changes = None if before is None else schema_changes(before, after) diff --git a/packages/kumo-engine/test/test_graph_mismatch.py b/packages/kumo-engine/test/test_graph_mismatch.py index d36e4a3db..df86cbb68 100644 --- a/packages/kumo-engine/test/test_graph_mismatch.py +++ b/packages/kumo-engine/test/test_graph_mismatch.py @@ -82,4 +82,4 @@ def config(location, columns, tables=('users', )): assert schema_changes( config('dev', [age]), config('dev', [{**age, 'dtype': 'float64'}], tables=('users', 'x')) - ) == ['table x added', 'users columns'] + ) == ['table x added', 'users columns (age)'] From 018033538b42253f032fdf57b5d0ea875538aafa Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 03:12:43 +0000 Subject: [PATCH 206/278] Lay out the MLOps doc in the upper PRs' order and words Loading moves to Reading artifacts back, scoring to Score, export and load under Reference; lineage and moving snapshots are condensed. on_error is documented. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/docs/mlops.md | 125 ++++++++++------------------- 1 file changed, 43 insertions(+), 82 deletions(-) diff --git a/packages/kumo-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index 338c2eb73..cb2cde19a 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -10,118 +10,79 @@ import kumo.engine as kumo cfg = kumo.MLOps("artifacts") # or s3://..., gs://..., /Volumes/..., (connector, "@DB.SCHEMA.STAGE/kumo") with kumo.mlops(cfg): snapshot = graph.snapshot() - result = kumo.Trainer(pquery.suggest_model_plan(snapshot)).fit( - pquery.generate_training_table(snapshot), snapshot=snapshot) -# Later, on any machine that reaches the root: -snapshot = kumo.GraphSnapshot.load_from_connector(cfg, id=result.snapshot_id) -trainer = kumo.Trainer.load_from_connector(cfg, path=f"models/{result.id}") + result = kumo.Trainer(pquery.suggest_model_plan(snapshot)).fit(pquery.generate_training_table(snapshot), + snapshot=snapshot) ``` ## How it fits together -| artifact | id | stored at (MLOps on) | -|---|---|---| -| snapshot | `snapshot.id` | `/snapshots///` | -| training table | `table.id` | `/training_tables//` | -| model | `result.id` (the bundle id) | `/models//` | +- The **artifact root** holds every artifact a block produces, each in its own directory with a `_provenance.json` manifest that embeds its parents' manifests. -Each carries a `_provenance.json` manifest that embeds its parents' manifests (see [Lineage](#lineage)). +| id | names | where to find it | +|---|---|---| +| `snapshot.id` | a graph snapshot | `/snapshots///`; `result.snapshot_id` | +| `table.id` | a training table | `/training_tables//` | +| `result.id` | a model (its bundle id) | `/models//` | ## Turning MLOps on -`kumo.MLOps` says where artifacts go. Use it as a block, or pass it as `mlops=` on a single call: +`kumo.MLOps` says where artifacts go. Use it as a block, or pass it as `mlops=` on a single call (`fit`, `graph.snapshot`, `generate_training_table`): ```python -cfg = kumo.MLOps("s3://bucket/kumo") +cfg = kumo.MLOps("s3://bucket/kumo", on_error="raise") with kumo.mlops(cfg): snapshot = graph.snapshot() trainer.fit(table, snapshot=snapshot, mlops=cfg) ``` -- **`artifact_root`** is required: a URI (`s3://`, `gs://`, `file://`, or a local path, including `/Volumes/...` on a Databricks cluster) or a `(connector, path)` pair for a store a URI cannot reach, such as `(DatabricksConnector(...), "/Volumes////kumo")` or `(SnowflakeConnector(...), "@DB.SCHEMA.STAGE/kumo")`. A connector whose sink cannot hold files there raises when `kumo.MLOps` is built. -- **`store`** names the artifact kinds stored when they are produced: `snapshot`, `training_table` and `model` (default: all). +- **`artifact_root`** (required): a URI (`s3://`, `gs://`, `file://`, or a local path, including `/Volumes/...` on a Databricks cluster) or a `(connector, path)` pair for a store a URI cannot reach: a volume through a `DatabricksConnector` (`"/Volumes////kumo"`) or a named stage through a `SnowflakeConnector` (`"@DB.SCHEMA.STAGE/kumo"`) once the installed sink holds files there. +- **`store`**: the artifact kinds stored when produced: `snapshot`, `training_table` and `model` (default: all). +- **`on_error`**: `"warn"` (default) warns about a failed store and returns; `"raise"`, for scheduled jobs, raises once the call has finished, with the failure as the `__cause__`. After `fit`, `trainer.result` still holds the model. - **Precedence:** a call's `mlops=` replaces the enclosing block, and an inner block replaces an outer one. A block covers the calling thread or asyncio task. -- **Off by default.** Without a block or `mlops=`, every call behaves as it always has, and plain exports (`export(config)`, `export_config=`) never write to the artifact root. +- **Off by default.** Without a block or `mlops=`, every call behaves as it always has; plain exports (`export(config)`, `export_config=`) never write to the artifact root. -## Exporting and loading artifacts - -Every artifact has one way out, `.export(OutputConfig(...))`, and one way back, `.load_from_connector(...)`. Pass the `kumo.MLOps` in place of the connector to read from the artifact root, with paths relative to it: - -| Artifact | Export | Load | -|---|---|---| -| Snapshot | `snapshot.export(SnapshotOutputConfig(connector, path))`, returns the location | `GraphSnapshot.load_from_connector(connector, path=location)`, or `(cfg, id=snapshot_id)` | -| Training table | `table.export(TrainingTableOutputConfig(connector, name))` | `TrainingTable.load_from_connector(connector, path=name, pquery=pquery, snapshot=snapshot)`, or `(cfg, path=f"training_tables/{id}", ...)` | -| Model | `result.export(ModelOutputConfig(connector, path=... or location=...))` | `Trainer.load_from_connector(connector, path=... or location=...)`, or `(cfg, path=f"models/{id}")` | +### Reading artifacts back -`fit(export_config=...)` is shorthand for fitting, then exporting. Loading needs a file destination (S3, GCS, local files, a volume or stage). Exporting again to a destination that already holds a complete copy of the same artifact reuses it (logged at info) instead of uploading again. +Pass the `kumo.MLOps` where a connector goes, with paths relative to the artifact root: -## Reference - -### Lineage +```python +snapshot = kumo.GraphSnapshot.load_from_connector(cfg, id=result.snapshot_id) +trainer = kumo.Trainer.load_from_connector(cfg, path=f"models/{result.id}") +table = kumo.TrainingTable.load_from_connector(cfg, path=f"training_tables/{table_id}", pquery=pquery, snapshot=snapshot) +``` -Every artifact carries a `_provenance.json` manifest, metadata only. The leading underscore keeps Spark and Hive Parquet readers from treating it as data. +A loaded model is a `Trainer`; its `result` is the `TrainingResult` (`.id`, `.manifest`, `.metrics`, `.holdout_df()`). -A graph snapshot's manifest, next to its index, records each source table's connector, location and version, plus the engine build that wrote it. `graph.snapshot(record_sources=False)` keeps only connector and table names, plus a SHA-256 of each full location so a source moved elsewhere is still noticed, and keeps Custom SQL only as a hash. +## Score -#### Training tables +`predict(snapshot)` scores the trained query's population at its prediction time; `entities=` scores given rows; `predict(prediction_table, snapshot=...)` scores a table generated beforehand. A scoring job needs no `Graph` code: `GraphSnapshot.load_from_connector(cfg, id=...)` and `Trainer.load_from_connector(cfg, ...)`, then `predict`. It warns when the snapshot was taken of a different graph than the model's, naming the schema changes. -A generated or caller-supplied training table carries `_provenance.json` in its staging directory: -- the query, the effective generation plan and the task type; -- the split seed; -- the snapshot it was generated from (that snapshot's manifest, embedded); -- for `TrainingTable.from_source_table`, the source table's connector, location and version. +## Reference -`TrainingTable.id` is deterministic: the same snapshot, query, plan, split seed and engine build give the same id. A caller-supplied source without a version gets a random id. An export to files (`TrainingTable.export`) writes the manifest next to the Parquet parts. `TrainingTable.load_from_connector(connector, path=..., pquery=..., snapshot=...)` materializes an export again. +### Exporting and loading artifacts -With MLOps on, each table is stored once under `/training_tables//`. +Every artifact has one way out and one way back. A `kumo.MLOps` stands in for the connector to read the artifact root. -#### Models +| Artifact | Export | Load | +|---|---|---| +| Snapshot | `snapshot.export(SnapshotOutputConfig(connector, path))`, returns the location | `GraphSnapshot.load_from_connector(connector, path=location)`, `(cfg, id=...)` | +| Training table | `table.export(TrainingTableOutputConfig(connector, name))` | `TrainingTable.load_from_connector(connector, path=name, pquery=pquery, snapshot=snapshot)` | +| Model | `result.export(ModelOutputConfig(connector, path=... or location=...))`, or `fit(export_config=...)` | `Trainer.load_from_connector(connector, path=... or location=...)` | -Every fitted bundle carries `_provenance.json` (`ModelManifest`): -- its `id`, the bundle id (a content hash that leaves the manifest file out); -- when it was fitted and with which fit mode; -- its parents: the training table's manifest, embedded with that table's snapshot, and the warm-start model's manifest, if any. +Loading needs a file destination; a table export is for consumers outside kumo-engine. Exporting again to a destination that holds a complete copy of the same artifact reuses it (logged at info). `ModelOutputConfig(connector, location=...)` publishes to a connector's registry outside a block. -An embedded model keeps only the ids of its own model parents, so a long chain of warm-start retrains does not grow each manifest. `fit` returns a new `TrainingResult` per call (`.id`, `.manifest`); `Trainer.result` is the latest fit's result or the loaded bundle's (`trainer.result.id`, `trainer.result.manifest`). +### Lineage and manifests -With MLOps on, each bundle is stored under `/models//`. A complete copy that is already there is kept; a partial one is replaced. +Every artifact carries `_provenance.json`, metadata only (the underscore keeps Spark and Hive readers off it). Parents are embedded; a parent without a manifest (a bundle from an older release) is an id stub. -The bundle also carries the published model's holdout predictions (`holdout.parquet`, with its truncation flag), so `result.holdout_df()` works on a model loaded with `Trainer.load` or `Trainer.load_from_connector` (a path or a registry), not only on the result of `fit`. Bundles fitted before this have no copy, and `holdout_df()` raises `FileNotFoundError` for them. +- **Snapshot:** each source table's connector, location and version, and the engine build. `graph.snapshot(record_sources=False)` keeps only connector and table names, a SHA-256 of each location, and Custom SQL as a hash. +- **Training table:** the query, generation plan, task type and split seed, its snapshot, and for `TrainingTable.from_source_table` the source table. `table.id` is deterministic over snapshot, query, plan, seed and engine build; a source without a version gets a random id. +- **Model** (`ModelManifest`, `result.manifest`): the bundle id (a content hash excluding the manifest), fit times and fit mode, and the training table and warm-start model (an embedded model keeps only ids of its own model parents). The bundle carries the holdout predictions (`holdout.parquet`), so `holdout_df()` works on a loaded model. ### Moving snapshots between machines -To build an index on one machine and train or score on another, export it to a file destination and load it there: - -```python -from kumo.engine import GraphSnapshot, S3Connector, SnapshotOutputConfig - -location = graph.snapshot().export(SnapshotOutputConfig(S3Connector(), "s3://bucket/snapshots")) -# On the training machine: -snapshot = GraphSnapshot.load_from_connector(S3Connector(), path=location) -``` - -- **Layout:** `export` writes `///`, the two halves of `snapshot.id`, and returns that location. It holds `kumo_snapshot.json` and one directory per export generation (`index/` and `_provenance.json`). -- **Transfer:** files move in parallel with a CRC32 per range. Each export uploads into a new generation and commits by writing `kumo_snapshot.json` last, naming that generation. Until then the previous export stays loadable, and concurrent exports never mix files: the last commit wins whole. Superseded generations, and abandoned ones over a day old, are removed best-effort. A load either completes or leaves nothing; if a newer export commits while it downloads, it switches to that one. -- **Reuse:** loading again reuses the local copy only if it holds this export's bytes (the same range checksums) and passes DiskGraph's integrity check; otherwise it is downloaded again. Exporting again skips the upload when the destination already holds a complete export (its `kumo_snapshot.json` and `_DONE`) with the local index's bytes, and logs that it reused it. -- **Staleness:** `snapshot.source_changes(graph)` reports, per table of `graph`, whether its source is `unchanged`, `changed` or `unknown` since the snapshot was built. `load_from_connector(..., graph=graph)` warns when the snapshot was taken of another graph, naming the schema changes (tables, columns, types, keys, edges), or when any of its tables changed. A graph id also covers each table's connector and location, so a graph that differs only there also warns, saying the schema is the same. -- **Manifest:** the snapshot's `_provenance.json` travels with the export. -- **Trust:** load only from paths whose writers you trust. Checksums detect corruption, not tampering. -- **Destinations:** any destination whose sink holds files at the path works: S3, GCS, local files, or a volume or stage once those sinks are installed. A warehouse connector raises. - -With MLOps on (`graph.snapshot(mlops=...)` or a `kumo.mlops` block), each snapshot is exported under `/snapshots/`, unless a complete export of it is already there. - -### Scoring without the graph - -A scoring job does not need the training job's `Graph` code. It loads a snapshot and the model from the artifact root, then scores: - -```python -snapshot = GraphSnapshot.load_from_connector(cfg, id=snapshot_id) -trainer = Trainer.load_from_connector(cfg, path=f"models/{bundle_id}") -result = trainer.predict(snapshot, binary_classification_threshold=0.5) -``` - -- `predict(snapshot)` reuses the trained query and table-generation plan, and scores its population at `prediction_time`. -- `entities=` scores given rows instead. -- It warns when the snapshot was taken of a different graph than the one the model was trained on, naming the schema changes, for example `Schema changes: events edges.` A graph that differs only in table connectors or locations (a dev graph and a prod graph of the same schema) warns that the schema is the same. The model records its training graph from this release on; for an older model the warning cannot say what differs. -- `predict(prediction_table, snapshot=...)` scores a prediction table generated beforehand, as before. -- `fit` results score too: `result.predict(snapshot, ...)`. +- **Layout:** `export` writes `///`, holding `kumo_snapshot.json` and one directory per export generation (`index/` and `_provenance.json`). +- **Transfer:** parallel, with a CRC32 per range. Each export commits by writing `kumo_snapshot.json` last, so the previous export stays loadable and concurrent exports never mix: the last commit wins. Superseded generations are removed best-effort. A load completes or leaves nothing, and follows a newer export that commits meanwhile. +- **Reuse:** a load reuses the local copy only if it holds the export's bytes and passes DiskGraph's integrity check; an export skips the upload when the destination holds a complete export with the same bytes. +- **Staleness:** `snapshot.source_changes(graph)` reports each table as `unchanged`, `changed` or `unknown`; `load_from_connector(..., graph=graph)` warns about another graph's schema or changed tables. +- **Trust:** load only from paths whose writers you trust; checksums detect corruption, not tampering. A warehouse connector cannot hold a snapshot. From c0d7000b5995695d7d857d0ea81ac7953a53fb54 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 03:14:03 +0000 Subject: [PATCH 207/278] Check the recorded fit mode of each published model Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/test/test_holdout_df.py | 1 + 1 file changed, 1 insertion(+) diff --git a/packages/kumo-engine/test/test_holdout_df.py b/packages/kumo-engine/test/test_holdout_df.py index c159967d2..e72322fb3 100644 --- a/packages/kumo-engine/test/test_holdout_df.py +++ b/packages/kumo-engine/test/test_holdout_df.py @@ -104,6 +104,7 @@ def test_holdout_is_the_published_models(graph, refit_trainval, refit_full, explain, model): result = _fit(graph, QUERY, refit_trainval, refit_full, explain) + assert result.manifest.fit_mode == (model or 'refit_full') if model is None: assert result.holdout_path is None else: From 25b67ddf9d4fd554b99b52fa826ff8bf7dcf1294 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 03:19:42 +0000 Subject: [PATCH 208/278] Log a fit's holdout as test/ and register those metrics Trackers and registries get the shipped model's holdout (before_refit/ for the pre-refit one); fits drop seconds and the alias tag. Snowflake keeps each trial's best tune metric under TRIAL_/ plus a run's duration_seconds, MLflow times runs from the record, and a Snowflake registry version defaults to the next V and records the holdout. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_tracking.py | 39 ++++++++-- .../kumo/engine/integrations/mlflow.py | 18 ++++- .../kumo/engine/integrations/snowflake.py | 59 ++++++++------ .../kumo/engine/trainer/_run_spec.py | 50 ++++++++---- .../kumo/engine/trainer/trainer.py | 23 +++--- .../kumo-engine/test/test_mlflow_model.py | 18 +++-- .../kumo-engine/test/test_mlflow_tracking.py | 10 +-- .../test/test_snowflake_tracking.py | 8 +- .../test/test_tracking_contract.py | 48 ++++++++---- .../kumo-engine/test/test_trainer_registry.py | 4 +- .../src/kumo_snowflake/registry.py | 65 ++++++++++++---- .../kumo-snowflake/tests/test_registry.py | 78 +++++++++++++++++-- 12 files changed, 300 insertions(+), 120 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_tracking.py b/packages/kumo-engine/kumo/engine/_tracking.py index 3ca520c8f..c44dbd93e 100644 --- a/packages/kumo-engine/kumo/engine/_tracking.py +++ b/packages/kumo-engine/kumo/engine/_tracking.py @@ -61,6 +61,7 @@ class RunRecord: summaries: Tuple[Summary, ...] = () model_label: Optional[str] = None started: Optional[datetime] = None + finished: Optional[datetime] = None name_suffix: Optional[str] = None @@ -82,20 +83,41 @@ class RunSpec: error: Optional[str] = None model_label: Optional[str] = None started: Optional[datetime] = None + finished: Optional[datetime] = None name_suffix: Optional[str] = None +def metric_name(metric: str) -> str: + r"""The plan's acc is logged as accuracy.""" + return 'accuracy' if metric == 'acc' else metric + + def logged_name(name: str) -> str: r"""``val_auroc`` as ``val/auroc``.""" best = name.startswith('best_') split, _, metric = name.removeprefix('best_').partition('_') - if split not in ('train', 'val', 'test') or not metric: + if split not in ('train', 'val') or not metric: return name - return f"{split}/{'best_' if best else ''}{metric}" - - -def logged_names(metrics: Mapping[str, Any]) -> Dict[str, Any]: - return {logged_name(name): value for name, value in metrics.items()} + return f"{split}/{'best_' if best else ''}{metric_name(metric)}" + + +def logged_names(metrics: Mapping[str, Any], + refit_full: bool = False) -> Dict[str, Any]: + r"""Tracker names of metrics; test/ is the shipped model's holdout.""" + refit = any(name.startswith('test_refit_') for name in metrics) + named = {} + for name, value in metrics.items(): + if not name.startswith('test_'): + named[logged_name(name)] = value + continue + metric = name.removeprefix('test_') + latest = metric.startswith('refit_') == refit + metric = metric_name(metric.removeprefix('refit_')) + if not latest and refit_full: + continue + named[f'test/{metric}' if latest and not refit_full else + f'before_refit/{metric}'] = value + return named def by_step(steps: Mapping[str, Sequence[Tuple[int, float]]] @@ -127,6 +149,7 @@ def run_record(spec: RunSpec) -> RunRecord: summaries=tuple(spec.summaries), model_label=spec.model_label, started=spec.started, + finished=spec.finished, name_suffix=spec.name_suffix, ) @@ -158,8 +181,8 @@ def trial_spec(events: Sequence[Event]) -> RunSpec: error=None if 'error' not in event.params else str( event.params['error'])) metrics.update({ - name if name == 'seconds' else f'best_{name}': value - for name, value in event.metrics.items() + f'best_{name}': value + for name, value in event.metrics.items() if name != 'seconds' }) return RunSpec( kind='trial', diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index 98969d532..0e6bd2636 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -14,6 +14,7 @@ import time import warnings from contextlib import ExitStack, contextmanager +from datetime import datetime from pathlib import Path from typing import ( TYPE_CHECKING, @@ -283,6 +284,7 @@ def start(self, record: RunRecord) -> None: else record.name) run = client.create_run( experiment_id, run_name=name, + start_time=_millis(record.started), tags=context.resolve_tags(tags)) self._parent = self.run_id = run.info.run_id self.logged_run_name = run.info.run_name @@ -354,8 +356,10 @@ def log_model( def finish(self, record: Optional[RunRecord], error: Optional[BaseException]) -> None: status = 'FAILED' - fit_run = (record is not None and self._parent is not None - and record.tags.get('kumo.kind') == 'fit') + kind = None if record is None else record.tags.get('kumo.kind') + started = self._parent is not None + fit_run = started and kind == 'fit' + ended = None if record is None else _millis(record.finished) try: if record is not None and self._parent is not None: for epoch, metrics in by_step(record.steps): @@ -374,7 +378,7 @@ def finish(self, record: Optional[RunRecord], _quietly(self._client.set_terminated, run_id, 'FAILED') if self._parent is not None: _quietly(self._client.set_terminated, self._parent, - status) + status, ended) self._reset() hint = _sagemaker_ui_hint(self.target.tracking_uri) if fit_run and hint: @@ -507,6 +511,10 @@ def _digest(name: str, source: Mapping[str, Optional[str]], default=str).encode()).hexdigest()[:8] +def _millis(moment: Optional[datetime]) -> Optional[int]: + return None if moment is None else int(moment.timestamp() * 1000) + + _BATCH_METRICS = 1000 _SYSTEM_METRICS_FAILURES = 6 _COLUMN_TYPES = {'int64': 'long', 'int32': 'integer', 'float64': 'double', @@ -634,6 +642,7 @@ def log_model( from kumo.engine._manifest import artifact_summaries from kumo.engine._mlops import snapshot_uri from kumo.engine.trainer._provenance import bundle_manifest + from kumo.engine.trainer._run_spec import registry_metrics mlflow = import_provider('mlflow', 'mlflow') result = trainer.result @@ -666,7 +675,8 @@ def log_model( bundle_manifest(result.bundle_dir, result.id), snapshot, run_id, experiment_id, name, snapshot_uri(serving_snapshot)) - tracking._log_metrics(run_id, result.metrics, model_id=info.model_id) + tracking._log_metrics(run_id, registry_metrics(result), + model_id=info.model_id) if created: with target.scope(): client.set_terminated(run_id, 'FINISHED') diff --git a/packages/kumo-engine/kumo/engine/integrations/snowflake.py b/packages/kumo-engine/kumo/engine/integrations/snowflake.py index 3f1107411..777a53b33 100644 --- a/packages/kumo-engine/kumo/engine/integrations/snowflake.py +++ b/packages/kumo-engine/kumo/engine/integrations/snowflake.py @@ -8,8 +8,9 @@ import re import tempfile from contextlib import ExitStack +from datetime import datetime, timezone from pathlib import Path -from typing import Any, Dict, List, Mapping, Optional, Sequence, Set, Tuple +from typing import Any, List, Mapping, Optional, Sequence, Set, Tuple from kumo_core.instrumentation import format_error from kumo_core.sources.sql import resolve_snowflake_identifier @@ -79,8 +80,9 @@ def _reset(self) -> None: self._tracking: Any = None self._run: Any = None self._run_context = ExitStack() - self._trial_tags: Dict[Tuple[int, int], Set[str]] = {} + self._tune = 'loss' self._metrics: Set[str] = set() + self._started = _now() self._params = 0 self._start_tags: Set[str] = set() self._artifacts: Set[str] = set() @@ -145,35 +147,32 @@ def start(self, record: RunRecord) -> None: for row in tracking._sql_client.show_runs_in_experiment( experiment_name=experiment.name) ] + self._started = record.started or _now() self._run = tracking.start_run(self._new_run_name( existing, _identifier(record_name(record, self.run_name)))) self.run_id = self.logged_run_name = str(self._run.name) if self.native_model_lineage: self._run_context.enter_context(self._run) self._start_tags = set(record.tags) + self._tune = record.tags.get('kumo.tune_metric', 'loss') self._log_params({**record.params, **record.tags}) self._log_artifacts(record.artifacts) def on_trial(self, record: RunRecord, event: Event) -> None: - prefix = f'{record.name}/' - key = (event.trial or 0, int(event.params['attempt'])) + if self._run is None: + return + attempt = int(event.params['attempt']) + prefix = f'TRIAL_{event.trial or 0}' + ( + f'_ATTEMPT_{attempt}/' if attempt > 1 else '/') if event.name == 'trial_started': - self._trial_tags[key] = set(record.tags) - self._log_params({prefix + k: v - for k, v in {**record.params, - **record.tags}.items()}) - elif event.name == 'epoch_ended': - self._log_metrics({prefix + k: v - for k, v in event.metrics.items()}, - _EPOCH_METRICS, step=event.epoch or 0) + self._log_params({prefix + k: v for k, v in record.params.items()}) elif event.name == 'trial_ended': - self._log_metrics({prefix + k: v - for k, v in record.metrics.items()}, - _SUMMARY_METRICS) - logged = self._trial_tags.pop(key, set()) - self._log_params({prefix + k: v - for k, v in record.tags.items() - if k not in logged}) + best = f'val/best_{self._tune}' + self._log_metrics({prefix + best: value + for name, value in record.metrics.items() + if name == best}, _SUMMARY_METRICS) + self._log_params({prefix + k: v for k, v in record.tags.items() + if k in ('kumo.status', 'kumo.error')}) def finish(self, record: Optional[RunRecord], error: Optional[BaseException]) -> None: @@ -188,11 +187,10 @@ def finish(self, record: Optional[RunRecord], if k not in self._start_tags}, _MAX_PARAMS) self._log_artifacts(record.artifacts) self._log_summaries(record.summaries) - dropped = {DROPPED_METRICS: len(self._dropped_metrics), - DROPPED_PARAMS: self._dropped_params} - if any(dropped.values()): - self._tracking.log_params( - {name: n for name, n in dropped.items() if n}) + finished = record.finished or _now() + self._log_metrics({'duration_seconds': ( + finished - self._started).total_seconds()}, _MAX_METRICS) + self._log_dropped() if error is None: status = 'FINISHED' finally: @@ -240,6 +238,15 @@ def _log_artifacts(self, artifacts: Mapping[str, str]) -> None: str(path), posixpath.dirname(name) or None) self._artifacts.add(name) + def _log_dropped(self) -> None: + dropped = { + DROPPED_METRICS: len(self._dropped_metrics), + DROPPED_PARAMS: self._dropped_params, + } + dropped = {name: count for name, count in dropped.items() if count} + if dropped: + self._tracking.log_params(dropped) + def _log_params(self, params: Mapping[str, str], limit: int = _MAX_PARAMS - _FINISH_PARAMS) -> None: kept = {} @@ -269,6 +276,10 @@ def _log_metrics(self, metrics: Mapping[str, float], limit: int, self._tracking.log_metrics(kept, step=step) +def _now() -> datetime: + return datetime.now(timezone.utc) + + def _identifier(name: str) -> str: return re.sub(r'[^A-Za-z0-9_]+', '_', name).upper() diff --git a/packages/kumo-engine/kumo/engine/trainer/_run_spec.py b/packages/kumo-engine/kumo/engine/trainer/_run_spec.py index a4a3b95ce..78bcfe4e8 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_run_spec.py +++ b/packages/kumo-engine/kumo/engine/trainer/_run_spec.py @@ -7,13 +7,14 @@ from dataclasses import dataclass, replace from datetime import datetime, timezone from pathlib import Path -from typing import Any, Iterator, Mapping, Optional, Sequence, Tuple +from typing import Any, Dict, Iterator, Mapping, Optional, Sequence, Tuple from kumo_core.artifacts import PROVENANCE_FILE -from kumo_core.summaries import Summary +from kumo_core.summaries import Scalar, Summary from kumo.engine._run_names import fit_name, model_label -from kumo.engine._tracking import RunSpec +from kumo.engine._text import metric_key +from kumo.engine._tracking import RunSpec, logged_names from kumo.engine._version import __version__ @@ -51,31 +52,46 @@ def fit_spec(context: FitContext) -> RunSpec: ) -def fitted_spec(spec: RunSpec, result: Any, seconds: Optional[float] = None, - alias: Optional[str] = None, +def fitted_spec(spec: RunSpec, result: Any, summaries: Sequence[Summary] = ()) -> RunSpec: - r"""``spec`` once the fit returned ``result``.""" + r"""The fit run's spec once the fit returned.""" artifacts = dict(spec.artifacts) provenance = Path(result.bundle_dir) / PROVENANCE_FILE if provenance.is_file(): artifacts[PROVENANCE_FILE] = provenance.read_text(encoding='utf-8') - tested = {name: value for name, value in result.metrics.items() - if not name.startswith(('train_', 'val_'))} return replace( - spec, metrics={ - **tested, **({} if seconds is None else { - 'seconds': seconds - }) - }, parents={**spec.parents, 'bundle_id': result.id, - **_registered(result.registered_model, alias)}, + spec, metrics=holdout_metrics(result), + parents={**spec.parents, 'bundle_id': result.id, + **_registered(result.registered_model)}, artifacts=artifacts, summaries=summaries) -def _registered(model: Any, alias: Optional[str]) -> Mapping[str, Any]: +def holdout_metrics(result: Any) -> Dict[str, float]: + r"""The trainer's test metrics under their logged names.""" + return logged_names({ + name: value for name, value in result.metrics.items() + if not name.startswith(('train_', 'val_')) + }, refit_full=_fit_mode(result) == 'refit_full') + + +def registry_metrics(result: Any) -> Dict[str, float]: + r"""Holdout metrics and model scalars, as registries record them.""" + return {metric_key(name): value for name, value in { + **holdout_metrics(result), + **{summary.name: summary.value + for summary in getattr(result, '_summaries', ()) + if isinstance(summary, Scalar) and summary.role == 'model'} + }.items()} + + +def _fit_mode(result: Any) -> Optional[str]: + return getattr(getattr(result, 'manifest', None), 'fit_mode', None) + + +def _registered(model: Any) -> Mapping[str, Any]: if model is None: return {} - return {'model_name': model.name, 'model_version': model.version, - 'model_alias': alias} + return {'model_name': model.name, 'model_version': model.version} def _flatten(value: Any, prefix: str = '') -> Iterator[Tuple[str, Any]]: diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 625437941..5bfe1d9b0 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -95,7 +95,12 @@ read_provenance, record_derivation, ) -from kumo.engine.trainer._run_spec import FitContext, fit_spec, fitted_spec +from kumo.engine.trainer._run_spec import ( + FitContext, + fit_spec, + fitted_spec, + registry_metrics, +) from kumo.engine.trainer._scoring import ( ScoringDefaults, normalized_scoring_defaults, @@ -187,11 +192,12 @@ def export(self, config: Any) -> PublishedModel: tracker = as_tracking(config.tracking) if tracker is not None and tracker.register is not None: return cast(PublishedModel, _with_metrics( - tracker.register(self, config.registry), self.metrics)) + tracker.register(self, config.registry), + registry_metrics(self))) config = _connector_registry(config) _require_publisher(config) - return _publish(self.bundle_dir, self.bundle_id, self.metrics, - config)[0] + return _publish(self.bundle_dir, self.bundle_id, + registry_metrics(self), config)[0] def with_scoring_defaults(self, defaults: ScoringDefaults) -> 'TrainingResult': @@ -755,17 +761,14 @@ def fit( registered, error = _store_and_register( result.bundle_dir, bundle_manifest(result.bundle_dir, result.bundle_id), settings, - session=session, - snapshot=snapshot, metrics=result.metrics, warn=warn, + session=session, snapshot=snapshot, metrics=registry_metrics(result), warn=warn, stored=(deps.errors or [None])[0]) if registered is not None: result = replace(result, registered_model=registered) self._training_result = result - alias = None if settings is None or settings.registry is None else ( - settings.registry.alias) session.finish( lambda: fitted_spec( - cast(RunSpec, session.spec), result, seconds, alias, + cast(RunSpec, session.spec), result, _fit_summaries(snapshot, train_table, result, settings)), error) raise_on_error(settings, session.error or error, @@ -1053,7 +1056,7 @@ def _finish_fit( training_result = replace( training_result, published_model=_publish( training_result.bundle_dir, training_result.bundle_id, - training_result.metrics, export_config)[0]) + registry_metrics(training_result), export_config)[0]) training_result = replace(training_result, _trainer=self) self._training_result = training_result diff --git a/packages/kumo-engine/test/test_mlflow_model.py b/packages/kumo-engine/test/test_mlflow_model.py index a093bf5d3..897a09be1 100644 --- a/packages/kumo-engine/test/test_mlflow_model.py +++ b/packages/kumo-engine/test/test_mlflow_model.py @@ -21,6 +21,7 @@ import kumo.engine as kumo from kumo.engine._scratch import get_scratch_dir from kumo.engine._tracking import logged_names +from kumo.engine.trainer._run_spec import registry_metrics mlflow = pytest.importorskip('mlflow') @@ -60,19 +61,20 @@ def test_fit_logs_registers_and_serves_the_model(table_and_snapshot, model, = mlflow.search_logged_models(output_format='list') assert model.source_run_id == parent.info.run_id metrics = {metric.key: metric.value for metric in model.metrics} - assert metrics == { - **{name: value for name, value in logged_names(result.metrics).items() - if name.startswith('test/')}, 'seconds': metrics['seconds']} + holdout = {name: value for name, value in logged_names( + result.metrics).items() if not name.startswith(('train/', 'val/'))} + assert metrics == holdout version = mlflow.MlflowClient().get_model_version_by_alias('m', 'champion') assert version.source == model.model_uri assert parent.data.tags.items() >= { 'kumo.kind': 'fit', 'kumo.bundle_id': result.id, - 'kumo.model_name': 'm', 'kumo.model_version': str(version.version), - 'kumo.model_alias': 'champion'}.items() + 'kumo.model_name': 'm', + 'kumo.model_version': str(version.version)}.items() assert result.registered_model == PublishedModel( name='m', version=str(version.version), uri=f'models:/m/{version.version}', - metadata={'kumo_bundle_id': result.id}, metrics=result.metrics) + metadata={'kumo_bundle_id': result.id}, + metrics=registry_metrics(result)) assert version.tags == {'kumo_bundle_id': result.id} assert trainer.result is result @@ -140,9 +142,9 @@ def test_logged_model_scores_the_query_population_and_timed_rows( requirements = mlflow.pyfunc.get_model_dependencies(info.model_uri) assert f'kumo-engine[gcs,s3]=={kumo.__version__}' in open( requirements).read().splitlines() - assert {metric.key + assert {metric.key: metric.value for metric in mlflow.get_logged_model(info.model_id).metrics - } == set(trainer.result.metrics) + } == registry_metrics(trainer.result) assert [c['name'] for c in info.signature.outputs.to_dict() ] == ['ENTITY', 'TARGET_PRED', 'False_PROB', 'True_PROB'] pyfunc = mlflow.pyfunc.load_model(info.model_uri) diff --git a/packages/kumo-engine/test/test_mlflow_tracking.py b/packages/kumo-engine/test/test_mlflow_tracking.py index 2bb4918d5..6db9314e1 100644 --- a/packages/kumo-engine/test/test_mlflow_tracking.py +++ b/packages/kumo-engine/test/test_mlflow_tracking.py @@ -153,10 +153,9 @@ def test_fit_logs_a_run_per_fit_and_per_trial(table_and_snapshot, tmp_path, if not name.startswith('system/') } tested = {name: value for name, value in logged_names( - result.metrics).items() if name.startswith('test/')} - assert {name for name in metrics if not name.startswith(( - 'train/', 'val/', 'data/', 'graph/'))} == { - *tested, 'seconds', 'search/best_val_mae'} + result.metrics).items() if not name.startswith(('train/', 'val/'))} + assert {name for name in metrics + if name.startswith(('test/', 'before_refit/'))} == set(tested) assert metrics['graph/rows/ents'] == 400 assert {name: metrics[name] for name in tested} == tested assert 'off for the rest of this run' not in caplog.text @@ -171,8 +170,7 @@ def test_fit_logs_a_run_per_fit_and_per_trial(table_and_snapshot, tmp_path, for child in children: assert child.info.status == 'FINISHED' assert child.data.tags['kumo.status'] == 'succeeded' - assert {'val/best_mae', 'seconds', - 'val/mae'} <= set(child.data.metrics) + assert {'val/best_mae', 'val/mae'} <= set(child.data.metrics) history = client.get_metric_history(child.info.run_id, 'val/mae') assert sorted(metric.step for metric in history) == [1, 2] selected = min(children, key=lambda run: run.data.metrics['val/best_mae']) diff --git a/packages/kumo-engine/test/test_snowflake_tracking.py b/packages/kumo-engine/test/test_snowflake_tracking.py index be6b766e2..37eae359d 100644 --- a/packages/kumo-engine/test/test_snowflake_tracking.py +++ b/packages/kumo-engine/test/test_snowflake_tracking.py @@ -71,7 +71,7 @@ def _replay(tracking, result=RESULT, error=None, during=None): during() session.finish( None if result is None else - lambda: fitted_spec(session.spec, result, EVENTS[-1].metrics['seconds']), + lambda: fitted_spec(session.spec, result), error) return session @@ -188,10 +188,10 @@ def test_metric_and_param_limits(fake, monkeypatch): run, = fake._instance.runs.values() assert sorted({name for name, _, _ in run['metrics']}) == [ - 'seconds', 'test/mae', 'trial-1-attempt-1/train/loss', - 'trial-1-attempt-1/val/best_mae', 'trial-1-attempt-1/val/mae' + 'TRIAL_1/val/best_mae', 'TRIAL_2_ATTEMPT_2/val/best_mae', + 'duration_seconds', 'search/best_val_mae', 'test/mae' ] - assert run['params']['kumo.dropped_metrics'] == '8' + assert run['params']['kumo.dropped_metrics'] == '2' assert run['params']['kumo.dropped_params'] == '1' assert 'note' not in run['params'] assert run['params']['level'] == '2' diff --git a/packages/kumo-engine/test/test_tracking_contract.py b/packages/kumo-engine/test/test_tracking_contract.py index 68daed1e8..0fcfba8e5 100644 --- a/packages/kumo-engine/test/test_tracking_contract.py +++ b/packages/kumo-engine/test/test_tracking_contract.py @@ -86,7 +86,7 @@ def _epochs(trial, *values): BUNDLE = 'b' * 64 RESULT = SimpleNamespace(metrics={'val_mae': 0.3, 'test_mae': 0.32}, id=BUNDLE, bundle_dir=Path('no-bundle'), - registered_model=None) + registered_model=None, registered_alias=None) class _Recorder: @@ -135,15 +135,14 @@ def replay(tracker=None, events=EVENTS, result=RESULT, error=None, session.on_event(event) session.finish( None if result is None else lambda: fitted_spec( - session.spec, result, events[-1].metrics.get('seconds'), - summaries=summaries), error) + session.spec, result, summaries=summaries), error) records = { record.name: record for hook, record in recorder.calls if hook == 'on_trial' } spec = replace(fit_spec(context), tags=USER_TAGS) if result is not None: - spec = fitted_spec(spec, result, events[-1].metrics.get('seconds')) + spec = fitted_spec(spec, result) spec = replace(spec, steps=series if series is not None else { **SEARCH, **({} if error else SELECTED)}) if series is None and not error: @@ -174,11 +173,13 @@ class Adapter: step_metrics: bool = False artifacts: bool = False run_name: Callable[[str], Optional[str]] = str + expected: Callable[[Dict[str, RunRecord]], Dict[str, RunRecord]] = dict def gaps(expected: Dict[str, RunRecord], observed: Dict[str, Observed], adapter: Adapter) -> List[str]: kept_params = adapter.kept_params + expected = adapter.expected(expected) found = [] if set(observed) != set(expected): return [f'runs {sorted(observed)} != {sorted(expected)}'] @@ -295,7 +296,7 @@ def observe(): def bucket(key): prefix, _, name = key.partition('/') - if not prefix.startswith('trial-'): + if not prefix.startswith('TRIAL_'): prefix, name = '', key return names.setdefault(prefix or None, { 'params': {}, 'tags': {}, 'metrics': {}, 'steps': {} @@ -309,6 +310,8 @@ def bucket(key): else: runs['params'][name] = value for key, value, step in run['metrics']: + if key == 'duration_seconds': + continue runs, name = bucket(key) runs['metrics'][name] = value runs['steps'][name] = (*runs['steps'].get(name, ()), @@ -320,11 +323,27 @@ def bucket(key): **runs) return observed + def expected(records): + r"""What Snowflake keeps of each trial.""" + kept = {} + for key, record in records.items(): + if record.tags['kumo.kind'] == 'trial': + trial, attempt = key.split('-')[1::2] + key = f'TRIAL_{trial}' + ( + f'_ATTEMPT_{attempt}' if attempt != '1' else '') + record = replace(record, steps={}, metrics={ + name: value for name, value in record.metrics.items() + if name == 'val/best_mae'}, tags={ + name: value for name, value in record.tags.items() + if name in ('kumo.status', 'kumo.error')}) + kept[key] = record + return kept + return Adapter( make=lambda: SnowflakeTracking(_snowflake_experiment_fake.FakeSession(), experiment_name='ents', tags=USER_TAGS), observe=observe, step_metrics=True, artifacts=True, - run_name=lambda name: None) + run_name=lambda name: None, expected=expected) @pytest.fixture(params=['mlflow', 'snowflake']) @@ -356,16 +375,17 @@ def test_records_of_the_recorded_fit(): spoofed = run_record(replace(fit_spec(CONTEXT), tags={ 'kumo.bundle_id': 'x', 'training_job.tune_metric': 'y', 'a': 1})) assert spoofed.tags.keys() - fit.tags.keys() == {'a'} - assert fit.metrics == {'test/mae': 0.32, 'seconds': 9.0} + assert fit.metrics == {'test/mae': 0.32} assert fit.steps == {'search/best_val_mae': ((1, 0.35), (2, 0.3)), 'train/loss': ((1, 0.6), (2, 0.4)), 'val/mae': ((1, 0.5), (2, 0.3))} - registered = SimpleNamespace(**{**vars(RESULT), 'registered_model': - SimpleNamespace(name='ents', version='3')}) - tags = run_record(fitted_spec(fit_spec(CONTEXT), registered, - alias='candidate')).tags - assert tags.items() >= {'kumo.model_name': 'ents', 'kumo.model_version': - '3', 'kumo.model_alias': 'candidate'}.items() + registered = SimpleNamespace(**{ + **vars(RESULT), 'registered_alias': 'candidate', + 'registered_model': SimpleNamespace(name='ents', version='3')}) + tags = run_record(fitted_spec(fit_spec(CONTEXT), registered)).tags + assert tags.items() >= {'kumo.model_name': 'ents', + 'kumo.model_version': '3'}.items() + assert 'kumo.model_alias' not in tags assert fit.artifacts['query.pql'] == QUERY assert json.loads(fit.artifacts['model_plan.json']) == CONTEXT.model_plan @@ -379,7 +399,7 @@ def test_records_of_the_recorded_fit(): }) retry = records['trial-2-attempt-2'] assert retry.params['attempt'] == '2' - assert retry.metrics == {'val/best_mae': 0.3, 'seconds': 4.0} + assert retry.metrics == {'val/best_mae': 0.3} assert retry.steps == {'train/loss': ((1, 0.6), (2, 0.4)), 'val/mae': ((1, 0.5), (2, 0.3))} diff --git a/packages/kumo-engine/test/test_trainer_registry.py b/packages/kumo-engine/test/test_trainer_registry.py index a7c77f0ff..13c236981 100644 --- a/packages/kumo-engine/test/test_trainer_registry.py +++ b/packages/kumo-engine/test/test_trainer_registry.py @@ -231,7 +231,7 @@ def test_finish_fit_forwards_publication_details_and_returns_assigned_version( tmp_path, path, location): bundle = _fake_published_bundle(tmp_path) trainer = Trainer(ModelPlan()) - original = TrainingResult({'val_auroc': .9}, {}, bundle) + original = TrainingResult({'val_auroc': .9, 'test_auroc': .8}, {}, bundle) published = PublishedModel(name='CHURN_GNN', version='assigned-v7', uri='provider://models/CHURN_GNN/assigned-v7') calls = [] @@ -262,7 +262,7 @@ def publish_model(self, connector, bundle_dir, **kwargs): assert calls[0][1] == str(bundle) assert calls[0][2]['path'] == path assert calls[0][2]['location'] == location - assert calls[0][2]['metrics'] == {'val_auroc': .9} + assert calls[0][2]['metrics'] == {'test/auroc': .8} assert calls[0][2]['metadata']['kumo_bundle_id'] == result.bundle_id assert result.published_model is published assert result.metrics == original.metrics diff --git a/packages/kumo-snowflake/src/kumo_snowflake/registry.py b/packages/kumo-snowflake/src/kumo_snowflake/registry.py index cca96ef2d..42a9d55de 100644 --- a/packages/kumo-snowflake/src/kumo_snowflake/registry.py +++ b/packages/kumo-snowflake/src/kumo_snowflake/registry.py @@ -34,6 +34,7 @@ import json import logging +import math import re import shutil from contextlib import contextmanager @@ -92,6 +93,8 @@ #: a reader can look at. _EXPORT_DIR = '_registry-export' +_METRICS_LINE = 'Holdout: ' + def _require_registry() -> type["Registry"]: r"""Returns the ``Registry`` class, or explains how to install it.""" @@ -245,6 +248,20 @@ def _version_record(metadata: Mapping[str, str]) -> Dict[str, str]: return record +def _finite(metrics: Mapping[str, float]) -> Dict[str, float]: + return {name: float(value) for name, value in metrics.items() + if math.isfinite(value)} + + +def _comment(user: Optional[str], + metrics: Mapping[str, float]) -> Optional[str]: + r"""The user's comment, then a line of the model's holdout metrics.""" + line = ' · '.join(f'{name} {value:.3g}' + for name, value in _finite(metrics).items()) + return '\n'.join(text for text in (user, line and _METRICS_LINE + line) + if text) or None + + def _recorded_lineage(version: "ModelVersion") -> Dict[str, str]: try: return registry_lineage(dict(version.show_metrics() or {})) @@ -272,6 +289,16 @@ def _restore_alias(model: Any, version: str, alias: str) -> None: alias, version) +def _next_version(registry: Any, model_name: str) -> str: + try: + names = [str(version.version_name) + for version in registry.get_model(model_name).versions()] + except ValueError: + names = [] + numbers = [int(name[1:]) for name in names if re.fullmatch(r'V\d+', name)] + return f'V{max(numbers, default=0) + 1}' + + class SnowflakeModelRegistryMixin: r"""The model half of :class:`SnowflakeSinkBackend`. @@ -338,12 +365,10 @@ def publish_model( path: Refused. A registry addresses a model by name. location: The :class:`~kumoapi.plugin.ModelRegistryLocation` naming - the model, and optionally the version. A version of ``None`` - lets the registry assign one -- the usual case, and the reason - this reports back what it created. + the model, and optionally the version; by default the next + ``V``. metadata: Returned with the result; its bundle id is recorded. - metrics: Returned with the publication result. Training metrics - remain in the bundle's metrics.json. + metrics: The model's holdout metrics, recorded on the version. project_id: Ignored; a local runtime has no projects. Returns: @@ -361,17 +386,27 @@ def publish_model( f"{bundle} does not look like a published model bundle: it " f"has no {_SPEC_FILE}, which an export writes into every one.") - record = _version_record(metadata) + record = {**_finite(metrics), **_version_record(metadata)} with self._registry(connector) as registry: - version = registry.log_model( - _bundle_model(str(bundle)), - model_name=location.model_name, - version_name=location.version, - comment=location.comment, - metrics=record, - signatures=_carrier_signature(), - **_task(bundle), - ) + name = location.version or _next_version(registry, + location.model_name) + for attempt in range(3): + try: + version = registry.log_model( + _bundle_model(str(bundle)), + model_name=location.model_name, + version_name=name, + comment=_comment(location.comment, metrics), + metrics=record, + signatures=_carrier_signature(), + **_task(bundle), + ) + break + except Exception: + taken, name = name, location.version or _next_version( + registry, location.model_name) + if name == taken or attempt == 2: + raise return _published(connector, version, metadata=metadata, metrics=metrics) diff --git a/packages/kumo-snowflake/tests/test_registry.py b/packages/kumo-snowflake/tests/test_registry.py index 7c6e80dbc..4d612f8ad 100644 --- a/packages/kumo-snowflake/tests/test_registry.py +++ b/packages/kumo-snowflake/tests/test_registry.py @@ -86,7 +86,7 @@ def test_invalid_metrics_preserve_previous_destination(fetch, payload): 'kumo_bundle_id': 'b1', 'kumo_engine': '0' }]) -def test_publish_records_only_pointers_in_catalog( +def test_publish_records_pointers_and_holdout_metrics_in_catalog( tmp_path, monkeypatch, metadata): bundle = tmp_path / 'bundle' bundle.mkdir() @@ -99,23 +99,85 @@ def log_model(*args, **kwargs): calls.append(kwargs) return version + def get_model(_): + raise ValueError('Unable to find model CHURN') + @contextmanager def registry(_): - yield SimpleNamespace(log_model=log_model) + yield SimpleNamespace(log_model=log_model, get_model=get_model) sink = SnowflakeModelRegistryMixin() monkeypatch.setattr(sink, '_registry', registry) monkeypatch.setattr(registry_module, '_bundle_model', lambda _: object()) monkeypatch.setattr(registry_module, '_carrier_signature', lambda: {}) - published = sink.publish_model(SimpleNamespace(), str(bundle), - location=ModelRegistryLocation('CHURN'), - metadata=metadata, - metrics={'auc': 0.92, 'loss': float('nan')}) + metrics = {'auroc': 0.9213, 'holdout_loss': 0.5, 'brier': float('nan')} + published = sink.publish_model( + SimpleNamespace(), str(bundle), + location=ModelRegistryLocation('CHURN', comment='Weekly retrain.'), + metadata=metadata, metrics=metrics) assert calls[0]['metrics'] == { + 'auroc': 0.9213, 'holdout_loss': 0.5, **({'kumo_bundle_id': 'b1'} if metadata else {}), 'kumo_engine': '1' } + assert calls[0]['version_name'] == 'V1' + assert calls[0]['comment'] == ('Weekly retrain.\n' + 'Holdout: auroc 0.921 · holdout_loss 0.5') assert published.metadata == metadata - assert published.metrics['auc'] == 0.92 + assert published.metrics == metrics + + +def test_publish_names_the_next_version_and_retries_a_taken_one( + tmp_path, monkeypatch): + bundle = tmp_path / 'bundle' + bundle.mkdir() + (bundle / 'training-spec.json').write_text('{}') + names = ['V1', 'SLIMY_DONKEY_1', 'V3'] + tried = [] + + def log_model(*args, version_name, **kwargs): + tried.append(version_name) + if version_name == 'V4': + names.append('V4') + raise ValueError('Model CHURN version V4 already existed.') + return SimpleNamespace(model_name='CHURN', version_name=version_name, + fully_qualified_model_name='DB.SCHEMA.CHURN') + + model = SimpleNamespace(versions=lambda: [ + SimpleNamespace(version_name=name) for name in names]) + + @contextmanager + def registry(_): + yield SimpleNamespace(log_model=log_model, get_model=lambda _: model) + + sink = SnowflakeModelRegistryMixin() + monkeypatch.setattr(sink, '_registry', registry) + monkeypatch.setattr(registry_module, '_bundle_model', lambda _: object()) + monkeypatch.setattr(registry_module, '_carrier_signature', lambda: {}) + published = sink.publish_model(SimpleNamespace(), str(bundle), + location=ModelRegistryLocation('CHURN')) + assert tried == ['V4', 'V5'] + assert published.version == 'V5' + with pytest.raises(ValueError, match='already existed'): + sink.publish_model(SimpleNamespace(), str(bundle), + location=ModelRegistryLocation('CHURN', + version='V4')) + + def fail_after_creating(*args, version_name, **kwargs): + tried.append(version_name) + names.append(version_name) + raise RuntimeError('upload failed') + + @contextmanager + def failing(_): + yield SimpleNamespace(log_model=fail_after_creating, + get_model=lambda _: model) + + tried.clear() + monkeypatch.setattr(sink, '_registry', failing) + with pytest.raises(RuntimeError, match='upload failed'): + sink.publish_model(SimpleNamespace(), str(bundle), + location=ModelRegistryLocation('CHURN')) + assert len(tried) == 3 def test_listing_filters_engine_marker_and_returns_lineage(monkeypatch): @@ -290,7 +352,7 @@ def registry(_): monkeypatch.setattr(registry_module, '_bundle_model', lambda _: object()) monkeypatch.setattr(registry_module, '_carrier_signature', lambda: {}) sink.publish_model(SimpleNamespace(), str(bundle), - location=ModelRegistryLocation('CHURN')) + location=ModelRegistryLocation('CHURN', version='V1')) assert calls[0].get('task') == expected From 356fecfa54b93f8eb3207ddc8974cb6074d706a7 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 03:25:15 +0000 Subject: [PATCH 209/278] Register a fit's logged model, or log one in no run, and name it after the model register() reuses the model the fit run logged, else logs it with no run in the configured or model-named experiment; load_model resolves an alias to its version first. Callers look up the held version once. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_tracking.py | 6 + .../kumo/engine/integrations/mlflow.py | 156 +++++++++++++----- .../kumo-engine/test/test_mlflow_model.py | 52 +++++- .../kumo-engine/test/test_mlflow_target.py | 16 +- 4 files changed, 178 insertions(+), 52 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_tracking.py b/packages/kumo-engine/kumo/engine/_tracking.py index c44dbd93e..1e3629bec 100644 --- a/packages/kumo-engine/kumo/engine/_tracking.py +++ b/packages/kumo-engine/kumo/engine/_tracking.py @@ -194,6 +194,12 @@ def trial_spec(events: Sequence[Event]) -> RunSpec: ) +def fit_run_id(manifest: Any) -> Optional[str]: + r"""The tracker's run of the fit a `ModelManifest` records.""" + tracking = getattr(manifest, 'tracking', None) + return None if tracking is None else tracking.run_id + + class Tracking(Protocol): r"""Records one `Trainer.fit` in an experiment tracker.""" def start(self, record: RunRecord) -> None: diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index 0e6bd2636..ff73c0dd9 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -16,6 +16,7 @@ from contextlib import ExitStack, contextmanager from datetime import datetime from pathlib import Path +from types import SimpleNamespace from typing import ( TYPE_CHECKING, Any, @@ -40,7 +41,7 @@ from kumo.engine._events import Event from kumo.engine._mlops import full_error -from kumo.engine._run_names import record_name, trial_name +from kumo.engine._run_names import model_label, record_name, trial_name from kumo.engine._summary_tables import Routed, route, table_json from kumo.engine._text import metric_key from kumo.engine._tracking import ( @@ -49,6 +50,7 @@ RunRecord, TrackingBase, by_step, + fit_run_id, ) from kumo.engine._version import __version__ from kumo.engine.connector._extras import import_provider @@ -192,12 +194,14 @@ def prepare(self) -> None: self._outer = (active.info.run_id if active is not None and same else None) - def _experiment_for(self, client: Any) -> str: - if self.experiment_name is None: + def _experiment_for(self, client: Any, + name: Optional[str] = None) -> str: + name = name or self.experiment_name + if name is None: return self._default_experiment or '0' - experiment = client.get_experiment_by_name(self.experiment_name) - return (experiment.experiment_id if experiment else - client.create_experiment(self.experiment_name)) + experiment = client.get_experiment_by_name(name) + return (experiment.experiment_id + if experiment else client.create_experiment(name)) def _require_experiment(self) -> None: if self._named and self.experiment_name is None: @@ -210,43 +214,76 @@ def load_model(self, location: ModelRegistryLocation, if (location.alias is None) == (location.version is None): raise ValueError('Pass a ModelRegistryLocation with either an ' 'alias or a version.') - reference = (f'@{location.alias}' - if location.alias else f'/{location.version}') - return load_model(f'models:/{location.model_name}{reference}', - target=self.target, **kwargs) + name, alias, version = (location.model_name, location.alias, + location.version) + if alias: + with self.target.scope(): + version = self.target.client().get_model_version_by_alias( + name, alias).version + uri = f'models:/{name}/{version}' + trainer = load_model(uri, target=self.target, **kwargs) + return trainer def register(self, result: Any, location: ModelRegistryLocation) -> PublishedModel: r"""Registers a `TrainingResult` at ``location``.""" from kumo.engine.graph import GraphSnapshot from kumo.engine.trainer._provenance import bundle_manifest + from kumo.engine.trainer._run_spec import registry_metrics from kumo.engine.trainer._spec import read_training_spec from kumo.engine.trainer.trainer import registry_pointers self._require_experiment() - client = self.target.client() + self._client = client = self.target.client() pointers = registry_pointers(result.id) + metrics = registry_metrics(result) with _limited_http(): with self.target.scope(): held = _registered_version(client, location.model_name, pointers) - info = None - if held is None: + fit, model = self._fit_of(fit_run_id(result.manifest)) + info = None if model is None else SimpleNamespace( + model_uri=f'models:/{model.model_id}', run_id=fit.run_id, + model_id=model.model_id) + if held is None and info is None: spec = read_training_spec(result.bundle_dir) snapshot = GraphSnapshot(spec.snapshot_id, Path(), spec.graph_config or {}) - with self.target.scope(): - experiment_id = self._experiment_for(client) - run_id = client.create_run( - experiment_id, - tags={'kumo.bundle_id': str(result.id)}).info.run_id - info = _log_model(self.target, client, result.bundle_dir, - bundle_manifest(result.bundle_dir, - result.id), - snapshot, run_id, experiment_id) - with self.target.scope(): - client.set_terminated(run_id, 'FINISHED') - return _register(self.target, client, info, location, pointers) + info = _log_model( + self.target, client, result.bundle_dir, + bundle_manifest(result.bundle_dir, result.id), snapshot, + name=_logged_model_name(location.model_name, + result.manifest), + run_id=None if fit is None else fit.run_id, + experiment_id=(self._registry_experiment( + location.model_name) if fit is None else + fit.experiment_id), serving_snapshot=None) + if fit is not None: + self._log_metrics(fit.run_id, metrics, + model_id=info.model_id) + published = _register(self.target, client, info, location, + pointers, held) + return published + + def _logged_model(self, fit: Any) -> Any: + try: + with self.target.scope(): + models = self._client.search_logged_models( + [fit.experiment_id], + filter_string=f"source_run_id = '{fit.run_id}'", + max_results=1) + except Exception: + logger.debug('No logged model of %s', fit.run_id, exc_info=True) + return None + return models[0] if models else None + + def _registry_experiment(self, model_name: str) -> str: + r"""The experiment for a model logged in no run, never Default.""" + name = self.experiment_name or model_label(model_name, None, None) + if self.experiment_name is None and self.platform == 'databricks': + name = f'/Shared/kumo/{name}' + with self.target.scope(): + return self._experiment_for(self._client, name) def start(self, record: RunRecord) -> None: self._require_experiment() @@ -336,7 +373,11 @@ def log_model( os.environ.update(upload) try: info = _log_model(self.target, self._client, bundle_dir, manifest, - snapshot, run_id, self._experiment_id) + snapshot, name=_logged_model_name( + registry and registry.model_name, manifest), + run_id=run_id, + experiment_id=self._experiment_id, + serving_snapshot=None) finally: if self._parent == run_id: os.environ.update({name: _HTTP_LIMITS[name] for name in upload}) @@ -347,9 +388,13 @@ def log_model( self._model_id = info.model_id if registry is None: return None + pointers = registry_pointers(_id(manifest)) try: + with self.target.scope(): + held = _registered_version(self._client, + registry.model_name, pointers) return _register(self.target, self._client, info, registry, - registry_pointers(_id(manifest))) + pointers, held) except Exception as error: raise RegistrationError(full_error(error)) from error @@ -384,6 +429,18 @@ def finish(self, record: Optional[RunRecord], if fit_run and hint: warnings.warn(hint, UserWarning, stacklevel=2) + def _fit_of(self, fit_id: Optional[str]) -> Tuple[Any, Any]: + r"""The fit run's info and its logged model, each None if missing.""" + if not fit_id: + return None, None + try: + with self.target.scope(): + fit = self._client.get_run(fit_id).info + except Exception: + logger.debug('No fit run %s here', fit_id, exc_info=True) + return None, None + return fit, self._logged_model(fit) + def log_metrics(self, metrics: Mapping[str, float], step: int) -> None: self._log_metrics(self._parent, metrics, step) @@ -626,7 +683,7 @@ def log_model( trainer: 'Trainer', *, snapshot: 'GraphSnapshot', - name: str = 'model', + name: Optional[str] = None, run_id: Optional[str] = None, serving_snapshot: Any = None, target: Any = None, @@ -652,7 +709,7 @@ def log_model( target = tracking.target if run_id is None: tracking._require_experiment() - tracking._client = client = target.client() + client = tracking._client = target.client() active = mlflow.active_run() if run_id is None and active is not None and ( mlflow.get_tracking_uri() == target.tracking_uri): @@ -671,10 +728,13 @@ def log_model( except Exception as error: logger.warning('Could not log the fit\'s source tables to MLflow: %s', format_error(error), exc_info=True) + registered = getattr(result.registered_model, 'name', None) info = _log_model(target, client, result.bundle_dir, bundle_manifest(result.bundle_dir, result.id), - snapshot, run_id, experiment_id, name, - snapshot_uri(serving_snapshot)) + snapshot, name=name or _logged_model_name( + registered, result.manifest), run_id=run_id, + experiment_id=experiment_id, + serving_snapshot=snapshot_uri(serving_snapshot)) tracking._log_metrics(run_id, registry_metrics(result), model_id=info.model_id) if created: @@ -734,11 +794,20 @@ def _download(mlflow: Any, target: Any, model_uri: str, tracking_uri=target.tracking_uri, registry_uri=target.registry_uri) +def _logged_model_name(registered: Optional[str], manifest: Any) -> str: + r"""model_label without the characters MLflow rejects.""" + table = getattr(getattr(manifest, 'parents', None), 'training_table', + None) + return re.sub(r'[/:.%"\']', '-', model_label( + registered, getattr(table, 'query', None), + getattr(table, 'task_type', None))) + + def _log_model(target: Any, client: Any, bundle_dir: Path, manifest: Optional['ModelManifest'], - snapshot: Optional['GraphSnapshot'], run_id: str, - experiment_id: str, name: str = 'model', - serving_snapshot: Optional[str] = None) -> Any: + snapshot: Optional['GraphSnapshot'], *, name: str, + run_id: Optional[str], experiment_id: Optional[str], + serving_snapshot: Optional[str]) -> Any: from kumo_core.artifacts import PROVENANCE_FILE from kumo.engine.integrations import _mlflow_model, _mlflow_model_code @@ -757,10 +826,11 @@ def _log_model(target: Any, client: Any, bundle_dir: Path, pointers = registry_pointers(_id(manifest)) provenance = bundle_dir / PROVENANCE_FILE with target.scope(): + run = None if run_id is None else client.get_run(run_id) model = client.create_logged_model( - experiment_id, - name=name, source_run_id=run_id, - params=dict(client.get_run(run_id).data.params)) + experiment_id or run.info.experiment_id, name=name, + source_run_id=run_id, + params=None if run is None else dict(run.data.params)) mlflow_model = mlflow.models.Model( artifact_path=model.artifact_location, model_uuid=model.model_id, run_id=run_id, metadata=pointers, model_id=model.model_id) @@ -798,9 +868,10 @@ def _log_model(target: Any, client: Any, bundle_dir: Path, repository.log_artifacts(local) logged = client.finalize_logged_model(model.model_id, 'READY') - client.log_outputs(run_id, [ - mlflow.entities.LoggedModelOutput(model.model_id, 0) - ]) + if run_id is not None: + client.log_outputs(run_id, [ + mlflow.entities.LoggedModelOutput(model.model_id, 0) + ]) except BaseException: with target.scope(): _quietly(client.finalize_logged_model, model.model_id, 'FAILED') @@ -828,10 +899,11 @@ def _requirement(config: Any) -> str: def _register(target: Any, client: Any, info: Any, location: ModelRegistryLocation, - metadata: Mapping[str, str]) -> PublishedModel: + metadata: Mapping[str, str], + version: Optional[str]) -> PublishedModel: + r"""Registers ``info`` as a new version unless ``version`` holds it.""" name = location.model_name with target.scope(): - version = _registered_version(client, name, metadata) if version is None: try: client.create_registered_model(name) diff --git a/packages/kumo-engine/test/test_mlflow_model.py b/packages/kumo-engine/test/test_mlflow_model.py index 897a09be1..e3b51d17a 100644 --- a/packages/kumo-engine/test/test_mlflow_model.py +++ b/packages/kumo-engine/test/test_mlflow_model.py @@ -59,12 +59,12 @@ def test_fit_logs_registers_and_serves_the_model(table_and_snapshot, if 'mlflow.parentRunId' not in run.data.tags ] model, = mlflow.search_logged_models(output_format='list') - assert model.source_run_id == parent.info.run_id + assert (model.name, model.source_run_id) == ('m', parent.info.run_id) + version = mlflow.MlflowClient().get_model_version_by_alias('m', 'champion') metrics = {metric.key: metric.value for metric in model.metrics} holdout = {name: value for name, value in logged_names( result.metrics).items() if not name.startswith(('train/', 'val/'))} assert metrics == holdout - version = mlflow.MlflowClient().get_model_version_by_alias('m', 'champion') assert version.source == model.model_uri assert parent.data.tags.items() >= { 'kumo.kind': 'fit', 'kumo.bundle_id': result.id, @@ -108,6 +108,7 @@ def test_fit_logs_registers_and_serves_the_model(table_and_snapshot, registry=kumo.ModelRegistryLocation('m2')) first = champion.export(other) assert first.version == '1' and champion.export(other) == first + assert len(mlflow.search_logged_models(output_format='list')) == 1 assert load_model('models:/m2/1').result.id == result.id original = mlflow.models.Model.load @@ -145,6 +146,8 @@ def test_logged_model_scores_the_query_population_and_timed_rows( assert {metric.key: metric.value for metric in mlflow.get_logged_model(info.model_id).metrics } == registry_metrics(trainer.result) + assert mlflow.get_logged_model(info.model_id).name == ( + 'ents-binary_classification') assert [c['name'] for c in info.signature.outputs.to_dict() ] == ['ENTITY', 'TARGET_PRED', 'False_PROB', 'True_PROB'] pyfunc = mlflow.pyfunc.load_model(info.model_uri) @@ -367,7 +370,7 @@ def test_model_upload_gets_a_longer_http_timeout_and_a_late_one_is_dropped( location = kumo.ModelRegistryLocation('m') seen = [] - def log_model(*args): + def log_model(target, client, bundle_dir, manifest, snapshot, **kwargs): seen.append(os.environ.get('MLFLOW_HTTP_REQUEST_TIMEOUT')) if len(seen) == 2: tracker.finish(None, None) @@ -523,3 +526,46 @@ def no_folder(self, name, *args, **kwargs): 'earlier failure (MlflowException: RESOURCE_DOES_NOT_EXIST: Parent ' 'of /NoSuchFolder/churn does not exist.)' ] + + +def test_a_logged_model_is_named_after_its_model(): + from kumo.engine.integrations.mlflow import _logged_model_name + + manifest = SimpleNamespace(parents=SimpleNamespace( + training_table=SimpleNamespace( + query='PREDICT users.CHURN FOR EACH users.ID', + task_type='binary_classification'))) + assert _logged_model_name('main.ml.churn', manifest) == 'churn' + assert _logged_model_name('a/b:c%d"e\'f', None) == 'a-b-c-d-e-f' + assert _logged_model_name(None, manifest) == 'users-binary_classification' + assert _logged_model_name(None, None) == 'model' + + +def test_a_later_registration_logs_the_model_in_its_fit_run_else_in_none( + table_and_snapshot, tmp_path, tracking_uri): + from kumo.engine.integrations.mlflow import MlflowTracking + + table, snapshot = table_and_snapshot + tracking = MlflowTracking('fits', log_system_metrics=False) + result = kumo.Trainer(_plan()).fit(table, snapshot=snapshot, mlops=( + kumo.MLOps(str(tmp_path / 'root'), tracking=tracking))) + fits = [mlflow.get_experiment_by_name('fits').experiment_id] + assert mlflow.MlflowClient().search_logged_models(fits) == [] + published = tracking.register(result, kumo.ModelRegistryLocation('m')) + + fit = mlflow.get_run(result.tracking.run_id) + model, = mlflow.MlflowClient().search_logged_models(fits) + assert (published.version, model.source_run_id, model.name) == ( + '1', fit.info.run_id, 'm') + assert {m.key: m.value for m in model.metrics} == registry_metrics(result) + + other = f'sqlite:///{tmp_path / "other.db"}' + mlflow.set_tracking_uri(other) + untracked = MlflowTracking(log_system_metrics=False) + assert untracked.register(result, kumo.ModelRegistryLocation('m')) == ( + published) + client = mlflow.MlflowClient(other) + experiment = client.get_experiment_by_name('m').experiment_id + model, = client.search_logged_models([experiment]) + assert (model.source_run_id, client.search_runs([experiment])) == ( + None, []) diff --git a/packages/kumo-engine/test/test_mlflow_target.py b/packages/kumo-engine/test/test_mlflow_target.py index 01f05dad2..731e91723 100644 --- a/packages/kumo-engine/test/test_mlflow_target.py +++ b/packages/kumo-engine/test/test_mlflow_target.py @@ -293,7 +293,7 @@ def test_pin_databricks_auth_variables(): def test_registering_a_registered_bundle_reuses_its_version( tracking_uri, caplog): - from kumo.engine.integrations.mlflow import _register + from kumo.engine.integrations.mlflow import _register, _registered_version target = resolve(None) client = target.client() @@ -302,13 +302,15 @@ def test_registering_a_registered_bundle_reuses_its_version( run_id=run.info.run_id, model_id=None) location = kumo.ModelRegistryLocation('churn', alias='champion') - first = _register(target, client, info, location, - {'kumo_bundle_id': 'a'}) + def register(bundle): + pointers = {'kumo_bundle_id': bundle} + return _register(target, client, info, location, pointers, + _registered_version(client, 'churn', pointers)) + + first = register('a') with caplog.at_level('INFO'): - again = _register(target, client, info, location, - {'kumo_bundle_id': 'a'}) - other = _register(target, client, info, location, - {'kumo_bundle_id': 'b'}) + again = register('a') + other = register('b') assert (first.version, again.version, other.version) == ('1', '1', '2') assert 'churn version 1 already holds bundle a; reused it' in caplog.text From f76c34968fb0cdaede5c24a4f806b037a93a558e Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 03:26:43 +0000 Subject: [PATCH 210/278] Sample MLflow system metrics every second and log them in batches Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/integrations/mlflow.py | 53 ++++++++++++++----- .../kumo-engine/test/test_mlflow_tracking.py | 27 ++++++---- 2 files changed, 57 insertions(+), 23 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index ff73c0dd9..fd82f8eb7 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -327,8 +327,9 @@ def start(self, record: RunRecord) -> None: self.logged_run_name = run.info.run_name self._experiment_id = run.info.experiment_id if self.log_system_metrics and record.tags.get('kumo.kind') == 'fit': - self._guard.enter_context( - _SystemMetrics(self.log_metrics)) + self._guard.enter_context(_SystemMetrics( + lambda readings: self._log_readings(client, run.info.run_id, + readings))) self._log_params(self._parent, record.params) self._log_artifacts(record.artifacts) @@ -444,6 +445,19 @@ def _fit_of(self, fit_id: Optional[str]) -> Tuple[Any, Any]: def log_metrics(self, metrics: Mapping[str, float], step: int) -> None: self._log_metrics(self._parent, metrics, step) + def _log_readings(self, client: Any, run_id: str, + readings: Sequence['_Reading']) -> None: + entities = self._mlflow.entities + metrics = [ + entities.Metric(metric_key(name), float(value), timestamp, step) + for values, step, timestamp in readings + for name, value in values.items() + ] + with self.target.scope(): + for start in range(0, len(metrics), _BATCH_METRICS): + client.log_batch( + run_id, metrics=metrics[start:start + _BATCH_METRICS]) + def _log_summaries(self, summaries: Sequence[Summary]) -> None: self._log_routed(route(summaries), self._parent) @@ -579,11 +593,17 @@ def _millis(moment: Optional[datetime]) -> Optional[int]: 'time': 'datetime'} +_Reading = Tuple[Mapping[str, float], int, int] # values, second, time +_FLUSH_SECONDS = 10.0 +_MIN_READINGS = 2 + + class _SystemMetrics: - def __init__(self, log: Callable[[Mapping[str, float], int], None], - interval: float = 10.0) -> None: + def __init__(self, log: Callable[[List[_Reading]], None], + interval: Optional[float] = None) -> None: + fixed = os.environ.get('MLFLOW_SYSTEM_METRICS_SAMPLING_INTERVAL') self._log = log - self._interval = interval + self._interval = interval or float(fixed or 1) self._monitors = _metric_monitors() self._failed = False self._stop = threading.Event() @@ -598,10 +618,13 @@ def __enter__(self) -> '_SystemMetrics': def __exit__(self, *exc: Any) -> None: self._stop.set() if self._thread.is_alive(): - self._thread.join(self._interval + 5) + self._thread.join(_FLUSH_SECONDS + 5) def _run(self) -> None: - step = failures = 0 + began = flushed = time.monotonic() + readings: List[_Reading] = [] + taken = failures = 0 + step = -1 stopped = False while not stopped: stopped = self._stop.wait(self._interval) @@ -610,21 +633,27 @@ def _run(self) -> None: monitor.collect_metrics() metrics.update(monitor.aggregate_metrics()) monitor.clear_metrics() + step = max(round(time.monotonic() - began), step + 1) + readings.append(({f'system/{k}': v for k, v in metrics.items()}, + step, int(time.time() * 1000))) + taken += 1 + if taken < _MIN_READINGS or not ( + stopped or time.monotonic() - flushed >= _FLUSH_SECONDS): + continue + batch, readings, flushed = readings, [], time.monotonic() try: - self._log({f'system/{k}': v - for k, v in metrics.items()}, step) + self._log(batch) failures = 0 except Exception as error: if not self._failed: - logger.info('Could not log system metrics, skipping that ' - 'reading: %s', format_error(error)) + logger.info('Could not log system metrics, skipping those ' + 'readings: %s', format_error(error)) self._failed = True failures += 1 if failures == _SYSTEM_METRICS_FAILURES: logger.info('System metrics stopped after %d failed logs ' 'in a row', failures) return - step += 1 def _metric_monitors() -> List[Any]: diff --git a/packages/kumo-engine/test/test_mlflow_tracking.py b/packages/kumo-engine/test/test_mlflow_tracking.py index 6db9314e1..a205bdce8 100644 --- a/packages/kumo-engine/test/test_mlflow_tracking.py +++ b/packages/kumo-engine/test/test_mlflow_tracking.py @@ -437,32 +437,37 @@ def test_system_metrics_skip_a_failed_log_and_stop_after_repeated_ones( from kumo.engine.integrations import mlflow as integration - calls = [] + batches = [] - def log(metrics, step): - calls.append(step) - if step == 0 or step >= 2: + def log(readings): + batches.append([(values, step) for values, step, _ in readings]) + if len(batches) == 1 or len(batches) >= 3: raise ConnectionError('blip') monkeypatch.setattr(integration, '_metric_monitors', lambda: [_Monitor()]) + monkeypatch.setattr(integration, '_FLUSH_SECONDS', 0.0) caplog.set_level(logging.INFO, integration.logger.name) with integration._SystemMetrics(log, interval=0.001) as metrics: metrics._thread.join(10) - assert calls == list(range(2 + integration._SYSTEM_METRICS_FAILURES)) + reading = {'system/cpu': 1.0} + assert batches == [[(reading, 0), (reading, 1)], *[ + [(reading, step)] + for step in range(2, 3 + integration._SYSTEM_METRICS_FAILURES)]] assert caplog.text.count('Could not log system metrics') == 1 assert 'ConnectionError: blip' in caplog.text assert 'stopped after 6 failed logs in a row' in caplog.text -def test_system_metrics_log_a_reading_when_a_run_ends_within_the_interval( +def test_system_metrics_sample_every_second_unless_the_environment_says( monkeypatch): from kumo.engine.integrations import mlflow as integration - calls = [] monkeypatch.setattr(integration, '_metric_monitors', lambda: [_Monitor()]) - with integration._SystemMetrics(lambda metrics, step: calls.append( - (metrics, step)), interval=60): + batches = [] + with integration._SystemMetrics(batches.append) as metrics: pass - - assert calls == [({'system/cpu': 1.0}, 0)] + assert batches == [] # fewer than 2 readings + assert metrics._interval == 1.0 + monkeypatch.setenv('MLFLOW_SYSTEM_METRICS_SAMPLING_INTERVAL', '2') + assert integration._SystemMetrics(batches.append)._interval == 2.0 From 1ca53953e9b4bbd00108a91f22890c1ef0b5258e Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 03:28:22 +0000 Subject: [PATCH 211/278] Hint once per server how to open SageMaker's MLflow UI, naming the run SageMaker scopes also stop MLflow printing run links, and the registry version error names the tracker rather than MLflow. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_mlops.py | 4 ++-- .../engine/integrations/_mlflow_target.py | 4 +++- .../kumo/engine/integrations/mlflow.py | 10 +++++--- .../kumo-engine/test/test_mlflow_target.py | 23 +++++++++++++++++++ .../kumo-engine/test/test_mlflow_tracking.py | 14 ++++++++--- 5 files changed, 46 insertions(+), 9 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_mlops.py b/packages/kumo-engine/kumo/engine/_mlops.py index 713c56ea3..8abfc8532 100644 --- a/packages/kumo-engine/kumo/engine/_mlops.py +++ b/packages/kumo-engine/kumo/engine/_mlops.py @@ -120,8 +120,8 @@ def _check_registry(self, tracker: Any) -> None: check_model_alias(registry.alias) if tracker is not None and tracker.log_model is not None: if registry.version is not None: - raise ValueError('MLflow assigns registered model versions; ' - 'leave ModelRegistryLocation.version unset.') + raise ValueError('the tracker assigns registered model ' + 'versions; leave ModelRegistryLocation.version unset.') return connector = None if tracker is None else tracker.connector if connector is None: diff --git a/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py b/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py index fb956759d..82e5d9ee1 100644 --- a/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py +++ b/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py @@ -16,6 +16,7 @@ SAGEMAKER_ARN = re.compile(r'arn:aws[a-z-]*:sagemaker:([^:]*):[^:]*:' r'(mlflow-tracking-server|mlflow-app)/(.+)') _LOCK = threading.RLock() +_NO_LINKS = 'MLFLOW_SUPPRESS_PRINTING_URL_TO_STDOUT' _HEADERS = logging.getLogger('mlflow.tracking.request_header.registry') _quiet = 0 _QUIET_LOCK = threading.Lock() @@ -101,7 +102,8 @@ def explicit(self) -> bool: @contextmanager def scope(self) -> Iterator[None]: r"""Runs the block with this target's credentials.""" - with _without_header_warnings(): + links = {_NO_LINKS: 'true'} if self.platform == 'sagemaker' else {} + with _without_header_warnings(), environ(links): if self._swap is None: yield else: diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index fd82f8eb7..8ab2e4d34 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -84,6 +84,7 @@ _PATH_CONNECTORS = ('LOCAL_FILE', 'S3', 'GCS') _SAGEMAKER_MODEL = re.compile(r'[a-zA-Z0-9](-*[a-zA-Z0-9]){0,56}') _UC_FORBIDDEN = re.compile(r'[ /\x00-\x1f\x7f]') +_HINTED: set = set() class MlflowTracking(TrackingBase): @@ -404,7 +405,6 @@ def finish(self, record: Optional[RunRecord], status = 'FAILED' kind = None if record is None else record.tags.get('kumo.kind') started = self._parent is not None - fit_run = started and kind == 'fit' ended = None if record is None else _millis(record.finished) try: if record is not None and self._parent is not None: @@ -427,8 +427,12 @@ def finish(self, record: Optional[RunRecord], status, ended) self._reset() hint = _sagemaker_ui_hint(self.target.tracking_uri) - if fit_run and hint: - warnings.warn(hint, UserWarning, stacklevel=2) + if (hint and started and kind in ('fit', 'prediction') + and self.target.tracking_uri not in _HINTED): + _HINTED.add(self.target.tracking_uri) + warnings.warn(f'MLflow run {self.logged_run_name!r} in experiment ' + f'{self.experiment!r}: {hint}', UserWarning, + stacklevel=2) def _fit_of(self, fit_id: Optional[str]) -> Tuple[Any, Any]: r"""The fit run's info and its logged model, each None if missing.""" diff --git a/packages/kumo-engine/test/test_mlflow_target.py b/packages/kumo-engine/test/test_mlflow_target.py index 731e91723..d461a61a0 100644 --- a/packages/kumo-engine/test/test_mlflow_target.py +++ b/packages/kumo-engine/test/test_mlflow_target.py @@ -207,6 +207,29 @@ def lookup(self, name): assert mlflow.get_tracking_uri() == tracking_uri +NO_LINKS = 'MLFLOW_SUPPRESS_PRINTING_URL_TO_STDOUT' + + +def test_sagemaker_calls_print_no_mlflow_run_links(monkeypatch): + from contextlib import nullcontext + + monkeypatch.delenv(NO_LINKS, raising=False) + seen = [] + for uri in (ARN, 'http://mlflow'): + with target_module.Target(uri, uri, swap=nullcontext).scope(): + seen.append(os.environ.get(NO_LINKS)) + + assert (seen, os.environ.get(NO_LINKS)) == (['true', None], None) + + +def test_pin_mlflow_reads_its_run_link_switch_per_call(): + from mlflow.environment_variables import ( + MLFLOW_SUPPRESS_PRINTING_URL_TO_STDOUT, + ) + + assert MLFLOW_SUPPRESS_PRINTING_URL_TO_STDOUT.name == NO_LINKS + + def test_an_unreachable_server_fails_when_mlops_is_built( tmp_path, monkeypatch): from kumo.engine.integrations.mlflow import MlflowTracking diff --git a/packages/kumo-engine/test/test_mlflow_tracking.py b/packages/kumo-engine/test/test_mlflow_tracking.py index a205bdce8..65dded1d3 100644 --- a/packages/kumo-engine/test/test_mlflow_tracking.py +++ b/packages/kumo-engine/test/test_mlflow_tracking.py @@ -296,9 +296,12 @@ def refuse(self, name, *args, **kwargs): ('sqlite:///mlflow.db', None), ]) def test_a_tracked_fit_on_sagemaker_says_how_to_open_the_ui( - tracking_uri, uri, command): + tracking_uri, uri, command, monkeypatch): + from kumo.engine.integrations import mlflow as integration from kumo.engine.integrations.mlflow import MlflowTracking + monkeypatch.setattr(integration, '_HINTED', set()) + tracker = MlflowTracking(log_system_metrics=False) tracker.start(run_record(fit_spec(CONTEXT))) tracker.target.tracking_uri = uri @@ -307,8 +310,13 @@ def test_a_tracked_fit_on_sagemaker_says_how_to_open_the_ui( tracker.finish(run_record(fit_spec(CONTEXT)), None) hints = [str(w.message) for w in caught if 'Studio' in str(w.message)] assert hints == ([] if command is None else [ - "MLflow's run links do not open on SageMaker: open the MLflow UI " - f'from SageMaker Studio or at the URL `{command}` prints.']) + "MLflow run 'fit-ents-regression-20261029-061800' in experiment " + "'Default': MLflow's run links do not open on SageMaker: open the " + f'MLflow UI from SageMaker Studio or at the URL `{command}` prints.']) + with warnings.catch_warnings(record=True) as again: + warnings.simplefilter('always') + tracker.finish(run_record(fit_spec(CONTEXT)), None) + assert not [w for w in again if 'Studio' in str(w.message)] def test_user_disabled_autologging_stays_disabled(tracking_uri): From 8bce1d4d391cfcc5a6e08e64ad030edd2f1ed8d0 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 03:30:34 +0000 Subject: [PATCH 212/278] Track trial runs only for searches, under the fit's snapshot and task type A single-trial fit logs its epochs on the fit run alone, and searches no longer log a search/best_* series. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_tracking.py | 28 ++++++++++++------- .../test/test_snowflake_tracking.py | 2 +- .../test/test_tracking_contract.py | 14 ++++++---- 3 files changed, 27 insertions(+), 17 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_tracking.py b/packages/kumo-engine/kumo/engine/_tracking.py index 1e3629bec..79551a8c7 100644 --- a/packages/kumo-engine/kumo/engine/_tracking.py +++ b/packages/kumo-engine/kumo/engine/_tracking.py @@ -129,6 +129,9 @@ def by_step(steps: Mapping[str, Sequence[Tuple[int, float]]] return sorted(grouped.items()) +_TRIAL_PARENTS = ('snapshot_id', 'task_type') + + def run_record(spec: RunSpec) -> RunRecord: params = {key: _param_value(value) for key, value in spec.params.items()} return RunRecord( @@ -280,7 +283,6 @@ class _Search: metric: str lower: bool best: Optional[Tuple[float, int]] = None - trials: int = 0 @classmethod def of(cls, params: Mapping[str, Any]) -> '_Search': @@ -301,12 +303,11 @@ def on(self, event: Event) -> Optional[Tuple[Dict[str, float], int]]: if self.best is None or (value < self.best[0] if self.lower else value > self.best[0]): self.best = (value, event.trial or 0) - self.trials += 1 - return {f'search/best_{self.metric}': self.best[0]}, self.trials + return None class RegistrationError(Exception): - r"""Registering, not logging, the model failed.""" + r"""Registering the logged model failed.""" class TrackingSession: @@ -371,11 +372,18 @@ def on_event(self, event: Event) -> None: **event.params, 'attempt': events[0].params['attempt'] }) events.append(event) - record = run_record(trial_spec(events)) - logged = dataclasses.replace(event, - metrics=logged_names(event.metrics)) - self._submit(lambda: self._call( - lambda tracker: tracker.on_trial(record, logged))) + if self._search is not None and not self._search.single: + trial = trial_spec(events) + record = run_record(dataclasses.replace( + trial, model_label=self.spec.model_label, parents={ + **{name: self.spec.parents.get(name) + for name in _TRIAL_PARENTS}, + **trial.parents + })) + logged = dataclasses.replace(event, + metrics=logged_names(event.metrics)) + self._submit(lambda: self._call( + lambda tracker: tracker.on_trial(record, logged))) series = None if self._search is None else self._search.on(event) if series is not None: metrics, step = logged_names(series[0]), series[1] @@ -443,7 +451,7 @@ def log() -> None: def finish(self, spec: Optional[Callable[[], RunSpec]], error: Optional[BaseException]) -> None: - r"""Ends the run with ``spec()`` (`None`: the started spec).""" + r"""Ends the run, failed if error.""" worker = self._worker if self.spec is None or worker is None: return diff --git a/packages/kumo-engine/test/test_snowflake_tracking.py b/packages/kumo-engine/test/test_snowflake_tracking.py index 37eae359d..4ccc5ea0e 100644 --- a/packages/kumo-engine/test/test_snowflake_tracking.py +++ b/packages/kumo-engine/test/test_snowflake_tracking.py @@ -189,7 +189,7 @@ def test_metric_and_param_limits(fake, monkeypatch): run, = fake._instance.runs.values() assert sorted({name for name, _, _ in run['metrics']}) == [ 'TRIAL_1/val/best_mae', 'TRIAL_2_ATTEMPT_2/val/best_mae', - 'duration_seconds', 'search/best_val_mae', 'test/mae' + 'duration_seconds', 'test/mae' ] assert run['params']['kumo.dropped_metrics'] == '2' assert run['params']['kumo.dropped_params'] == '1' diff --git a/packages/kumo-engine/test/test_tracking_contract.py b/packages/kumo-engine/test/test_tracking_contract.py index 0fcfba8e5..e98e642dc 100644 --- a/packages/kumo-engine/test/test_tracking_contract.py +++ b/packages/kumo-engine/test/test_tracking_contract.py @@ -81,7 +81,6 @@ def _epochs(trial, *values): _event('fit_ended', metrics={'val_mae': 0.3, 'test_mae': 0.32, 'seconds': 9.0}, status='succeeded'), ] -SEARCH = {'search/best_val_mae': ((1, 0.35), (2, 0.3))} SELECTED = {'train_loss': ((1, 0.6), (2, 0.4)), 'val_mae': ((1, 0.5), (2, 0.3))} BUNDLE = 'b' * 64 RESULT = SimpleNamespace(metrics={'val_mae': 0.3, 'test_mae': 0.32}, @@ -144,7 +143,7 @@ def replay(tracker=None, events=EVENTS, result=RESULT, error=None, if result is not None: spec = fitted_spec(spec, result) spec = replace(spec, steps=series if series is not None else { - **SEARCH, **({} if error else SELECTED)}) + **({} if error else SELECTED)}) if series is None and not error: spec = replace(spec, parents={**spec.parents, 'best_trial': 2}) if 'error' in events[-1].params: @@ -376,8 +375,7 @@ def test_records_of_the_recorded_fit(): 'kumo.bundle_id': 'x', 'training_job.tune_metric': 'y', 'a': 1})) assert spoofed.tags.keys() - fit.tags.keys() == {'a'} assert fit.metrics == {'test/mae': 0.32} - assert fit.steps == {'search/best_val_mae': ((1, 0.35), (2, 0.3)), - 'train/loss': ((1, 0.6), (2, 0.4)), + assert fit.steps == {'train/loss': ((1, 0.6), (2, 0.4)), 'val/mae': ((1, 0.5), (2, 0.3))} registered = SimpleNamespace(**{ **vars(RESULT), 'registered_alias': 'candidate', @@ -391,9 +389,11 @@ def test_records_of_the_recorded_fit(): oom = records['trial-2-attempt-1'] assert oom == RunRecord( - name='trial-2-attempt-1', + name='trial-2-attempt-1', model_label='ents-regression', params={'optimization.base_lr': '0.001', 'attempt': '1'}, tags={ 'kumo.kind': 'trial', + 'kumo.snapshot_id': SNAPSHOT, + 'kumo.task_type': 'regression', 'kumo.status': 'failed', 'kumo.error': 'OutOfMemoryError: CUDA out of memory', }) @@ -413,7 +413,8 @@ def test_adapter_records_the_same_runs(adapter, tmp_path): assert gaps(expected, adapter.observe(), adapter) == [] -def test_adapter_records_a_single_trial_fit_epoch_by_epoch(adapter): +def test_adapter_records_a_single_trial_fit_epoch_by_epoch_without_trials( + adapter): plan = {**CONTEXT.model_plan, 'training_job': {'tune_metric': 'mae', 'num_experiments': 1}} events = [EVENTS[0], *_epochs(1, (0.5, 0.4), (0.3, 0.35)), EVENTS[3], @@ -423,6 +424,7 @@ def test_adapter_records_a_single_trial_fit_epoch_by_epoch(adapter): 'train_loss': ((1, 0.5), (2, 0.3)), 'val_mae': ((1, 0.4), (2, 0.35))}) + assert list(expected) == [FIT] assert gaps(expected, adapter.observe(), adapter) == [] From e6df805a2300a5da8aba7a624b927db76dd9fa62 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 03:35:31 +0000 Subject: [PATCH 213/278] Tag every run with the engine version and a fit with its entity Fit params skip unset plan values, and the tune metric uses the logged metric name. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_tracking.py | 2 ++ .../kumo/engine/trainer/_run_spec.py | 17 ++++++++++------- .../kumo-engine/test/test_mlflow_tracking.py | 6 +++--- .../test/test_model_agnostic_mlops.py | 6 ++++-- .../kumo-engine/test/test_tracking_contract.py | 3 ++- 5 files changed, 21 insertions(+), 13 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_tracking.py b/packages/kumo-engine/kumo/engine/_tracking.py index 79551a8c7..5131d62d9 100644 --- a/packages/kumo-engine/kumo/engine/_tracking.py +++ b/packages/kumo-engine/kumo/engine/_tracking.py @@ -32,6 +32,7 @@ from kumo.engine._events import Event, _finite_floats from kumo.engine._mlops import full_error, warn_failure +from kumo.engine._version import __version__ if TYPE_CHECKING: from kumoapi.plugin import ModelRegistryLocation, PublishedModel @@ -143,6 +144,7 @@ def run_record(spec: RunSpec) -> RunRecord: tags={ **_user_tags(spec.tags, params), 'kumo.kind': spec.kind, + 'kumo.engine_version': __version__, **{f'kumo.{name}': str(value) for name, value in spec.parents.items() if value}, **({} if spec.status is None else {'kumo.status': spec.status}), diff --git a/packages/kumo-engine/kumo/engine/trainer/_run_spec.py b/packages/kumo-engine/kumo/engine/trainer/_run_spec.py index 78bcfe4e8..e0fde9478 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_run_spec.py +++ b/packages/kumo-engine/kumo/engine/trainer/_run_spec.py @@ -13,9 +13,8 @@ from kumo_core.summaries import Scalar, Summary from kumo.engine._run_names import fit_name, model_label -from kumo.engine._text import metric_key -from kumo.engine._tracking import RunSpec, logged_names -from kumo.engine._version import __version__ +from kumo.engine._text import entity, metric_key +from kumo.engine._tracking import RunSpec, logged_names, metric_name @dataclass(frozen=True) @@ -30,20 +29,24 @@ class FitContext: model_name: Optional[str] = None +_UNSET = '???' + + def fit_spec(context: FitContext) -> RunSpec: started = context.started or datetime.now(timezone.utc) model = model_label(context.model_name, context.query, context.task_type) + tune = (context.model_plan.get('training_job') or {}).get('tune_metric') return RunSpec( kind='fit', name=fit_name(model, started), model_label=model, started=started, - params=dict(_flatten(context.model_plan)), + params={key: value for key, value in _flatten(context.model_plan) + if value is not None and value != _UNSET}, parents={'snapshot_id': context.snapshot_id, 'task_type': context.task_type, - 'tune_metric': (context.model_plan.get('training_job') - or {}).get('tune_metric'), - 'engine_version': __version__}, + 'entity': entity(context.query), + 'tune_metric': tune and metric_name(tune)}, artifacts={ 'query.pql': context.query, 'model_plan.json': json.dumps(dict(context.model_plan), indent=2, diff --git a/packages/kumo-engine/test/test_mlflow_tracking.py b/packages/kumo-engine/test/test_mlflow_tracking.py index 65dded1d3..124e6b61a 100644 --- a/packages/kumo-engine/test/test_mlflow_tracking.py +++ b/packages/kumo-engine/test/test_mlflow_tracking.py @@ -125,9 +125,9 @@ def test_fit_logs_a_run_per_fit_and_per_trial(table_and_snapshot, tmp_path, assert parent.info.status == 'FINISHED' assert {k for k in parent.data.tags if k.startswith('kumo.')} == { 'kumo.kind', 'kumo.bundle_id', 'kumo.dropped_params', - 'kumo.snapshot_id', 'kumo.task_type', 'kumo.tune_metric', - 'kumo.engine_version', 'kumo.best_trial', 'kumo.bundle_uri', - 'kumo.artifact_root'} + 'kumo.snapshot_id', 'kumo.task_type', 'kumo.entity', + 'kumo.tune_metric', 'kumo.engine_version', 'kumo.best_trial', + 'kumo.bundle_uri', 'kumo.artifact_root'} assert parent.data.tags['kumo.bundle_uri'] == str( tmp_path / 'root' / 'models' / result.id) assert parent.data.tags['kumo.snapshot_id'] == snapshot.id diff --git a/packages/kumo-engine/test/test_model_agnostic_mlops.py b/packages/kumo-engine/test/test_model_agnostic_mlops.py index ef94de3b8..9b3563b29 100644 --- a/packages/kumo-engine/test/test_model_agnostic_mlops.py +++ b/packages/kumo-engine/test/test_model_agnostic_mlops.py @@ -34,6 +34,7 @@ from kumo.engine._manifest import ModelManifest, read_manifest, write_manifest from kumo.engine._predictor import load_predictor from kumo.engine._tracking import RunSpec, TrackingSession, run_record +from kumo.engine._version import __version__ from kumo.engine.tracking import Event from kumo.engine.trainer._spec import MANIFEST_FILE, PROVENANCE_ROLE from kumo.engine.trainer.trainer import _store_and_register @@ -96,9 +97,10 @@ def test_any_run_is_tracked_from_its_spec(): started, trial, finished = (record for _, record in tracker.calls) assert started == run_record(spec) assert trial.name == 'trial-1-attempt-1' - assert trial.tags == {'kumo.kind': 'trial'} + assert trial.tags == {'kumo.kind': 'trial', 'kumo.engine_version': __version__} assert finished.metrics == {'auc': 0.9} - assert finished.tags == {'kumo.kind': 'custom', 'kumo.bundle_id': 'b'} + assert finished.tags == {'kumo.kind': 'custom', 'kumo.bundle_id': 'b', + 'kumo.engine_version': __version__} class _RegistryTracker(_Tracker): diff --git a/packages/kumo-engine/test/test_tracking_contract.py b/packages/kumo-engine/test/test_tracking_contract.py index e98e642dc..8c91f930f 100644 --- a/packages/kumo-engine/test/test_tracking_contract.py +++ b/packages/kumo-engine/test/test_tracking_contract.py @@ -360,7 +360,7 @@ def test_records_of_the_recorded_fit(): assert fit.tags == {'team': 'ents', 'kumo.kind': 'fit', 'kumo.bundle_id': BUNDLE, 'kumo.snapshot_id': SNAPSHOT, 'kumo.task_type': 'regression', - 'kumo.tune_metric': 'mae', + 'kumo.entity': 'ents', 'kumo.tune_metric': 'mae', 'kumo.engine_version': __version__, 'kumo.best_trial': '2'} assert fit.params == { @@ -392,6 +392,7 @@ def test_records_of_the_recorded_fit(): name='trial-2-attempt-1', model_label='ents-regression', params={'optimization.base_lr': '0.001', 'attempt': '1'}, tags={ 'kumo.kind': 'trial', + 'kumo.engine_version': __version__, 'kumo.snapshot_id': SNAPSHOT, 'kumo.task_type': 'regression', 'kumo.status': 'failed', From 5f3d13b965c51566c848fed465c0a842e8a5bfb8 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 03:42:13 +0000 Subject: [PATCH 214/278] Summarize a fit's data as splits, health and column tables plus problem notes Data warnings become one Note per problem with a count tag; a manifest's artifacts are a locations table, its datasets are named by kind and date, and the snapshot and training table get schema and size details. Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/_data_summaries.py | 265 ++++++++++++------ packages/kumo-engine/kumo/engine/_manifest.py | 36 ++- .../kumo/engine/trainer/trainer.py | 26 +- packages/kumo-engine/test/test_artifacts.py | 28 +- .../kumo-engine/test/test_data_summaries.py | 97 ++++--- .../kumo-engine/test/test_mlflow_tracking.py | 6 +- 6 files changed, 300 insertions(+), 158 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_data_summaries.py b/packages/kumo-engine/kumo/engine/_data_summaries.py index 9e24d9a64..b34f065ee 100644 --- a/packages/kumo-engine/kumo/engine/_data_summaries.py +++ b/packages/kumo-engine/kumo/engine/_data_summaries.py @@ -6,69 +6,126 @@ import json import logging from pathlib import Path -from typing import Any, Dict, List, Mapping, Optional +from typing import Any, Dict, List, Mapping, Optional, Tuple +from kumo_core.artifacts import PROVENANCE_FILE from kumo_core.instrumentation import format_error -from kumo_core.summaries import Scalar, Summary, Table, Tag +from kumo_core.summaries import Note, Summary, Table, Tag from kumo.engine._manifest import Details +from kumo.engine._text import entity logger = logging.getLogger(__name__) _SPLITS = ('train', 'val', 'test') -_LABEL_STATS = {'target_mean': 'label_mean', 'target_median': 'label_q50', - 'target_min': 'label_min', 'target_max': 'label_max'} _COLUMN_STATS = ('count', 'null_rate', 'num_unique', 'min', 'max', 'mean', 'std', 'q25', 'q50', 'q75', 'avg_tokens', 'max_tokens') +_TITLES = {'null_rate': 'null rate', 'num_unique': 'distinct≈', + 'q25': 'q25≈', 'q50': 'q50≈', 'q75': 'q75≈', + 'avg_tokens': 'avg tokens≈', 'max_tokens': 'max tokens≈', + 'top': 'top values≈'} _TOP_CATEGORIES = 10 +_TARGET_BINS = 50 # as DiskGraph's node-regression histogram +_QUARTILES = ('target_min', 'target_percentile_25', 'target_median', + 'target_percentile_75', 'target_max') _EDGE_STATS = {'FRAC_SRC_NONZERO_DEGREE': 'src_match_frac', 'FRAC_DST_NONZERO_DEGREE': 'dst_match_frac'} +_HEALTH = ('table or link', 'rows', 'columns', 'dropped rows', + 'duplicate keys', 'rows without time', 'edges', + 'source matched', 'target matched', 'failed lookups') _BAD_TIME_FRACTION = 0.05 _INVALID_FRACTION = 0.10 _MIN_MATCH_FRACTION = 0.20 _MIN_ROWS = 10 -_TAG_LENGTH = 5000 def split_summaries(staging_dir: Path) -> List[Summary]: - r"""Rows, entities and label balance per split.""" + r"""Brief data/splits table and a note on a numerical target's bins.""" stats = _read(staging_dir / 'split_stats.json') + tables = {split: stats.get(f'{split}_table_stats') for split in _SPLITS} + tables = {split: table for split, table in tables.items() + if isinstance(table, dict)} + targets = {split: table.get('target_distribution') or {} + for split, table in tables.items()} + if any(_number(target.get('target_mean')) is not None + for target in targets.values()): + columns, label = ('entities', 'target mean'), { + 'target mean': 'target_mean'} + elif any(set(target.get('count') or {}) <= {'true', 'false'} + and (target.get('count') or {}) for target in targets.values()): + columns, label = ('entities', 'positive rate'), {} + else: + columns, label = ('entities',), {} + rows = tuple( + (split, _whole(table.get('num_train_rows')), *( + _value(table, targets[split], name, label) for name in columns)) + for split, table in tables.items()) + query = _read(staging_dir / PROVENANCE_FILE).get('query') + entities = f"distinct {entity(query) or 'entities'} (approx.)" summaries: List[Summary] = [] - for split in _SPLITS: - table = stats.get(f'{split}_table_stats') - if not isinstance(table, dict): - continue - prefix = f'data/{split}' - for key, name in (('num_train_rows', 'rows'), - ('num_entities', 'entities')): - if _number(table.get(key)) is not None: - summaries.append(Scalar(f'{prefix}/{name}', table[key])) - target = table.get('target_distribution') or {} + if rows: + summaries.append(Table('data/splits', ('split', 'rows', *( + entities if name == 'entities' else name for name in columns)), + rows, brief=True)) + note = _bins_note((targets.get('train') or {})) + return summaries + ([Note('notes', note)] if note else []) + + +def _value(table: Mapping[str, Any], target: Mapping[str, Any], name: str, + label: Mapping[str, str]) -> Optional[float]: + if name == 'entities': + return _at_most(_whole(table.get('num_entities')), + _whole(table.get('num_train_rows'))) + if name == 'positive rate': counts = target.get('count') or {} - if sum(counts.values()): - summaries.append(Scalar(f'{prefix}/positive_rate', - counts.get('true', 0) / - sum(counts.values()))) - summaries += [ - Scalar(f'{prefix}/{name}', target[key]) - for key, name in _LABEL_STATS.items() - if _number(target.get(key)) is not None - ] - return summaries + total = sum(counts.values()) + return counts.get('true', 0) / total if total else None + return _number(target.get(label[name])) + + +def _target_bins(target: Mapping[str, Any]) -> Tuple[Tuple[float, ...], str]: + bins = (target.get('histogram') or {}).get('bins') or () + source = f"the training target's {len(bins) - 1} histogram bins" + low, high = (_number(target.get(key)) + for key in ('target_min', 'target_max')) + if not bins and low is not None and high is not None and high > low: + bins = [low + (high - low) * i / _TARGET_BINS + for i in range(_TARGET_BINS + 1)] + source = f'{_TARGET_BINS} equal bins' + edges = {*bins, *(target[key] for key in _QUARTILES + if _number(target.get(key)) is not None)} + parts = [source] if bins else [] + extra = len(edges) - len(set(bins)) + if extra: + parts.append(f"{extra} quartile edge{'s' * (extra > 1)}") + return ((tuple(sorted(edges)), ' + '.join(parts)) if len(edges) > 1 + else ((), '')) + + +def _bins_note(target: Mapping[str, Any]) -> Optional[str]: + edges, source = _target_bins(target) + if not edges: + return None + return f'Target histograms: {source}.' def snapshot_summaries(index: Path, graph_config: Mapping[str, Any]) -> List[Summary]: - r"""The snapshot's tables, relations, data quality and column stats.""" + r"""data/health, data/columns and a problems note per failed data check.""" ingest = _read(index / 'ingest_stats.json').get('stats') or {} columns: Dict[str, Dict[str, Any]] = _read(index / 'col_stats.json') relations = _relations(_read(index / 'edge_stats.json'), ingest) profiles = _profiles(columns) + problems = _problems(ingest, columns, relations, graph_config) return [ - *_graph_scalars(ingest, relations), - *_quality(ingest, columns, relations, graph_config), - Table('graph/col_stats', - ('table', 'column', 'stype', *_COLUMN_STATS, 'top'), + _health(ingest, columns, relations, graph_config), + *(Note('problems', problem) for problem in problems), + *([Tag('kumo.data_warnings', str(len(problems)))] if problems + else []), + Table('data/columns', + ('table', 'column', 'stype', + *(_TITLES.get(name, name) for name in _COLUMN_STATS), + _TITLES['top']), tuple((table, column, *profile.values()) for table, table_profiles in profiles.items() for column, profile in table_profiles.items())), @@ -106,39 +163,52 @@ def _relations(edge_stats: Mapping[str, Any], return relations -def _graph_scalars(ingest: Mapping[str, Any], - relations: Mapping[str, Mapping[str, float]] - ) -> List[Summary]: - rows = ingest.get('node_counts') or {} - return [ - *(Scalar(f'graph/rows/{table}', count) - for table, count in sorted(rows.items())), - *(Scalar(f'graph/edges/{relation}/{name}', value) - for relation, stats in sorted(relations.items()) - for name, value in stats.items() if _number(value) is not None), - Scalar('graph/tables', len(rows)), - Scalar('graph/edge_types', len(relations)), - ] +def _links(table: Mapping[str, Any], + relations: Mapping[str, Mapping[str, float]] + ) -> List[Tuple[Mapping[str, Any], Mapping[str, float]]]: + return [(fkey, relations.get( + f"{table['table_name']}_{fkey['name']}_to_{fkey['dst_table']}_" + f"{fkey['dst_key']}") or {}) for fkey in table.get('fkeys') or ()] -def _quality(ingest: Mapping[str, Any], columns: Mapping[str, Any], - relations: Mapping[str, Mapping[str, float]], - graph_config: Mapping[str, Any]) -> List[Summary]: - r"""Per-table data quality, and ``kumo.data_warnings`` if any fail.""" +def _health(ingest: Mapping[str, Any], columns: Mapping[str, Any], + relations: Mapping[str, Mapping[str, float]], + graph_config: Mapping[str, Any]) -> Table: duplicates = _by_table(ingest.get('pk_duplicate_counts'), ':') null_time = _by_table(ingest.get('pk_null_timestamp_drops'), ':shard-') dropped = ingest.get('dropped_rows_counts') or {} + nodes = ingest.get('node_counts') or {} + rows: List[Tuple[Any, ...]] = [] + tables = graph_config.get('tables') or () + for table in tables: + name = table['table_name'] + rows.append((name, nodes.get(name), len(columns.get(name) or {}), + dropped.get(name, 0), duplicates.get(name, 0), + null_time.get(name, 0) if table.get('time_col') else '', + '', '', '', '')) + for table in tables: + for fkey, relation in _links(table, relations): + rows.append((_link(table, fkey), '', '', '', '', '', + *(relation.get(key) for key in ( + 'count', 'src_match_frac', 'dst_match_frac', + 'fk_lookup_failures')))) + return Table('data/health', _HEALTH, tuple(rows)) + + +def _link(table: Mapping[str, Any], fkey: Mapping[str, Any]) -> str: + return f"{table['table_name']}.{fkey['name']} → {fkey['dst_table']}" + + +def _problems(ingest: Mapping[str, Any], columns: Mapping[str, Any], + relations: Mapping[str, Mapping[str, float]], + graph_config: Mapping[str, Any]) -> List[str]: + duplicates = _by_table(ingest.get('pk_duplicate_counts'), ':') + null_time = _by_table(ingest.get('pk_null_timestamp_drops'), ':shard-') sources = ingest.get('source_rows_by_table') or {} failures = ingest.get('key_failures_by_context') or {} - summaries: List[Summary] = [] - warnings: List[str] = [] + problems: List[str] = [] for table in graph_config.get('tables') or (): name = table['table_name'] - counts = {'duplicate_pks': duplicates.get(name, 0), - 'dropped_rows': dropped.get(name, 0), - 'null_time_dropped': null_time.get(name, 0)} - summaries += [Scalar(f'graph/{name}/{key}', value) - for key, value in counts.items()] stats = columns.get(name) or {} nodes = (ingest.get('node_counts') or {}).get(name) rows = sources.get(name) or nodes @@ -147,56 +217,82 @@ def _quality(ingest: Mapping[str, Any], columns: Mapping[str, Any], or {}).get('parse_failures', 0) time = (table.get('time_col') or {}).get('name') time_stats = stats.get(time) or {} - bad_times = (counts['null_time_dropped'] + + bad_times = (null_time.get(name, 0) + time_stats.get('null_count', 0) + time_stats.get('parse_fail_count', 0)) invalid = [column for column, column_stats in stats.items() if column_stats.get('count') and _invalid(column_stats) >= _INVALID_FRACTION * column_stats['count']] - warnings += [ - *([f'{name}: {bad_times / rows:.1%} missing or invalid time'] + problems += [ + *([f'`{name}`: {bad_times / rows:.1%} of rows have a missing or ' + f'invalid time'] if time and rows and bad_times >= _BAD_TIME_FRACTION * rows else []), - *([f'{name}: {bad_pks} missing or invalid primary keys'] + *([f'`{name}`: {bad_pks} missing or invalid primary keys'] if bad_pks else []), - *([f"{name}: {counts['duplicate_pks']} duplicate primary keys"] - if counts['duplicate_pks'] else []), - *(f'{name}.{column}: at least {_INVALID_FRACTION:.0%} invalid ' + *([f'`{name}`: {duplicates[name]} duplicate primary keys'] + if duplicates.get(name) else []), + *(f'`{name}.{column}`: at least {_INVALID_FRACTION:.0%} invalid ' f'values' for column in sorted(invalid)), - *([f'{name}: {nodes} rows'] if nodes is not None + *([f'`{name}`: only {nodes} rows'] if nodes is not None and nodes < _MIN_ROWS else []), ] - for fkey in table.get('fkeys') or (): - relation = relations.get( - f"{name}_{fkey['name']}_to_{fkey['dst_table']}_" - f"{fkey['dst_key']}") or {} + for fkey, relation in _links(table, relations): for side, table_name in (('src', name), ('dst', fkey['dst_table'])): match = _number(relation.get(f'{side}_match_frac')) if match is not None and match < _MIN_MATCH_FRACTION: - warnings.append( - f"{name}.{fkey['name']} -> {fkey['dst_table']}: " - f"{match:.1%} of {table_name} rows match") - if warnings: - summaries.append(Tag('kumo.data_warnings', - '; '.join(warnings)[:_TAG_LENGTH])) - return summaries + problems.append( + f'link `{_link(table, fkey)}`: {match:.1%} of ' + f'`{table_name}` rows match') + return problems -def table_details(index: Path) -> Details: - r"""Each source table's schema and profile, for its `Dataset`.""" +def table_details(index: Path, snapshot_id: Optional[str] = None) -> Details: + r"""Each source table's schema and profile, and the snapshot's.""" rows = (_read(index / 'ingest_stats.json').get('stats') or {}).get('node_counts') or {} profiles = _profiles(_read(index / 'col_stats.json')) - return { + tables = _read(index / 'diskgraph_config.json').get('tables') or () + snapshot = {} if snapshot_id is None else {snapshot_id: ({}, { + 'tables': len(tables), + 'rows': sum(rows.get(table.get('name'), 0) for table in tables)})} + return {**snapshot, **{ table.get('kumo_source_table_name') or table.get('name'): ( {column['name']: {'stype': column.get('semantic_type'), 'dtype': column.get('source_dtype')} for column in table.get('columns') or ()}, {'rows': rows.get(table.get('name')), 'columns': profiles.get(table.get('name'), {})}) - for table in _read(index / 'diskgraph_config.json').get('tables') or () - } + for table in tables + }} + + +def training_table_details(table: Any) -> Details: + r"""A training table's schema and profile from its Parquet metadata.""" + from pyarrow import parquet + + manifest = table.manifest + path = Path(table.train_table_path) + if manifest is None or not path.exists(): + return {} + parts = sorted(path.rglob('*.parquet')) if path.is_dir() else [path] + if not parts: + return {} + schema = parquet.read_schema(parts[0]) + return {manifest.id: ( + {field.name: {'stype': None, 'dtype': _dtype(field.type)} + for field in schema}, + {'rows': sum(parquet.ParquetFile(part).metadata.num_rows + for part in parts), + 'columns': len(schema)})} + + +def _dtype(arrow_type: Any) -> str: + text = str(arrow_type) + return {'double': 'float64', 'float': 'float32'}.get(text, ( + 'time' if text.startswith('timestamp') else + 'date' if text.startswith('date') else text)) def _by_table(counts: Optional[Mapping[str, int]], @@ -219,6 +315,7 @@ def _profile(stats: Mapping[str, Any]) -> Dict[str, Any]: value = (stats.get('null_count', 0) / count if count else None) \ if name == 'null_rate' else _number(stats.get(name)) profile[name] = value + profile['num_unique'] = _at_most(profile['num_unique'], count) top = () if stats.get('semantic_type') == 'id' else ( stats.get('category_counts') or ()) profile['top'] = ', '.join(f'{value}: {n}' @@ -232,6 +329,16 @@ def _number(value: Any) -> Optional[float]: return float(value) +def _at_most(estimate: Optional[float], + exact: Optional[float]) -> Optional[float]: + return None if estimate is None else ( + estimate if exact is None else min(estimate, exact)) + + +def _whole(value: Any) -> Optional[int]: + return None if _number(value) is None else int(value) + + def _read(path: Path) -> Dict[str, Any]: try: value = json.loads(path.read_text()) diff --git a/packages/kumo-engine/kumo/engine/_manifest.py b/packages/kumo-engine/kumo/engine/_manifest.py index 4e513829e..d9985f78c 100644 --- a/packages/kumo-engine/kumo/engine/_manifest.py +++ b/packages/kumo-engine/kumo/engine/_manifest.py @@ -4,6 +4,7 @@ """Shared JSON manifest format for engine artifacts.""" import json +import re import types import warnings from collections.abc import Mapping as AbcMapping @@ -28,7 +29,7 @@ ) from kumo_core.io.filesystem import atomic_write_bytes -from kumo_core.summaries import Dataset, Summary, Tag +from kumo_core.summaries import Dataset, Summary, Table, Tag SCHEMA_VERSION = 1 _OMIT_NONE = 'omit_none' @@ -55,6 +56,7 @@ class Artifact: version: Optional[str] = None version_kind: Optional[str] = None context: Optional[str] = None + label: Optional[str] = None @dataclass(frozen=True, kw_only=True) @@ -173,8 +175,10 @@ class Manifest(Record): def as_artifact(self, key: str) -> Artifact: return Artifact(self.kind, self.id, self.location, self.IS_DATASET, - self.URI_TAG or f'kumo.{self.kind}_uri', - version=self.id, version_kind='exact') + self.URI_TAG, version=self.id, version_kind='exact', + label=self.created_at and + f"{self.kind.replace('_', '-')}-{self.created_at[:10]}-" + f"{re.split('[@-]', self.id)[-1][:6]}") def artifacts(manifest: Manifest) -> List[Artifact]: @@ -209,19 +213,21 @@ def _walk(value: Any, key: str, root: Manifest) -> Iterator[Artifact]: def artifact_summaries(manifest: Manifest, context: str, - settings: Any = None, details: Optional[Details] = None) -> List[Summary]: - r"""Tags locating ``manifest``'s artifacts, and its `Dataset`s.""" + r"""Location tags, the locations table and datasets of a manifest.""" located = [item for item in artifacts(manifest) if item.location] tags: Dict[str, str] = {} for item in located: if item.tag: tags.setdefault(item.tag, str(item.location)) - if settings is not None: - tags['kumo.artifact_root'] = settings.root.uri('').rstrip('/') - summaries: List[Summary] = [Tag(name, value) - for name, value in tags.items()] - for item in located: + summaries: List[Summary] = [ + *(Tag(name, value) for name, value in tags.items()), + *([Table('locations', ('kind', 'id', 'location'), tuple( + (item.kind, item.id, str(item.location)) for item in located), + brief=True)] if located else []), + ] + for item in sorted(located, key=lambda item: _DATASET_ORDER.get( + item.kind, len(_DATASET_ORDER))): if not item.is_dataset: continue location = str(item.location) @@ -229,11 +235,14 @@ def artifact_summaries(manifest: Manifest, context: str, 'location': location, 'version': item.version, 'version_kind': item.version_kind} schema, profile = (details or {}).get(item.id, ({}, {})) - summaries.append(Dataset(item.id, source, schema, profile, - item.context or context)) + summaries.append(Dataset(item.label or item.id, source, schema, + profile, item.context or context)) return summaries +_DATASET_ORDER = {'snapshot': 0, 'training_table': 1, 'source_table': 2} + + def _connector(location: str) -> Optional[str]: from urllib.parse import urlparse @@ -365,8 +374,7 @@ def as_artifact(self, key: str) -> Optional[Artifact]: kind = _kinds().get(self.kind) return None if kind is None else Artifact( self.kind, str(self.sha256), self.location, kind.IS_DATASET, - f'kumo.{self.asset_kind}_uri', version=self.sha256, - version_kind='exact') + version=self.sha256, version_kind='exact') @dataclass(frozen=True, kw_only=True) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 5bfe1d9b0..e9f766810 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -40,7 +40,6 @@ build_generated_training_task_dataset, resolve_query_forecast, ) -from kumo_core.summaries import Summary from kumo_core.train.artifact_manifest import PredictionSelectionRef from kumo_core.training.config import FitMode from kumo_core.training.explain_aggregation import aggregate_subgraph_features @@ -64,6 +63,7 @@ snapshot_summaries, split_summaries, table_details, + training_table_details, ) from kumo.engine._deps import Dependencies from kumo.engine._embeddings import ( @@ -418,18 +418,6 @@ def _fit_identity(snapshot: GraphSnapshot, return {} -def _fit_summaries(snapshot: GraphSnapshot, train_table: TrainingTable, - result: TrainingResult, - settings: Any) -> Tuple[Summary, ...]: - index = Path(snapshot.index_path) - return ( - *snapshot_summaries(index, snapshot.graph_config), - *([] if result.manifest is None else artifact_summaries( - result.manifest, 'training', settings, table_details(index))), - *split_summaries(train_table.staging_dir), - ) - - def _store_model(bundle_dir: Path, manifest: Optional[ModelManifest], settings: Any) -> None: r"""Copies the bundle under the root unless it is there.""" @@ -768,9 +756,15 @@ def fit( self._training_result = result session.finish( lambda: fitted_spec( - cast(RunSpec, session.spec), result, - _fit_summaries(snapshot, train_table, result, settings)), - error) + cast(RunSpec, session.spec), result, ( + *snapshot_summaries(Path(snapshot.index_path), + snapshot.graph_config), + *([] if result.manifest is None else artifact_summaries( + result.manifest, 'training', + {**table_details(Path(snapshot.index_path), + snapshot.id), + **training_table_details(train_table)})), + *split_summaries(train_table.staging_dir))), error) raise_on_error(settings, session.error or error, 'Trainer.fit (its model is in trainer.result)') return result diff --git a/packages/kumo-engine/test/test_artifacts.py b/packages/kumo-engine/test/test_artifacts.py index 1fc12bed2..ed5fe584d 100644 --- a/packages/kumo-engine/test/test_artifacts.py +++ b/packages/kumo-engine/test/test_artifacts.py @@ -6,7 +6,7 @@ from typing import List, Optional from kumo_core.artifacts import bundle_id_from_manifest -from kumo_core.summaries import Dataset, Tag +from kumo_core.summaries import Dataset, Table, Tag import kumo.engine as kumo from kumo.engine._deps import Dependencies, FileAsset @@ -61,15 +61,15 @@ def test_a_new_kind_names_its_artifacts_without_other_changes(): assert artifacts(report) == [ Artifact('report', 'r', 's3://root/reports/r', True, - 'kumo.report_uri', version='r', version_kind='exact'), + version='r', version_kind='exact'), Artifact('snapshot', 'g@d', 's3://root/snapshots/g/d', True, - 'kumo.snapshot_uri', version='g@d', version_kind='exact'), + version='g@d', version_kind='exact'), Artifact('source_table', 'orders', 's3://data/orders', True, connector='S3', version='7', version_kind='exact'), Artifact('model', 'm', 's3://root/models/m', False, 'kumo.bundle_uri', version='m', version_kind='exact'), - Artifact('snapshot', 'p', None, True, 'kumo.snapshot_uri', - version='p', version_kind='exact'), + Artifact('snapshot', 'p', None, True, version='p', + version_kind='exact'), ] @@ -108,17 +108,16 @@ def test_a_new_kind_and_asset_reach_tags_and_datasets(tmp_path): assert dependency.location == f'{root}/assets/dummy/{dependency.sha256}' assert ReportManifest.from_dict(report.to_dict()) == report - assert artifact_summaries(report, 'training', settings, { + assert artifact_summaries(report, 'training', { 'orders': ({'ID': {'stype': 'ID', 'dtype': 'int64'}}, {'rows': 3}) }) == [ - Tag('kumo.report_uri', f'{root}/reports/r'), - Tag('kumo.snapshot_uri', f'{root}/snapshots/g/d'), Tag('kumo.bundle_uri', f'{root}/models/m'), - Tag('kumo.dummy_uri', dependency.location), - Tag('kumo.artifact_root', root), - Dataset('r', {'connector': 'LOCAL_FILE', - 'location': f'{root}/reports/r', 'version': 'r', - 'version_kind': 'exact'}, {}, {}, 'training'), + Table('locations', ('kind', 'id', 'location'), ( + ('report', 'r', f'{root}/reports/r'), + ('snapshot', 'g@d', f'{root}/snapshots/g/d'), + ('source_table', 'orders', 's3://data/orders'), + ('model', 'm', f'{root}/models/m'), + ('asset', dependency.sha256, dependency.location)), brief=True), Dataset('g@d', {'connector': 'LOCAL_FILE', 'location': f'{root}/snapshots/g/d', 'version': 'g@d', 'version_kind': 'exact'}, {}, {}, @@ -127,4 +126,7 @@ def test_a_new_kind_and_asset_reach_tags_and_datasets(tmp_path): 'version': '7', 'version_kind': 'exact'}, {'ID': {'stype': 'ID', 'dtype': 'int64'}}, {'rows': 3}, 'training'), + Dataset('r', {'connector': 'LOCAL_FILE', + 'location': f'{root}/reports/r', 'version': 'r', + 'version_kind': 'exact'}, {}, {}, 'training'), ] diff --git a/packages/kumo-engine/test/test_data_summaries.py b/packages/kumo-engine/test/test_data_summaries.py index ce68db8f5..a69d52787 100644 --- a/packages/kumo-engine/test/test_data_summaries.py +++ b/packages/kumo-engine/test/test_data_summaries.py @@ -7,7 +7,7 @@ import numpy as np import pandas as pd import pytest -from kumo_core.summaries import Dataset, Scalar, Table, Tag +from kumo_core.summaries import Dataset, Note, Table, Tag import kumo.engine as kumo from kumo.engine._data_summaries import ( @@ -59,34 +59,42 @@ def test_summaries_match_the_snapshot_and_training_table(scratch): graph, 'PREDICT COUNT(orders.*, 0, 7, days) > 0 FOR EACH users.USER_ID' ).generate_training_table(snapshot) - splits = {summary.name: summary.value - for summary in split_summaries(table.staging_dir)} - for split in ('train', 'val', 'test'): + splits, = split_summaries(table.staging_dir) + assert splits.brief and splits.columns == ( + 'split', 'rows', 'distinct users (approx.)', 'positive rate') + for split, row in zip(('train', 'val', 'test'), splits.rows): seeds = pd.read_parquet(table.staging_dir / f'{split}_seeds') - assert splits[f'data/{split}/rows'] == len(seeds) - assert splits[f'data/{split}/positive_rate'] == pytest.approx( - seeds['TARGET1'].mean()) + assert row[:2] == (split, len(seeds)) + assert row[3] == pytest.approx(seeds['TARGET1'].mean()) - summaries = snapshot_summaries(Path(snapshot.index_path), - snapshot.graph_config) + summaries = {s.name: s for s in snapshot_summaries( + Path(snapshot.index_path), snapshot.graph_config)} + assert set(summaries) == {'data/health', 'data/columns'} matched = orders['USER_ID'] < 300 - relation = 'graph/edges/orders_USER_ID_to_users_USER_ID' - scalars = {s.name: s.value for s in summaries if isinstance(s, Scalar)} - assert scalars == { - 'graph/rows/users': 300, 'graph/rows/orders': 6000, - f'{relation}/count': matched.sum(), - f'{relation}/src_match_frac': pytest.approx(matched.mean()), - f'{relation}/dst_match_frac': pytest.approx( - orders.loc[matched, 'USER_ID'].nunique() / 300), - f'{relation}/fk_lookup_failures': (~matched).sum(), - **{f'graph/{table}/{name}': 0 for table in ('users', 'orders') - for name in ('duplicate_pks', 'dropped_rows', 'null_time_dropped')}, - 'graph/tables': 2, 'graph/edge_types': 1} - assert not [s for s in summaries if isinstance(s, Tag)] - stats, = (s for s in summaries if isinstance(s, Table)) + health = summaries['data/health'] + assert [dict(zip(health.columns, row)) for row in health.rows] == [ + {'table or link': 'users', 'rows': 300, 'columns': 2, + 'dropped rows': 0, 'duplicate keys': 0, 'rows without time': '', + 'edges': '', 'source matched': '', + 'target matched': '', 'failed lookups': ''}, + {'table or link': 'orders', 'rows': 6000, 'columns': 3, + 'dropped rows': 0, 'duplicate keys': 0, 'rows without time': 0, + 'edges': '', 'source matched': '', + 'target matched': '', 'failed lookups': ''}, + {'table or link': 'orders.USER_ID → users', 'rows': '', + 'columns': '', 'dropped rows': '', 'duplicate keys': '', + 'rows without time': '', 'edges': matched.sum(), + 'source matched': pytest.approx(matched.mean()), + 'target matched': pytest.approx( + orders.loc[matched, 'USER_ID'].nunique() / 300), + 'failed lookups': (~matched).sum()}] + stats = summaries['data/columns'] + assert stats.columns[3:] == ( + 'count', 'null rate', 'distinct≈', 'min', 'max', 'mean', 'std', + 'q25≈', 'q50≈', 'q75≈', 'avg tokens≈', 'max tokens≈', 'top values≈') age = next(dict(zip(stats.columns, row)) for row in stats.rows if row[:2] == ('users', 'AGE')) - assert (age['stype'], age['count'], age['null_rate'], age['min'], + assert (age['stype'], age['count'], age['null rate'], age['min'], age['max']) == ('numerical', 300, 1 / 300, users['AGE'].min(), users['AGE'].max()) assert age['mean'] == pytest.approx(users['AGE'].mean()) @@ -138,13 +146,34 @@ def write(name, value): 'dst_key': 'USER_ID'}]}, ]} - tag, = (s for s in snapshot_summaries(tmp_path, config) - if isinstance(s, Tag)) - assert tag == Tag('kumo.data_warnings', '; '.join([ - 'users: 1 missing or invalid primary keys', - 'users: 2 duplicate primary keys', - 'users: 9 rows', - 'orders: 5.0% missing or invalid time', - 'orders.AMOUNT: at least 10% invalid values', - 'orders.USER_ID -> users: 10.0% of orders rows match', - ])) + summaries = snapshot_summaries(tmp_path, config) + assert [s for s in summaries if isinstance(s, (Note, Tag))] == [ + Note('problems', '`users`: 1 missing or invalid primary keys'), + Note('problems', '`users`: 2 duplicate primary keys'), + Note('problems', '`users`: only 9 rows'), + Note('problems', + '`orders`: 5.0% of rows have a missing or invalid time'), + Note('problems', '`orders.AMOUNT`: at least 10% invalid values'), + Note('problems', + 'link `orders.USER_ID → users`: 10.0% of `orders` rows match'), + Tag('kumo.data_warnings', '6'), + ] + + +def test_approximate_counts_never_exceed_the_exact_ones(tmp_path): + (tmp_path / 'split_stats.json').write_text(json.dumps({ + 'val_table_stats': {'num_train_rows': 947, 'num_entities': 944}, + 'test_table_stats': {'num_train_rows': 976, 'num_entities': 985}})) + (tmp_path / 'col_stats.json').write_text(json.dumps({'users': { + 'BIO': {'semantic_type': 'text', 'count': 2000, 'num_unique': 2007}, + 'AGE': {'semantic_type': 'numerical', 'count': 2000, + 'num_unique': 61}}})) + + assert split_summaries(tmp_path) == [Table( + 'data/splits', ('split', 'rows', 'distinct entities (approx.)'), + (('val', 947, 944), ('test', 976, 976)), brief=True)] + columns = {s.name: s for s in snapshot_summaries(tmp_path, {})}[ + 'data/columns'] + distinct = columns.columns.index('distinct≈') + assert [(row[1], row[distinct]) for row in columns.rows] == [ + ('BIO', 2000), ('AGE', 61)] diff --git a/packages/kumo-engine/test/test_mlflow_tracking.py b/packages/kumo-engine/test/test_mlflow_tracking.py index 124e6b61a..9da9bb23a 100644 --- a/packages/kumo-engine/test/test_mlflow_tracking.py +++ b/packages/kumo-engine/test/test_mlflow_tracking.py @@ -127,7 +127,7 @@ def test_fit_logs_a_run_per_fit_and_per_trial(table_and_snapshot, tmp_path, 'kumo.kind', 'kumo.bundle_id', 'kumo.dropped_params', 'kumo.snapshot_id', 'kumo.task_type', 'kumo.entity', 'kumo.tune_metric', 'kumo.engine_version', 'kumo.best_trial', - 'kumo.bundle_uri', 'kumo.artifact_root'} + 'kumo.bundle_uri'} assert parent.data.tags['kumo.bundle_uri'] == str( tmp_path / 'root' / 'models' / result.id) assert parent.data.tags['kumo.snapshot_id'] == snapshot.id @@ -156,7 +156,9 @@ def test_fit_logs_a_run_per_fit_and_per_trial(table_and_snapshot, tmp_path, result.metrics).items() if not name.startswith(('train/', 'val/'))} assert {name for name in metrics if name.startswith(('test/', 'before_refit/'))} == set(tested) - assert metrics['graph/rows/ents'] == 400 + health = json.loads(mlflow.artifacts.load_text( + f'{parent.info.artifact_uri}/data/health.json')) + assert health['data'][0][:2] == ['ents', 400] assert {name: metrics[name] for name in tested} == tested assert 'off for the rest of this run' not in caplog.text ents, = parent.inputs.dataset_inputs From fe1d4713a338cf1506d078e412e2e4f81a1d62f6 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 03:43:01 +0000 Subject: [PATCH 215/278] Publish through TrainingResult.export in the filesystem-only sink test Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/test/test_filestore_sink.py | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/packages/kumo-engine/test/test_filestore_sink.py b/packages/kumo-engine/test/test_filestore_sink.py index d7a9d873e..0554203b8 100644 --- a/packages/kumo-engine/test/test_filestore_sink.py +++ b/packages/kumo-engine/test/test_filestore_sink.py @@ -127,16 +127,12 @@ def test_a_trainer_publishes_and_loads_through_a_filesystem_only_sink( trainer, *_ = fitted connector = SimpleNamespace(source_type='SNOWFLAKE', sink_backend=StageSink()) - result = trainer._finish_fit( - trainer._training_result, predictive_query=trainer._training_query, - snapshot_id=trainer._snapshot_id, - training_table_generation=trainer.training_table_generation_plan, - export_config=kumo.ModelOutputConfig(connector, - path='@stage/models/m')) - assert result.published_model.uri == '@stage/models/m' + published = trainer.result.export( + kumo.ModelOutputConfig(connector, path='@stage/models/m')) + assert published.uri == '@stage/models/m' loaded = kumo.Trainer.load_from_connector(connector, path='@stage/models/m') - assert loaded._snapshot_id == trainer._snapshot_id + assert loaded.result.id == trainer.result.id def test_predictions_export_through_a_filesystem_only_sink(fitted, connector): From 8c43cd8cd76338545c104657095825fb379c64b3 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 03:43:55 +0000 Subject: [PATCH 216/278] Share the serving query and snapshot lookup in integrations/_serving Co-Authored-By: Claude Opus 5.5 --- .../kumo-engine/kumo/engine/_predictor.py | 8 +++---- .../kumo/engine/integrations/_mlflow_model.py | 9 +------- .../kumo/engine/integrations/_serving.py | 23 +++++++++++++++++++ 3 files changed, 27 insertions(+), 13 deletions(-) create mode 100644 packages/kumo-engine/kumo/engine/integrations/_serving.py diff --git a/packages/kumo-engine/kumo/engine/_predictor.py b/packages/kumo-engine/kumo/engine/_predictor.py index c237363f8..6d3be32d4 100644 --- a/packages/kumo-engine/kumo/engine/_predictor.py +++ b/packages/kumo-engine/kumo/engine/_predictor.py @@ -126,12 +126,10 @@ def _in_input_order(rows: pd.DataFrame, scored: pd.DataFrame) -> pd.DataFrame: def _temporal(trainer: Any, snapshot: Any) -> bool: from kumo_core.pquery import query_is_temporal - from kumo.engine.pquery import PredictiveQuery + from kumo.engine.integrations._serving import parsed_query - return query_is_temporal( - PredictiveQuery._from_snapshot(snapshot, - trainer._training_query)._parsed, - snapshot._native_graph) + return query_is_temporal(parsed_query(trainer, snapshot), + snapshot._native_graph) def _seconds(values: pd.Series) -> pd.Series: diff --git a/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py b/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py index c80d8958d..b5589cfda 100644 --- a/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py +++ b/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py @@ -6,8 +6,8 @@ import pandas as pd from kumo.engine.connector._extras import import_provider -from kumo.engine.connector.base import Connector from kumo.engine.graph import GraphSnapshot +from kumo.engine.integrations._serving import snapshot_location _pyfunc = import_provider('mlflow.pyfunc', 'mlflow') _schema = import_provider('mlflow.types.schema', 'mlflow') @@ -53,13 +53,6 @@ def _snapshot(self, uri: str) -> GraphSnapshot: return self._snapshots[uri] -def snapshot_location(uri: str) -> Tuple[Connector, str]: - from kumo.engine.connector import connector_for_root - - parent, slash, name = uri.rstrip('/').rpartition('/') - return connector_for_root(parent or slash or '.'), name - - def signature(snapshot: GraphSnapshot, query: str) -> Tuple[Any, pd.DataFrame]: from kumo_core.train.formatting import prediction_output_columns from kumo_core.types import TaskType diff --git a/packages/kumo-engine/kumo/engine/integrations/_serving.py b/packages/kumo-engine/kumo/engine/integrations/_serving.py new file mode 100644 index 000000000..b457ce35f --- /dev/null +++ b/packages/kumo-engine/kumo/engine/integrations/_serving.py @@ -0,0 +1,23 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""What MLflow and Snowflake serving share: the query and snapshot.""" + +from typing import Any, Tuple + +from kumo.engine.connector.base import Connector +from kumo.engine.graph import GraphSnapshot + + +def parsed_query(trainer: Any, snapshot: GraphSnapshot) -> Any: + from kumo.engine.pquery import PredictiveQuery + + return PredictiveQuery._from_snapshot(snapshot, + trainer._training_query)._parsed + + +def snapshot_location(uri: str) -> Tuple[Connector, str]: + from kumo.engine.connector import connector_for_root + + parent, slash, name = uri.rstrip('/').rpartition('/') + return connector_for_root(parent or slash or '.'), name From 61a68df9be227fae8149b94df939006f2d9d6814 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 03:47:53 +0000 Subject: [PATCH 217/278] Read stored assets and snapshots with the identity the loading call was given A model whose assets sit on a volume, stage or bucket now loads outside a kumo.mlops block through load_model's target, load_from_connector's connector or the pyfunc's model_config['assets_from']; when none can read them, the error names the argument to pass. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_deps.py | 107 +++++++++++++++++- .../kumo/engine/integrations/_mlflow_model.py | 41 ++++++- .../engine/integrations/_mlflow_target.py | 17 ++- .../kumo/engine/integrations/mlflow.py | 16 ++- .../kumo/engine/trainer/trainer.py | 8 +- .../kumo-engine/test/test_mlflow_target.py | 2 + .../test/test_model_agnostic_mlops.py | 66 +++++++++++ 7 files changed, 239 insertions(+), 18 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_deps.py b/packages/kumo-engine/kumo/engine/_deps.py index 5cc78627f..7855af7c6 100644 --- a/packages/kumo-engine/kumo/engine/_deps.py +++ b/packages/kumo-engine/kumo/engine/_deps.py @@ -8,9 +8,13 @@ import os import shutil import tempfile +from contextlib import contextmanager +from contextvars import ContextVar from pathlib import Path from typing import ( Any, + Callable, + Iterator, List, Mapping, Optional, @@ -31,6 +35,29 @@ stamp, ) +Location = Union[str, Tuple[Any, str]] +Hint = Callable[[str, str], str] +_READERS = (('@', 'SnowflakeConnector'), ('/Volumes/', 'DatabricksConnector'), + ('s3://', 'S3Connector'), ('gs://', 'GCSConnector')) + +_LOCATIONS: ContextVar[Tuple[Location, ...]] = ContextVar( + 'kumo_asset_locations', default=()) +_HINT: ContextVar[Optional[Hint]] = ContextVar('kumo_asset_hint', + default=None) + + +@contextmanager +def asset_locations(locations: Sequence[Location], + hint: Optional[Hint] = None) -> Iterator[None]: + r"""Asset roots to look in first, and how to hint at a missing reader.""" + token = _LOCATIONS.set((*locations, *_LOCATIONS.get())) + hinted = _HINT.set(hint or _HINT.get()) + try: + yield + finally: + _HINT.reset(hinted) + _LOCATIONS.reset(token) + class Asset(Protocol): r"""A content-addressed file or directory a model predicts with.""" @@ -115,6 +142,8 @@ def __init__(self, records: Sequence[Dependency] = (), self.records: List[Dependency] = list(records) self.errors: List[BaseException] = [] self._settings = settings + self._locations = list(_LOCATIONS.get()) + self._hint = _HINT.get() @classmethod def of(cls, manifest: Optional[ModelManifest], @@ -162,6 +191,13 @@ def resolve(self, asset_kind: str, name: str) -> Path: continue if directory is not None and _verified(directory, record): return directory + location = record.location or '' + reader = next((name for prefix, name in _READERS + if location.startswith(prefix)), None) + if reader is not None and self._hint is not None: + raise MissingDependencyError( + f"Cannot read {asset_kind} {name!r} (sha256 {record.sha256}) " + f"at {location}: {self._hint(reader, location)}.") raise MissingDependencyError( f"The model needs {asset_kind} {name!r} (sha256 " f"{record.sha256}), which is in none of " @@ -171,8 +207,15 @@ def resolve(self, asset_kind: str, name: str) -> Path: def candidates(self, record: Dependency) -> List[Tuple[Any, str]]: from kumo.engine.connector import connector_for_root + relative = f'{record.asset_kind}/{record.sha256}' found: List[Tuple[Any, str]] = [] - if record.location: + for location in self._locations: + if isinstance(location, str): + found.append((connector_for_root(location), relative)) + else: + connector, path = location + found.append((connector, f"{path.rstrip('/')}/{relative}")) + if record.location and not record.location.startswith('@'): parent, _, leaf = record.location.rstrip('/').rpartition('/') try: found.append((connector_for_root(parent or '/'), leaf)) @@ -180,11 +223,66 @@ def candidates(self, record: Dependency) -> List[Tuple[Any, str]]: pass settings = self._settings if settings is not None: - found.append((settings.root.connector, settings.root.location( - f'assets/{record.asset_kind}/{record.sha256}'))) + found.append((settings.root.connector, + settings.root.location(f'assets/{relative}'))) + if record.location and reads(settings.root.connector, + record.location): + found.append((settings.root.connector, record.location)) return found +def reads(connector: Any, location: str) -> bool: + r"""Whether ``connector`` reads the stored ``location`` itself.""" + from kumo.engine.connector import ( + FileConnector, + GCSConnector, + S3Connector, + is_file_destination, + ) + + for kind, prefix in ((S3Connector, 's3://'), (GCSConnector, 'gs://'), + (FileConnector, None)): + if isinstance(connector, kind): + return prefix is not None and location.startswith(prefix) + return (location.startswith(('@', '/Volumes/')) + and is_file_destination(connector, location)) + + +def stored_roots(manifest: Optional[ModelManifest], + accepts: Callable[[str], bool]) -> List[str]: + r"""The asset roots of ``manifest``'s stored locations ``accepts``.""" + return sorted({ + record.location.rsplit('/', 2)[0] + for record in Dependencies.of(manifest).records + if record.location and accepts(record.location) + }) + + +def connector_roots(manifest: Optional[ModelManifest], + connector: Any) -> List[Tuple[Any, str]]: + r"""The asset roots of ``manifest`` that ``connector`` reads.""" + if connector is None: + return [] + return [(connector, root) for root in stored_roots( + manifest, lambda location: reads(connector, location))] + + +@contextmanager +def assets_through(connector: Any, bundle: Union[str, Path], + hint: Hint) -> Iterator[None]: + r"""Resolves ``bundle``'s stored assets through ``connector`` too.""" + from kumo.engine.trainer._provenance import read_provenance + + with asset_locations( + connector_roots(read_provenance(Path(bundle)), connector), hint): + yield + + +def asset_cache() -> Path: + base = os.environ.get('XDG_CACHE_HOME') or Path.home() / '.cache' + return Path(base) / 'kumo' / 'assets' + + def _store(asset: Asset, sha256: str, settings: Any) -> Optional[str]: from kumo_core.object_store_sink import _prepared @@ -237,8 +335,7 @@ def _fetch(connector: Any, path: str, sha256: str) -> Optional[Path]: return None if isinstance(filesystem, LocalFileSystem): return Path(root) - cache = os.environ.get('XDG_CACHE_HOME') or Path.home() / '.cache' - target = Path(cache) / 'kumo' / 'assets' / sha256 + target = asset_cache() / sha256 if (target / PROVENANCE_FILE).is_file(): return target target.parent.mkdir(parents=True, exist_ok=True) diff --git a/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py b/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py index b5589cfda..40a9e5e1c 100644 --- a/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py +++ b/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py @@ -1,6 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. # All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import os from typing import Any, Dict, Optional, Tuple import pandas as pd @@ -24,9 +25,15 @@ class KumoModel(_pyfunc.PythonModel): def load_context(self, context: Any) -> None: + from kumo.engine._deps import assets_through from kumo.engine._predictor import load_predictor - self._predictor = load_predictor(context.artifacts['bundle']) + self._reader = context.model_config.get('assets_from') + bundle = context.artifacts['bundle'] + with assets_through( + self._reader, bundle, lambda reader, _: + f'{_assets_from(reader)}, with an identity that reads it'): + self._predictor = load_predictor(bundle) self._columns = list(context.model_config['columns']) self._snapshots: Dict[str, GraphSnapshot] = {} self._default = context.model_config.get('snapshot') @@ -46,13 +53,39 @@ def predict(self, context: Any, model_input: pd.DataFrame, params) def _snapshot(self, uri: str) -> GraphSnapshot: + from kumo.engine._deps import reads + if uri not in self._snapshots: - connector, path = snapshot_location(uri) - self._snapshots[uri] = GraphSnapshot.load_from_connector( - connector, path=path) + connector, path = ((self._reader, uri) + if self._reader is not None + and reads(self._reader, uri) else + snapshot_location(uri)) + try: + self._snapshots[uri] = GraphSnapshot.load_from_connector( + connector, path=path) + except FileNotFoundError as error: + reader = _remote_reader(uri) + if reader is None or connector is self._reader: + raise + raise FileNotFoundError( + f'{error} To read {uri} from here, ' + f'{_assets_from(reader)}.') from error return self._snapshots[uri] +def _assets_from(reader: str) -> str: + return (f"pass model_config={{'assets_from': {reader}(...)}} to " + f"mlflow.pyfunc.load_model(model_uri, ...)") + + +def _remote_reader(uri: str) -> Optional[str]: + if uri.startswith('@'): + return 'SnowflakeConnector' + if uri.startswith('/Volumes/') and not os.path.isdir('/Volumes'): + return 'DatabricksConnector' + return None + + def signature(snapshot: GraphSnapshot, query: str) -> Tuple[Any, pd.DataFrame]: from kumo_core.train.formatting import prediction_output_columns from kumo_core.types import TaskType diff --git a/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py b/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py index 82e5d9ee1..47b25047c 100644 --- a/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py +++ b/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py @@ -66,8 +66,10 @@ class Target: def __init__(self, tracking_uri: str, registry_uri: Optional[str], swap: Optional[Callable[[], ContextManager[Any]]] = None, - aws_session: Any = None, host: Optional[str] = None) -> None: + aws_session: Any = None, host: Optional[str] = None, + assets: Any = None) -> None: self.tracking_uri = tracking_uri + self.assets = assets self.registry_uri = registry_uri self._swap = swap self._aws_session = aws_session @@ -177,12 +179,13 @@ def _databricks(connector: Any) -> Target: if (principal is None or principal.config.authentication is DatabricksAuthenticationKind.AMBIENT): return Target('databricks', 'databricks-uc', - swap=lambda: global_tracking_uri('databricks')) + swap=lambda: global_tracking_uri('databricks'), + assets=connector) client = import_provider('kumo_databricks.session', 'databricks')._client_for(principal) return Target('databricks', 'databricks-uc', swap=lambda: _databricks_token(client), - host=client.config.host) + host=client.config.host, assets=connector) @contextmanager @@ -206,15 +209,19 @@ def _databricks_token(client: Any) -> Iterator[None]: def _sagemaker(target: SageMakerMLflow) -> Target: from kumoapi.data_source import S3AmbientAuthentication + from kumo.engine.connector import S3Connector + arn = target.tracking_server_arn + assets = S3Connector(identity=target.identity) principal = _principal(target.identity) if principal is None or isinstance(principal.config.authentication, S3AmbientAuthentication): - return Target(arn, arn, swap=lambda: global_tracking_uri(arn)) + return Target(arn, arn, swap=lambda: global_tracking_uri(arn), + assets=assets) session = import_provider('kumo_s3.backend', 's3')._session_for(principal) return Target(arn, arn, swap=lambda: _sagemaker_session(arn, session), - aws_session=session) + aws_session=session, assets=assets) @contextmanager diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index 8ab2e4d34..6fc8be3e6 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -788,6 +788,7 @@ def load_model( Args: target: as in `MlflowTracking`. """ + from kumo.engine._deps import assets_through from kumo.engine._scratch import new_run_id, trainer_scratch_dir from kumo.engine.trainer.trainer import Trainer, warn_if_other_bundle @@ -799,12 +800,22 @@ def load_model( local = Path(_download(mlflow, target, model_uri, destination)) model = mlflow.models.Model.load(str(local)) bundle = model.flavors['python_function']['artifacts']['bundle']['path'] - trainer = Trainer.load(local / bundle, embedding_models, - accelerator=accelerator) + with assets_through(target.assets, local / bundle, _target_hint): + trainer = Trainer.load(local / bundle, embedding_models, + accelerator=accelerator) warn_if_other_bundle(model.metadata, trainer) return trainer +def _target_hint(reader: str, location: str) -> str: + target = { + 'DatabricksConnector': 'DatabricksConnector(identity=...)', + 'S3Connector': 'SageMakerMLflow(arn, identity=S3Identity(...))', + }.get(reader, f'a {reader}') + return (f'pass target={target} to load_model(model_uri, ...), with an ' + f'identity that reads {location.rsplit("/", 2)[0]}') + + def _download(mlflow: Any, target: Any, model_uri: str, destination: Path) -> str: if target.explicit and model_uri.startswith('models:/'): @@ -879,6 +890,7 @@ def _log_model(target: Any, client: Any, bundle_dir: Path, 'columns': [] if signature is None else signature.outputs.input_names(), 'snapshot': serving_snapshot, + 'assets_from': None, }, artifacts={ 'bundle': str(bundle_dir), diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index e9f766810..6eebb8916 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -1420,8 +1420,12 @@ def load_from_connector( published = sink.fetch_model(cast(SourceConnector, model_registry), str(bundle_dir), path=address.path, location=address.location) - trainer = cls._load_bundle(bundle_dir, embedding_models, - accelerator=accelerator) + with dependency.assets_through( + model_registry, bundle_dir, lambda reader, _: ( + f'pass a {reader}(...) with an identity that reads it ' + f'as Trainer.load_from_connector(model_registry, ...)')): + trainer = cls._load_bundle(bundle_dir, embedding_models, + accelerator=accelerator) warn_if_other_bundle(published.metadata, trainer) return trainer diff --git a/packages/kumo-engine/test/test_mlflow_target.py b/packages/kumo-engine/test/test_mlflow_target.py index d461a61a0..cb5071b4d 100644 --- a/packages/kumo-engine/test/test_mlflow_target.py +++ b/packages/kumo-engine/test/test_mlflow_target.py @@ -66,8 +66,10 @@ def test_ambient_identities_set_no_credentials(monkeypatch, tracking_uri): assert mlflow.get_tracking_uri() == tracking_uri assert dict(os.environ) == before + assert resolve(_connector(DatabricksIdentity())).assets.catalog == 'main' for target in (resolve(SageMakerMLflow(ARN, S3Identity())), resolve(SageMakerMLflow(ARN))): + assert type(target.assets).__name__ == 'S3Connector' with target.scope(): assert mlflow.get_tracking_uri() == ARN assert target.artifact_repository('s3://b/run') is None diff --git a/packages/kumo-engine/test/test_model_agnostic_mlops.py b/packages/kumo-engine/test/test_model_agnostic_mlops.py index 9b3563b29..c6b3de037 100644 --- a/packages/kumo-engine/test/test_model_agnostic_mlops.py +++ b/packages/kumo-engine/test/test_model_agnostic_mlops.py @@ -366,3 +366,69 @@ def refuse(*args): registered, error = _store_and_register(bundle, manifest, settings, warn=False, stored=asset_failure) assert registered is None and error is asset_failure + + +def test_volume_assets_load_through_the_connector_the_loader_was_given( + tmp_path, monkeypatch): + from pyarrow.fs import LocalFileSystem, SubTreeFileSystem + + from kumo.engine.integrations import _mlflow_model + from kumo.engine.integrations.mlflow import _target_hint + + monkeypatch.setenv('XDG_CACHE_HOME', str(tmp_path / 'cache')) + (tmp_path / 'volume').mkdir() + files = SubTreeFileSystem(str(tmp_path / 'volume'), LocalFileSystem()) + volume = SimpleNamespace(sink_backend=SimpleNamespace( + holds_files=lambda connector, path: path.startswith('/Volumes/'), + filesystem=lambda connector, path: (files, path.lstrip('/')))) + (tmp_path / 'vocab').mkdir() + (tmp_path / 'vocab' / 'words.txt').write_text('a b c') + recorded = Dependencies(settings=kumo.MLOps((volume, '/Volumes/c/s/v/a'))) + record = recorded.require(FileAsset('toy', 'vocab', tmp_path / 'vocab')) + assert record.location.startswith('/Volumes/c/s/v/a/assets/toy/') + manifest = ModelManifest(id='m', dependencies=[record]) + + with _deps.asset_locations([], _target_hint): + with pytest.raises(MissingDependencyError, + match=r'pass target=DatabricksConnector\(' + r'identity=\.\.\.\) to load_model\(model_uri, ' + r'\.\.\.\), with an identity that reads ' + r'/Volumes/c/s/v/a/assets\.$'): + Dependencies.of(manifest).resolve('toy', 'vocab') + with _deps.asset_locations(_deps.connector_roots(manifest, volume)): + assert (Dependencies.of(manifest).resolve('toy', 'vocab') / + 'words.txt').read_text() == 'a b c' + other_root = kumo.MLOps((volume, '/Volumes/c/s/v/b')) + assert Dependencies.of(manifest, other_root).resolve('toy', 'vocab') + + seen = [] + monkeypatch.setattr( + 'kumo.engine._predictor.load_predictor', lambda bundle: seen.append( + Dependencies.of(manifest).candidates(record)[0])) + monkeypatch.setattr('kumo.engine.trainer._provenance.read_provenance', + lambda bundle: manifest) + _mlflow_model.KumoModel().load_context(SimpleNamespace( + artifacts={'bundle': str(tmp_path)}, + model_config={'columns': [], 'assets_from': volume})) + assert seen == [(volume, record.location)] + + +def test_the_pyfunc_reads_a_volume_snapshot_through_assets_from( + monkeypatch): + from kumo.engine.graph import GraphSnapshot + from kumo.engine.integrations import _mlflow_model + + volume = SimpleNamespace(sink_backend=SimpleNamespace( + holds_files=lambda connector, path: path.startswith('/Volumes/'), + filesystem=lambda connector, path: None)) + loads = [] + monkeypatch.setattr(GraphSnapshot, 'load_from_connector', + lambda connector, *, path: loads.append( + (connector, path))) + model = _mlflow_model.KumoModel() + model._snapshots, model._reader = {}, volume + model._snapshot('/Volumes/c/s/v/a/snapshots/g/version') + model._reader = None + model._snapshot('/Volumes/c/s/v/b/snapshots/g/version') + assert loads[0] == (volume, '/Volumes/c/s/v/a/snapshots/g/version') + assert loads[1][0] is not volume and loads[1][1] == 'version' From 88b88507f620249225056a7accfdd22bf721f8e5 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 03:48:39 +0000 Subject: [PATCH 218/278] Build a Snowflake registry version's model arguments in one helper Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_snowflake/registry.py | 17 ++++++++++++++--- packages/kumo-snowflake/tests/test_registry.py | 6 ++---- 2 files changed, 16 insertions(+), 7 deletions(-) diff --git a/packages/kumo-snowflake/src/kumo_snowflake/registry.py b/packages/kumo-snowflake/src/kumo_snowflake/registry.py index 42a9d55de..9ac4d98c6 100644 --- a/packages/kumo-snowflake/src/kumo_snowflake/registry.py +++ b/packages/kumo-snowflake/src/kumo_snowflake/registry.py @@ -37,6 +37,7 @@ import math import re import shutil +import tempfile from contextlib import contextmanager from pathlib import Path from typing import ( @@ -186,6 +187,15 @@ def _carrier_signature() -> Dict[str, "model_signature.ModelSignature"]: } +def _model_args(bundle: Path, metadata: Mapping[str, str], + connector: SourceConnector, workdir: str) -> Dict[str, Any]: + r"""``log_model`` arguments for the bundle's carrier model.""" + return { + 'model': _bundle_model(str(bundle)), + 'signatures': _carrier_signature() + } + + def _locate_bundle(root: Path) -> Path: r"""Finds the published bundle inside a downloaded model tree. @@ -387,18 +397,19 @@ def publish_model( f"has no {_SPEC_FILE}, which an export writes into every one.") record = {**_finite(metrics), **_version_record(metadata)} - with self._registry(connector) as registry: + with tempfile.TemporaryDirectory() as workdir, \ + self._registry(connector) as registry: + args = _model_args(bundle, metadata, connector, workdir) name = location.version or _next_version(registry, location.model_name) for attempt in range(3): try: version = registry.log_model( - _bundle_model(str(bundle)), model_name=location.model_name, version_name=name, comment=_comment(location.comment, metrics), metrics=record, - signatures=_carrier_signature(), + **args, **_task(bundle), ) break diff --git a/packages/kumo-snowflake/tests/test_registry.py b/packages/kumo-snowflake/tests/test_registry.py index 4d612f8ad..7d0fe7256 100644 --- a/packages/kumo-snowflake/tests/test_registry.py +++ b/packages/kumo-snowflake/tests/test_registry.py @@ -108,8 +108,7 @@ def registry(_): sink = SnowflakeModelRegistryMixin() monkeypatch.setattr(sink, '_registry', registry) - monkeypatch.setattr(registry_module, '_bundle_model', lambda _: object()) - monkeypatch.setattr(registry_module, '_carrier_signature', lambda: {}) + monkeypatch.setattr(registry_module, '_model_args', lambda *_: {}) metrics = {'auroc': 0.9213, 'holdout_loss': 0.5, 'brier': float('nan')} published = sink.publish_model( SimpleNamespace(), str(bundle), @@ -349,8 +348,7 @@ def registry(_): sink = SnowflakeModelRegistryMixin() monkeypatch.setattr(sink, '_registry', registry) - monkeypatch.setattr(registry_module, '_bundle_model', lambda _: object()) - monkeypatch.setattr(registry_module, '_carrier_signature', lambda: {}) + monkeypatch.setattr(registry_module, '_model_args', lambda *_: {}) sink.publish_model(SimpleNamespace(), str(bundle), location=ModelRegistryLocation('CHURN', version='V1')) assert calls[0].get('task') == expected From 5063dedc14bcbd1d639c0fe0d37643e2558fccec Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 03:49:50 +0000 Subject: [PATCH 219/278] Correct the MLOps guide for run names, registration, SageMaker and asset loading Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/docs/mlops.md | 52 ++++++++++++++++++------------ 1 file changed, 31 insertions(+), 21 deletions(-) diff --git a/packages/kumo-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index 6ce8efefb..7268d2138 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -2,7 +2,7 @@ Turn MLOps on with one `kumo.MLOps` and the normal `kumo.engine` calls store their snapshots, training tables and models under an artifact root, record fits in your experiment tracker, and register models in its registry. -Install with `pip install "kumo-engine[mlflow]"` (add `s3`, `gcs`, `databricks`, `sagemaker` or `snowflake` as needed; Snowflake needs only `kumo-engine[snowflake]`). GPU, CUDA and Databricks Runtime notes are under [Install](#install). +Install the engine wheel by path, `pip install "/[mlflow]"` (add `s3`, `gcs`, `databricks`, `sagemaker` or `snowflake` as needed; Snowflake needs only `[snowflake]`): PyPI's `kumo-engine` 1.0.0 predates this API (see [Install](#install)). GPU, CUDA and Databricks Runtime notes are under [Install](#install). ## Quickstart @@ -99,7 +99,7 @@ To serve only a better model, move its alias yourself: [Promote a model](#promot | `table.id` | a training table | `/training_tables//` | | `result.id` | a model (its bundle id) | `/models//`; the `kumo.bundle_id` run tag; `kumo_bundle_id` on the registered version | | `result.registered_model.version` | the registered version | the registry | -| `result.tracking` | the tracker's own run of the fit: `platform`, `run_id`, `url` | the tracker's UI | +| `result.tracking` | the tracker's run of the fit: `platform`, `experiment`, `run_id`, `run_name`, `url` | the tracker's UI | Find runs with the platform's UI, or by tag: `kumo.kind` (`fit`, `trial`) and `kumo.bundle_id`. @@ -108,13 +108,13 @@ Find runs with the platform's UI, or by tag: `kumo.kind` (`fit`, `trial`) and `k `kumo.MLOps` says where artifacts go and what records them. Use it as a block, or pass it as `mlops=` on a single call (`fit`, `graph.snapshot`, `generate_training_table`): ```python -cfg = kumo.MLOps("s3://bucket/kumo", tracking=..., registry=..., on_error="raise") +cfg = kumo.MLOps("s3:///kumo", tracking=..., registry=..., on_error="raise") with kumo.mlops(cfg): snapshot = graph.snapshot() trainer.fit(table, snapshot=snapshot, mlops=cfg) ``` -- **`artifact_root`** (required): a URI (`s3://`, `gs://`, `file://`, or a local path, including `/Volumes/...` on a Databricks cluster) or a `(connector, path)` pair for a store a URI cannot reach: a volume through a `DatabricksConnector` (`"/Volumes////kumo"`) or a named stage through a `SnowflakeConnector` (`"@DB.SCHEMA.STAGE/kumo"`) once the installed sink holds files there. +- **`artifact_root`** (required): a URI (`s3://`, `gs://`, `file://`, or a local path, including `/Volumes/...` on a Databricks cluster) or a `(connector, path)` pair for a store a URI cannot reach: a volume through a `DatabricksConnector` (`"/Volumes////kumo"`) or a named stage through a `SnowflakeConnector` (`"@DB.SCHEMA.STAGE/kumo"`) once the installed sink holds files there. The connector's `root_dir` must contain the path; `kumo.MLOps` raises otherwise, so a data connector rooted at `s3:///data` cannot hold `s3:///artifacts`. - **`tracking`**: `MlflowTracking` or `SnowflakeTracking` (below), or your own [`Tracking`](#custom-trackers). - **`registry`**: a `ModelRegistryLocation(name, alias=...)`; every fit registers its model there, in the tracker's registry. - **`store`**: the artifact kinds stored when produced: `snapshot`, `training_table` and `model` (default: all). @@ -143,8 +143,8 @@ table = kumo.TrainingTable.load_from_connector(cfg, path=f"training_tables/{tabl - `result.registered_model` is the version `MLOps(registry=)` registered; `result.published_model` the one `fit(export_config=...)` published. - `result.tracking` is the tracker's run: `TrackingRef(platform="databricks", run_id="...", url="https://.../ml/experiments/.../runs/...")`. The model's `_provenance.json` records the same. - `fit(scoring_defaults=kumo.ScoringDefaults(...))` stores prediction defaults in the bundle before it is stored and registered: `binary_classification_threshold` (binary and multilabel classification) and `num_classes_to_return` (multiclass, multilabel, link prediction). `predict` and the MLflow pyfunc apply them when a call omits them; explicit arguments and a query's `TOP k` win. A binary model always returns `True_PROB`; with no threshold passed or stored, its `TARGET_PRED` is null. `result.with_scoring_defaults(...)` retunes a fitted model as a new bundle whose manifest names the original in `parents.derived_from`. -- `result.export(cfg)` (or `trainer.export(cfg)` for a loaded model) registers a model at `cfg.registry` through its tracker. On every platform, registering a bundle the registry already holds under that name returns its version instead of adding one. -- If registration fails, `fit` warns with the full error and returns with `registered_model=None` (with `on_error="raise"`, it raises), and the fit run ends failed with `kumo.error`. Registering a bundle the registry already holds under the name returns that version (logged at info) and moves the alias to it. +- `result.export(cfg)` (or `trainer.export(cfg)` for a loaded model) registers a model at `cfg.registry` through its tracker. On MLflow and Unity Catalog, registering a bundle the registry already holds under that name returns its version instead of adding one (found by its bundle id tag); the Snowflake Model Registry adds a new version each time. +- If registration fails, `fit` warns with the full error and returns with `registered_model=None` (with `on_error="raise"`, it raises), and the fit run ends failed with `kumo.error`. On MLflow and Unity Catalog, a bundle the registry already holds under the name returns that version (logged at info) and the alias moves to it. **Model names** follow the registry's rules, checked when `kumo.MLOps` is built: @@ -158,7 +158,7 @@ An alias is a letter or underscore followed by letters, digits or underscores; ` ### Promote a model -Register every fit under a candidate alias and move the serving alias yourself once the new model is better, comparing `result.metrics` with the champion's: +Register every fit under a candidate alias and move the serving alias yourself once the new model is better, comparing `result.metrics` (the holdout's `test_`, logged to MLflow as `test/`) with the champion's: ```python with cfg.tracking.client() as client: # the tracker's MLflow server, as its identity @@ -171,9 +171,10 @@ with cfg.tracking.client() as client: # the tracker's MLflow server, as In the Snowflake Model Registry a version holds at most one alias: take `champion` off the old version (there is none on the first promotion) and `candidate` off the new one, then set it: ```sql -ALTER MODEL ML.MODELS.CHURN_GNN VERSION champion UNSET ALIAS; -- skip on the first promotion -ALTER MODEL ML.MODELS.CHURN_GNN VERSION V3 UNSET ALIAS; -- V3 holds candidate -ALTER MODEL ML.MODELS.CHURN_GNN VERSION V3 SET ALIAS = champion; +SHOW VERSIONS IN MODEL ML.MODELS.CHURN_GNN; -- or result.registered_model.version +ALTER MODEL ML.MODELS.CHURN_GNN VERSION champion UNSET ALIAS; -- skip on the first promotion +ALTER MODEL ML.MODELS.CHURN_GNN VERSION UNSET ALIAS; -- it holds candidate +ALTER MODEL ML.MODELS.CHURN_GNN VERSION SET ALIAS = champion; ``` Jobs that load `alias="champion"` pick up the new version on their next run. @@ -183,6 +184,7 @@ Jobs that load `alias="champion"` pick up the new version on their next run. `MlflowTracking(experiment_name=None, run_name=None, log_system_metrics=True, tags=None, log_model=False, *, target=None)`: - With no `target`, it records where MLflow's own configuration points (`MLFLOW_TRACKING_URI`, `mlflow.set_tracking_uri`; automatic on Databricks). `target` names another server with its credentials: a `DatabricksConnector`, a `SageMakerMLflow(arn)`, or a tracking URI. With a `target`, `experiment_name` is required to record runs. - `log_model=True` logs the model on the fit run without registering it; `MLOps(registry=)` implies it. +- Building `kumo.MLOps` creates a missing experiment; if MLflow refuses (on Databricks, the experiment's folder must exist), it raises and says so. `SnowflakeTracking(source, *, experiment_name, database=None, schema_name=None, run_name=None, tags=None, native_model_lineage=False)`: - `source` is a `SnowflakeConnector`, whose identity opens the session and whose Snowflake Model Registry `MLOps(registry=)` registers in, or a Snowpark session, which tracks only. @@ -198,14 +200,17 @@ Jobs that load `alias="champion"` pick up the new version on their next run. A model registered through `MlflowTracking` is an MLflow pyfunc: it serves with `mlflow.pyfunc.load_model(uri)`, `spark_udf` and Model Serving. ```python -model = mlflow.pyfunc.load_model("models:/main.ml.users_churn@champion") +with cfg.tracking.client(): # MLflow's own calls reach the tracker's server, as its identity + model = mlflow.pyfunc.load_model("models:/main.ml.users_churn@champion") scores = model.predict(pd.DataFrame({"ENTITY": [1, 2, 3]}), params={ "snapshot": cfg.snapshot_uri(snapshot), # "/Volumes/main/ml/kumo/snapshots//" - "prediction_time": "2024-06-01"}) # temporal queries + "prediction_time": "2024-06-01"}) # temporal queries need it, or a TIMESTAMP column ``` -- Input rows have an `ENTITY` column and optionally a `TIMESTAMP`; scores come back in input order, one per scored row. Binary classification returns `True_PROB`; `TARGET_PRED` needs `params={"binary_classification_threshold": ...}` or a stored threshold, and is null without one. +- Input rows have an `ENTITY` column and optionally a `TIMESTAMP`; a temporal query needs one of `TIMESTAMP` or `prediction_time`. Scores come back in input order, one per scored row. Binary classification returns `True_PROB`; `TARGET_PRED` needs `params={"binary_classification_threshold": ...}` or a stored threshold, and is null without one. - Each replica loads a snapshot on its first `predict` with that URI and keeps it. `log_model(trainer, snapshot=..., serving_snapshot=...)` logs a model whose replicas load that snapshot at startup and use it when no `snapshot` param is given; `serving_snapshot` takes a URI, or a `GraphSnapshot` or snapshot id stored under the block's artifact root. +- Snapshot URIs: `s3://`, `gs://`, `file://` or a local path, read with the ambient credentials. A `/Volumes/...` path off a Databricks cluster, or a stage path, is read through `model_config={"assets_from": connector}` (below), e.g. the connector of a `(connector, "/Volumes/...")` root. +- A model with stored assets (word vectors for text columns) on a volume, stage or bucket reads them, and snapshots there, with the identity it is given: `mlflow.pyfunc.load_model(uri, model_config={"assets_from": connector})`, a connector for that store with its credentials; an asset it cannot find raises, naming this option. - Serving needs an image with kumo-engine installed: `mlflow models serve -m --env-manager local`. Its pip requirement is `kumo-engine[s3,gcs]` of the same version (plus `[text]` for text columns). ## Credentials and running off-platform @@ -228,16 +233,21 @@ snowflake = SnowflakeTracking(SnowflakeConnector(warehouse="", database="ML" ### SageMaker managed MLflow -1. Create a tracking server once (15-25 minutes): `aws sagemaker create-mlflow-tracking-server --tracking-server-name --artifact-store-uri s3:///mlflow --role-arn `; the server role needs read and write on that prefix. A stopped server starts with `aws sagemaker start-mlflow-tracking-server --tracking-server-name ` (about 20 minutes); `kumo.MLOps` fails until `aws sagemaker describe-mlflow-tracking-server --tracking-server-name ` reports `Started`. -2. Give the job's identity `sagemaker-mlflow:*` on the server ARN and S3 read and write on the artifact store and the kumo artifact root. Rights scoped to the root's prefix (`s3:GetObject`, `s3:PutObject` and `s3:DeleteObject` on `//*`, `s3:ListBucket` limited to that prefix) suffice: kumo-engine makes no bucket-level calls. +1. Create a tracking server once (15-25 minutes): `aws sagemaker create-mlflow-tracking-server --tracking-server-name --artifact-store-uri s3:///mlflow --role-arn `; the server role needs read and write on that prefix. A stopped server starts with `aws sagemaker start-mlflow-tracking-server --tracking-server-name ` (about 20 minutes); `kumo.MLOps` fails until `aws sagemaker describe-mlflow-tracking-server --tracking-server-name ` reports `Started` (`Created` for a new server). +2. Give the job's identity `sagemaker-mlflow:*` on the server ARN (an MLflow App instead needs `sagemaker:CallMlflowAppApi` on the app ARN; the sample policy on AWS's MLflow Apps setup page lists the Tracking Server actions, which don't apply) and S3 read and write on the artifact store and the kumo artifact root. Rights scoped to the root's prefix (`s3:GetObject`, `s3:PutObject` and `s3:DeleteObject` on `//*`, `s3:ListBucket` limited to that prefix) suffice: kumo-engine makes no bucket-level calls. 3. Install `kumo-engine[mlflow,sagemaker]` and pass `MlflowTracking(..., target=SageMakerMLflow(arn))`; the registry is the same server. -4. MLflow's run links do not open on SageMaker: open the UI from SageMaker Studio, or at the URL `aws sagemaker create-presigned-mlflow-tracking-server-url --tracking-server-name ` prints. A tracked fit warns with this command, and `result.tracking.url` is `None`. +4. MLflow's run links do not open on SageMaker. A tracked fit warns with the command below, and `result.tracking.url` is `None`. To open the UI: + - **SageMaker AI Studio** (not Studio Classic): Applications » MLflow, then Open MLflow on the server's menu. Studio acts as the domain's execution role, which needs `sagemaker-mlflow:AccessUI` (grant it on `*` if a policy scoped to the server ARN is still denied) and `sagemaker:ListMlflowTrackingServers` to list it). + - **SageMaker Unified Studio**: in an IAM-mode domain (tested), the project's MLflow page only lists servers connected to the project: Connect Tracking Server with its ARN. AWS documents Create MLflow Tracking Server and Create MLflow App on that page for other domains, deleted with the project (untested here). + - **Anywhere**, with your own credentials: `aws sagemaker create-presigned-mlflow-tracking-server-url --tracking-server-name ` prints a one-time URL valid for 5 minutes; it needs `sagemaker:CreatePresignedMlflowTrackingServerUrl` and `sagemaker-mlflow:AccessUI`. + - **MLflow App** instead of a server: `aws sagemaker create-mlflow-app --name --artifact-store-uri s3:///mlflow --role-arn ` (ready at `Created`, about 5 minutes); pass its ARN as the target. The UI: `aws sagemaker create-presigned-mlflow-app-url --arn ` (needs `sagemaker:CreatePresignedMlflowAppUrl`), then pick the workspace `default`. +5. In a SageMaker notebook, use a Python 3.11+ kernel with the wheel installed by path (`pip install --find-links "/[mlflow,sagemaker,s3,text]" ipykernel ipywidgets`, then `python -m ipykernel install --user --name kumo-engine`); `SageMakerMLflow(arn)` uses the notebook's role. Notebook instances need the AL2023 platform (`notebook-al2023-v1`): the wheel needs glibc 2.34. ## Reference ### Install -- A bare `kumo-engine[mlflow]` with `--find-links ` can resolve to an unrelated `kumo-engine` on PyPI; install the engine wheel by path: `pip install --find-links "/[mlflow]"`. +- Install the engine wheel by path, `pip install "/[mlflow]"`. PyPI's `kumo-engine` is the 1.0.0 release, which predates the MLOps API, and a bare `kumo-engine[mlflow]` resolves to it over a dev wheel; with it installed, `kumo.MLOps` raises `AttributeError: module 'kumo.engine' has no attribute 'MLOps'` and `pip show kumo-engine` says 1.0.0. The wheel `scripts/build_wheel.sh` builds bundles kumo-core, kumo-api, kumo-pql, DiskGraph and the connector packages; with separately built wheels, add `--find-links `. - **GPU:** on a plain GPU VM or SageMaker image, install a torch build that matches the NVIDIA driver first, e.g. `pip install --index-url https://download.pytorch.org/whl/cu126 "torch==2.12.*" torchvision`. PyPI's default torch is built for a recent CUDA (`torch.version.cuda`; CUDA 13 needs driver ≥ 580); cu126 runs on ≥ 525.60. With `accelerator="auto"`, a torch that cannot use the driver runs on CPU: check `torch.cuda.is_available()`. - **Databricks:** Runtime 17.3 LTS ML or later keeps its torch, torchvision and pyarrow (kumo-engine accepts pyarrow 16 through 24; a custom CA bundle for S3 needs 21 or later). On 16.4 LTS ML pip replaces torch and breaks torchvision. The parent folder of an experiment path must exist. Registering to Unity Catalog from outside Databricks also needs `boto3`. - **Disk:** `kumo.configure(scratch_dir=...)` holds snapshots, tables and bundles; Python's own temporary files (multiprocessing) follow `TMPDIR`, so point both at a large local disk. @@ -267,7 +277,7 @@ Every artifact carries `_provenance.json`, metadata only (the underscore keeps S ### Artifacts, bundles and dependencies - A model's **bundle** holds only its own state: weights, encoders, the training spec and the holdout. Its manifest's **`parents`** are lineage; its **`dependencies`** are what `predict` needs besides the bundle and a snapshot: the `predictor` entry point, the `kumo-engine` `package` (with `[text]` when a text column is encoded) and content-addressed `asset`s. -- **Assets** are stored once under `/assets///` and shared by every model that uses them. `Trainer.load` and the pyfunc look for each at the location recorded, then under the enclosing block's root, then in the embedding-model store, check its sha256 and cache non-local copies under `~/.cache/kumo/assets/`. A stored asset that cannot be found raises `MissingDependencyError`. +- **Assets** are stored once under `/assets///` and shared by every model that uses them. `Trainer.load` and the pyfunc look for each at the location recorded, then under the enclosing block's root, then in the embedding-model store, check its sha256 and cache non-local copies under `~/.cache/kumo/assets/`. An asset on a stage, a volume or a bucket is read with the identity the loading call was given: `Trainer.load_from_connector(connector, ...)`, `load_model(uri, target=DatabricksConnector(...) or SageMakerMLflow(arn, identity=...))`, the block root's connector, or `mlflow.pyfunc.load_model(uri, model_config={"assets_from": connector})`. A stored asset that cannot be found raises `MissingDependencyError` naming the argument that passes such a connector. - **Text encoders** are the first assets: word-vector models (`GloVeVectors`) are stored; other `TokenEmbeddingModel`s are recorded by fingerprint and must be registered wherever the model loads: ```python @@ -295,11 +305,11 @@ Other systems plug in through `kumo.engine.tracking.Tracking`, which receives re ### Progress events -`fit(callbacks=[...])` calls each callback with every `kumo.engine.tracking.Event`: `fit_started`, per trial `trial_started`, `epoch_ended` and `trial_ended`, then `fit_ended`. Ended events carry `params["status"]` and, on failure, `params["error"]`; a callback that raises is logged and not called again. For JSON lines: `callbacks=[lambda e: print(json.dumps(e.to_dict()))]`. `graph.snapshot(callbacks=[...])` reports `snapshot_started`, `table_staged` per table and `snapshot_ended` (with `snapshot_id` and `reused`), each with `params["seconds"]` since the start, and logs each at INFO. +`fit(callbacks=[...])` calls each callback with every `kumo.engine.tracking.Event`: `fit_started`, per trial `trial_started`, `epoch_ended` and `trial_ended`, then `fit_ended`. Ended events carry `params["status"]` and, on failure, `params["error"]`; a callback that raises is logged and not called again. For JSON lines: `callbacks=[lambda e: print(json.dumps(e.to_dict()))]`. `graph.snapshot(callbacks=[...])` reports `snapshot_started`, `table_staged` per table and `snapshot_ended` (with `snapshot_id` and `reused`), each with `params["seconds"]`, the time since the snapshot started, and logs each at INFO. `callbacks=[print]` prints each as its name and fields, e.g. `table_staged table=users seconds=12.0`. ### Limits - Tracker calls run on a background thread; `fit` waits at most 30 s for tracking to finish, and up to an hour for model logging and registration (600 s MLflow HTTP timeout during the upload). During a fit MLflow HTTP calls use 2 retries and a 30 s timeout unless `MLFLOW_HTTP_REQUEST_MAX_RETRIES` / `MLFLOW_HTTP_REQUEST_TIMEOUT` are set. Only global rank 0 tracks. -- MLflow rejects `@` in metric names: `test_map@10` is logged as `test_map_at_10`. Params MLflow rejects are counted in `kumo.dropped_params`. -- Snowflake keeps 200 metrics and 1000 params per run; the rest are counted in `kumo.dropped_metrics` / `kumo.dropped_params`. Snowflake has no tags (they are params of the same name), no child runs and no dataset inputs. A taken `run_name` gets a random suffix; an active run elsewhere in the process is never resumed. +- MLflow rejects `@` in metric names: the trainer's `test_map@10` is logged as `test/map_at_10`. Params MLflow rejects are counted in `kumo.dropped_params`. +- Snowflake keeps 200 metrics and 1000 params per run; the rest are counted in `kumo.dropped_metrics` / `kumo.dropped_params`. Snowflake has no tags (they are params of the same name), no child runs and no dataset inputs. A taken run name gets `_2`, `_3` and so on; an active run elsewhere in the process is never resumed. - The fit's source tables are MLflow run inputs: Databricks tables as Delta tables at the snapshot's version (so Unity Catalog links them to the registered model), file sources by URI. From 1476ea739893716c5673b319694d02d2ef0756a3 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Mon, 2 Nov 2026 19:04:13 +0000 Subject: [PATCH 220/278] Skip summary types a tracker does not know instead of raising Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/docs/mlops.md | 2 +- packages/kumo-engine/kumo/engine/_summary_tables.py | 4 +--- packages/kumo-engine/test/test_tracking_contract.py | 6 ++++++ 3 files changed, 8 insertions(+), 4 deletions(-) diff --git a/packages/kumo-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index 7268d2138..ff9fe3e5a 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -301,7 +301,7 @@ A text column `infer_metadata()` reads as categorical (short values from a small ### Custom trackers -Other systems plug in through `kumo.engine.tracking.Tracking`, 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 attempt after each of its events, and `finish(record, error)` with the final record (`None` when tracking stopped early). Optional members: `prepare()`, `tags`, `check(registry)` (when `kumo.MLOps` is built), `log_model(bundle_dir, manifest, snapshot, registry)` (logs and registers, returning the version) or `connector` (a connector whose model registry `registry` registers in), `load_model(location)` and `register(result, location)` (for loading and `result.export(cfg)`), `platform`/`run_id`/`run_url` (recorded as `result.tracking`). A tracker that raises is not called again for that fit and its run ends failed. Source tables arrive in `finish` as `Dataset` summaries in `record.summaries`, each with its `name`, `source` (`connector`, `location`, `version`, `version_kind`), `schema`, `profile` and `context`; a tracker's `tags` are added except `kumo.*` and param names. +Other systems plug in through `kumo.engine.tracking.Tracking`, 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 attempt after each of its events, and `finish(record, error)` with the final record (`None` when tracking stopped early). Optional members: `prepare()`, `tags`, `check(registry)` (when `kumo.MLOps` is built), `log_model(bundle_dir, manifest, snapshot, registry)` (logs and registers, returning the version) or `connector` (a connector whose model registry `registry` registers in), `load_model(location)` and `register(result, location)` (for loading and `result.export(cfg)`), `platform`/`run_id`/`run_url` (recorded as `result.tracking`). A tracker that raises is not called again for that fit and its run ends failed. Source tables arrive in `finish` as `Dataset` summaries in `record.summaries`, each with its `name`, `source` (`connector`, `location`, `version`, `version_kind`), `schema`, `profile` and `context`; a tracker's `tags` are added except `kumo.*` and param names. Skip summary types your tracker does not handle; new ones may be added. ### Progress events diff --git a/packages/kumo-engine/kumo/engine/_summary_tables.py b/packages/kumo-engine/kumo/engine/_summary_tables.py index 91f26a06d..bbc7bf219 100644 --- a/packages/kumo-engine/kumo/engine/_summary_tables.py +++ b/packages/kumo-engine/kumo/engine/_summary_tables.py @@ -25,7 +25,7 @@ class Routed: def route(summaries: Sequence[Summary]) -> Routed: - r"""Sorts summaries by how a tracker writes them.""" + r"""Sorts summaries by how a tracker writes them; skips unknown types.""" routed, tables = Routed(), [] for summary in summaries: if isinstance(summary, Scalar): @@ -39,8 +39,6 @@ def route(summaries: Sequence[Summary]) -> Routed: routed.card.append(summary) elif isinstance(summary, Table): tables.append(summary) - else: - raise TypeError(f'No route for {type(summary).__name__}') routed.tables = tables return routed diff --git a/packages/kumo-engine/test/test_tracking_contract.py b/packages/kumo-engine/test/test_tracking_contract.py index 8c91f930f..0ef99f5b5 100644 --- a/packages/kumo-engine/test/test_tracking_contract.py +++ b/packages/kumo-engine/test/test_tracking_contract.py @@ -528,3 +528,9 @@ def start(self, record): assert elapsed < 1.5 assert isinstance(session.error, TimeoutError) + + +def test_a_summary_type_no_tracker_knows_is_skipped(): + from kumo.engine._summary_tables import Routed, route + + assert route([object()]) == Routed() From 3d3c9d896b8d3a0490f19864419457828c15dd91 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Mon, 2 Nov 2026 19:54:12 +0000 Subject: [PATCH 221/278] Round-trip an artifact through an ADLS file destination Co-Authored-By: Claude Opus 5.5 --- .../kumo-engine/test/test_adls_connector.py | 24 +++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/packages/kumo-engine/test/test_adls_connector.py b/packages/kumo-engine/test/test_adls_connector.py index c3d0c6603..b356a75f5 100644 --- a/packages/kumo-engine/test/test_adls_connector.py +++ b/packages/kumo-engine/test/test_adls_connector.py @@ -189,3 +189,27 @@ def test_invalid_prediction_output_fails_before_scoring(monkeypatch, tmp_path, work.assert_not_called() azure.assert_not_called() assert not (tmp_path / 'runs').exists() + + +def test_an_adls_root_holds_files_and_round_trips_an_artifact( + storage, tmp_path): + from kumo.engine.connector import is_file_destination + + connector = ADLSConnector(root_dir=ROOT) + sink = connector.sink_backend + assert is_file_destination(connector, ROOT + 'artifacts/') + bundle = tmp_path / 'bundle' + (bundle / 'weights').mkdir(parents=True) + (bundle / '_provenance.json').write_bytes(b'{"id": 1}') + (bundle / 'weights' / 'w.bin').write_bytes(b'\x00\x01\x02') + + published = sink.publish_model(connector, str(bundle), + path=ROOT + 'artifacts/m') + fetched = tmp_path / 'fetched' + sink.fetch_model(connector, str(fetched), path=published.uri) + + stored = storage / ('curated@companylake.dfs.core.windows.net/' + 'orders/artifacts/m') + assert (stored / 'weights' / 'w.bin').read_bytes() == b'\x00\x01\x02' + for name in ('_provenance.json', 'weights/w.bin'): + assert (fetched / name).read_bytes() == (bundle / name).read_bytes() From 3dfba070199706029e70b4d354280f5310cf91e7 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Mon, 2 Nov 2026 20:07:05 +0000 Subject: [PATCH 222/278] Treat a chunked writer as failed until it opens, so a failed open never publishes an empty file Co-Authored-By: Claude Opus 5.5 --- .../kumo-core/src/kumo_core/chunked_fs.py | 44 ++++++++++++------- packages/kumo-core/tests/test_chunked_fs.py | 21 +++++++++ 2 files changed, 48 insertions(+), 17 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/chunked_fs.py b/packages/kumo-core/src/kumo_core/chunked_fs.py index 2124096ca..7b1c654c9 100644 --- a/packages/kumo-core/src/kumo_core/chunked_fs.py +++ b/packages/kumo-core/src/kumo_core/chunked_fs.py @@ -313,27 +313,35 @@ def _abandon(stream: Optional[BinaryIO], close: bool = False) -> None: class _Writer(io.RawIOBase): def __init__(self, handler: ChunkedFileHandler, path: str, size: Optional[int] = None) -> None: + self._failed = True self._handler = handler self._path = path self._chunks = (f'{path}{CHUNKS}/{time.time_ns():020d}-' f'{uuid.uuid4().hex[:8]}') - self._dir = tempfile.mkdtemp(prefix='kumo-chunked-', - dir=spool_root(handler)) - self._spool = open(os.path.join(self._dir, 'spool'), 'wb') - self._filled = 0 - self._parts = 0 - self._written = 0 - self._part: Optional[BinaryIO] = None - self._streams = type(handler)._stream is not ChunkedFileHandler._stream - self._limit = (min(handler.plain_bytes, handler.part_bytes) - if self._streams else handler.part_bytes) self._plain: Optional[BinaryIO] = None + self._part: Optional[BinaryIO] = None + self._spool: Optional[BinaryIO] = None + self._dir: Optional[str] = None + try: + self._dir = tempfile.mkdtemp(prefix='kumo-chunked-', + dir=spool_root(handler)) + self._spool = open(os.path.join(self._dir, 'spool'), 'wb') + self._filled = 0 + self._parts = 0 + self._written = 0 + self._streams = (type(handler)._stream + is not ChunkedFileHandler._stream) + self._limit = (min(handler.plain_bytes, handler.part_bytes) + if self._streams else handler.part_bytes) + if self._streams and size is not None: + if size <= handler.part_bytes: + self._plain = handler._stream(path) + else: + self._part = self._open_part() + except BaseException: + self.close() + raise self._failed = False - if self._streams and size is not None: - if size <= handler.part_bytes: - self._plain = handler._stream(path) - else: - self._part = self._open_part() def writable(self): return True @@ -442,6 +450,8 @@ def close(self): self._handler._remove(self._path + CHUNKS, tree=True) finally: self._handler._forget(self._path) - self._spool.close() - shutil.rmtree(self._dir, ignore_errors=True) + if self._spool is not None: + self._spool.close() + if self._dir is not None: + shutil.rmtree(self._dir, ignore_errors=True) super().close() diff --git a/packages/kumo-core/tests/test_chunked_fs.py b/packages/kumo-core/tests/test_chunked_fs.py index e569d4bef..f4a43c558 100644 --- a/packages/kumo-core/tests/test_chunked_fs.py +++ b/packages/kumo-core/tests/test_chunked_fs.py @@ -379,3 +379,24 @@ def test_reads_spanning_parts_fill_the_callers_buffer(store): assert reader.readinto(buffer) == 2 * PART assert bytes(buffer) == data[PART - 3:3 * PART - 3] assert reader.read_at(4, 2) == data[2:6] + + +@pytest.mark.parametrize('size', [1, 3 * PART]) +def test_a_writer_that_fails_to_open_keeps_the_old_file( + monkeypatch, tmp_path, size): + import gc + + handler = StreamingHandler() + handler.spool_dir = str(tmp_path) + fs = PyFileSystem(handler) + _write(fs, 'root/a.bin', b'preserve-me') + + def fail(path): + raise OSError('cannot open') + + monkeypatch.setattr(handler, '_stream', fail) + with pytest.raises(OSError, match='cannot open'): + chunked_fs._Writer(handler, 'root/a.bin', size=size) + gc.collect() + assert _read(fs, 'root/a.bin') == b'preserve-me' + assert os.listdir(tmp_path) == [] From 5f75e66a5f4a3cb71accb9944f15015d950352f2 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Mon, 2 Nov 2026 20:07:33 +0000 Subject: [PATCH 223/278] Drop only older committed generations, and read the newest one whose parts are complete Co-Authored-By: Claude Opus 5.5 --- .../kumo-core/src/kumo_core/chunked_fs.py | 42 +++++++++++-------- packages/kumo-core/tests/test_chunked_fs.py | 34 +++++++++++++++ 2 files changed, 58 insertions(+), 18 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/chunked_fs.py b/packages/kumo-core/src/kumo_core/chunked_fs.py index 7b1c654c9..61d0e0375 100644 --- a/packages/kumo-core/src/kumo_core/chunked_fs.py +++ b/packages/kumo-core/src/kumo_core/chunked_fs.py @@ -187,17 +187,19 @@ def open_output_stream(self, path, metadata): def _logical_parts(objects: Dict[str, int], directory: str) -> Optional[List[Tuple[str, int]]]: - committed = [ + committed = sorted({ posixpath.dirname(path) for path in objects if posixpath.basename(path) == DONE and posixpath.dirname(posixpath.dirname(path)) == directory - ] - if not committed: - return None - generation = max(committed) - return sorted((path, size) for path, size in objects.items() - if posixpath.dirname(path) == generation - and posixpath.basename(path) != DONE) + }, reverse=True) + for generation in committed: + parts = sorted((path, size) for path, size in objects.items() + if posixpath.dirname(path) == generation + and posixpath.basename(path) != DONE) + if [posixpath.basename(path) for path, _ in parts] == [ + f'{index:05d}' for index in range(len(parts))]: + return parts + return None def _logical_files(objects: Dict[str, int]) -> Dict[str, int]: @@ -396,12 +398,16 @@ def _commit_parts(self) -> None: self._upload(f'{self._chunks}/{DONE}') if self._handler._size(self._path) is not None: self._handler._remove(self._path, tree=False) - stale = { - posixpath.dirname(path) - for path in self._handler._list(self._path + CHUNKS) - } - {self._chunks} - for generation in stale: - self._handler._remove(generation, tree=True) + self._drop_older() + + def _drop_older(self) -> None: + objects = self._handler._list(self._path + CHUNKS) + for generation in { + posixpath.dirname(path) for path in objects + if posixpath.basename(path) == DONE + }: + if generation < self._chunks: + self._handler._remove(generation, tree=True) def _flush_part(self) -> None: self._upload(f'{self._chunks}/{self._parts:05d}') @@ -432,11 +438,12 @@ def close(self): if self._failed: _abandon(self._plain) _abandon(self._part, close=True) + with contextlib.suppress(Exception): + self._handler._remove(self._chunks, tree=True) elif self._plain is not None: self._plain.close() self._plain = None - if self._handler._list(self._path + CHUNKS): - self._handler._remove(self._path + CHUNKS, tree=True) + self._drop_older() elif self._part is not None: self._part.close() self._part = None @@ -446,8 +453,7 @@ def close(self): self._commit_parts() else: self._upload(self._path) - if self._handler._list(self._path + CHUNKS): - self._handler._remove(self._path + CHUNKS, tree=True) + self._drop_older() finally: self._handler._forget(self._path) if self._spool is not None: diff --git a/packages/kumo-core/tests/test_chunked_fs.py b/packages/kumo-core/tests/test_chunked_fs.py index f4a43c558..2eea963eb 100644 --- a/packages/kumo-core/tests/test_chunked_fs.py +++ b/packages/kumo-core/tests/test_chunked_fs.py @@ -400,3 +400,37 @@ def fail(path): gc.collect() assert _read(fs, 'root/a.bin') == b'preserve-me' assert os.listdir(tmp_path) == [] + + +@pytest.mark.parametrize('newer_commits_first', [True, False]) +def test_overlapping_writers_on_one_path_never_break_each_other( + store, newer_commits_first): + handler, fs = store + older = fs.open_output_stream('root/a.bin') + older.write(b'o' * 19) + newer = fs.open_output_stream('root/a.bin') + newer.write(b'n' * 21) + first, second = (newer, older) if newer_commits_first else (older, newer) + first.close() + older_rest = b'o' * 18 if newer_commits_first else b'' + if older_rest: + older.write(older_rest) + second.close() + handler._forget('root/a.bin') + assert _read(fs, 'root/a.bin') == b'n' * 21 + assert fs.get_file_info('root/a.bin').size == 21 + + +def test_a_committed_generation_with_missing_parts_is_skipped(store): + handler, fs = store + _write(fs, 'root/a.bin', b'1' * (2 * PART)) + (good, _), = _generations(handler).items() + newer = good.rsplit('/', 1)[0] + '/99999999999999999999-ffffffff' + handler.objects[f'{newer}/00001'] = b'2' * PART + handler.objects[f'{newer}/_SUCCESS'] = b'' + handler._forget('root/a.bin') + assert _read(fs, 'root/a.bin') == b'1' * (2 * PART) + for name in [n for n in handler.objects if n.startswith(good + '/')]: + del handler.objects[name] + handler._forget('root/a.bin') + assert fs.get_file_info('root/a.bin').type is FileType.NotFound From 923000ead76fad9399ef2ed53cdf4a9c364f56d5 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Mon, 2 Nov 2026 20:10:05 +0000 Subject: [PATCH 224/278] Document that MLOps applies per call or block, not per model Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/docs/mlops.md | 1 + 1 file changed, 1 insertion(+) diff --git a/packages/kumo-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index cb2cde19a..1a0badcb9 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -40,6 +40,7 @@ trainer.fit(table, snapshot=snapshot, mlops=cfg) - **`on_error`**: `"warn"` (default) warns about a failed store and returns; `"raise"`, for scheduled jobs, raises once the call has finished, with the failure as the `__cause__`. After `fit`, `trainer.result` still holds the model. - **Precedence:** a call's `mlops=` replaces the enclosing block, and an inner block replaces an outer one. A block covers the calling thread or asyncio task. - **Off by default.** Without a block or `mlops=`, every call behaves as it always has; plain exports (`export(config)`, `export_config=`) never write to the artifact root. +- **Per call, not per model.** MLOps applies only to a call given `mlops=` or run inside `with kumo.mlops(...)`. A later call outside the block, even on a model trained inside one, runs without MLOps, and `on_error="raise"` does not detect this: it covers only MLOps work that ran. There is no per-call off switch; to run one call without MLOps, call it outside the block. For scheduled jobs, wrap the whole job in one block. ### Reading artifacts back From 2f735365b5e3d16b9e7d94cb7fd480aa20c77b5c Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Mon, 2 Nov 2026 20:23:21 +0000 Subject: [PATCH 225/278] Describe the pyfunc as a platform adapter, not a serving or BP replacement Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/docs/mlops.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/kumo-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index 926b54cc7..5e51b1224 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -198,7 +198,7 @@ Jobs that load `alias="champion"` pick up the new version on their next run. ### MLflow pyfunc -A model registered through `MlflowTracking` is an MLflow pyfunc: it serves with `mlflow.pyfunc.load_model(uri)`, `spark_udf` and Model Serving. +A model registered through `MlflowTracking` is also an MLflow pyfunc. It is a platform adapter: a way to load and score a registered Kumo model through MLflow's model API (`mlflow.pyfunc.load_model(uri)`, `spark_udf`, or `mlflow models serve`). It does not replace `Trainer.predict`, the batch prediction API. It is not a low-latency online service: each call scores the rows it is given against a snapshot. ```python with cfg.tracking.client(): # MLflow's own calls reach the tracker's server, as its identity @@ -212,7 +212,7 @@ scores = model.predict(pd.DataFrame({"ENTITY": [1, 2, 3]}), params={ - Each replica loads a snapshot on its first `predict` with that URI and keeps it. `log_model(trainer, snapshot=..., serving_snapshot=...)` logs a model whose replicas load that snapshot at startup and use it when no `snapshot` param is given; `serving_snapshot` takes a URI, or a `GraphSnapshot` or snapshot id stored under the block's artifact root. - Snapshot URIs: `s3://`, `gs://`, `file://` or a local path, read with the ambient credentials. A `/Volumes/...` path off a Databricks cluster, or a stage path, is read through `model_config={"assets_from": connector}` (below), e.g. the connector of a `(connector, "/Volumes/...")` root. - A model with stored assets (word vectors for text columns) on a volume, stage or bucket reads them, and snapshots there, with the identity it is given: `mlflow.pyfunc.load_model(uri, model_config={"assets_from": connector})`, a connector for that store with its credentials; an asset it cannot find raises, naming this option. -- Serving needs an image with kumo-engine installed: `mlflow models serve -m --env-manager local`. Its pip requirement is `kumo-engine[s3,gcs]` of the same version (plus `[text]` for text columns). +- `mlflow models serve` needs an image with kumo-engine installed: `mlflow models serve -m --env-manager local`. Its pip requirement is `kumo-engine[s3,gcs]` of the same version (plus `[text]` for text columns). ## Credentials and running off-platform From f1c6773459474b198d5552c158c9b600e7c87be2 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Mon, 2 Nov 2026 20:28:58 +0000 Subject: [PATCH 226/278] Stamp artifact columns as Arrow constants instead of per-row Python lists Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_core/artifact_export.py | 3 +-- .../kumo-core/tests/test_artifact_export.py | 25 +++++++++++++++++++ 2 files changed, 26 insertions(+), 2 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/artifact_export.py b/packages/kumo-core/src/kumo_core/artifact_export.py index 4c99e5927..cac3cfa28 100644 --- a/packages/kumo-core/src/kumo_core/artifact_export.py +++ b/packages/kumo-core/src/kumo_core/artifact_export.py @@ -307,8 +307,7 @@ def prepared_artifact_file(path, extra_columns=None, *, parquet_options=None, # Pandas turns all-null bool columns into Arrow nulls. table = pa.Table.from_batches([batch]) for name, value in (extra_columns or {}).items(): - column = pa.array([value] * table.num_rows, - type=pa.scalar(value).type) + 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 diff --git a/packages/kumo-core/tests/test_artifact_export.py b/packages/kumo-core/tests/test_artifact_export.py index e74b28ff2..244a121ec 100644 --- a/packages/kumo-core/tests/test_artifact_export.py +++ b/packages/kumo-core/tests/test_artifact_export.py @@ -153,6 +153,31 @@ def test_stamped_columns_keep_the_source_types(tmp_path): pa.field('BUNDLE_ID', pa.string())) +def test_stamping_builds_arrow_columns_not_python_lists(tmp_path, monkeypatch): + import pyarrow as pa + import pyarrow.parquet as pq + from kumo_core.artifact_export import prepared_artifact_file + + source = tmp_path / 'scores.parquet' + pq.write_table(pa.table({'ENTITY': np.arange(200000, dtype=np.int64)}), + source) + real_array = pa.array + + def guarded(values, *args, **kwargs): + assert not isinstance(values, list) or len(values) < 1000 + return real_array(values, *args, **kwargs) + + monkeypatch.setattr(pa, 'array', guarded) + extra = {'BUNDLE_ID': 'b', 'RANK': 7, 'FLAG': True} + with prepared_artifact_file(str(source), extra) as staged: + got = pq.read_table(staged) + assert got.schema.types[1:] == [pa.string(), pa.int64(), pa.bool_()] + assert got.num_rows == 200000 + assert got['BUNDLE_ID'].unique().to_pylist() == ['b'] + assert got['RANK'].unique().to_pylist() == [7] + assert got['FLAG'].unique().to_pylist() == [True] + + @pytest.mark.parametrize('existing_empty', [False, True]) def test_multipart_failure_and_failed_recopy_remain_retryable( tmp_path, monkeypatch, existing_empty): From 8143a79696fdcc1538bff23ae1f8181fae8985d6 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Mon, 2 Nov 2026 20:28:58 +0000 Subject: [PATCH 227/278] Document the pyfunc population param and its memory cost Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/docs/mlops.md | 1 + 1 file changed, 1 insertion(+) diff --git a/packages/kumo-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index 5e51b1224..90bf084da 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -209,6 +209,7 @@ scores = model.predict(pd.DataFrame({"ENTITY": [1, 2, 3]}), params={ ``` - Input rows have an `ENTITY` column and optionally a `TIMESTAMP`; a temporal query needs one of `TIMESTAMP` or `prediction_time`. Scores come back in input order, one per scored row. Binary classification returns `True_PROB`; `TARGET_PRED` needs `params={"binary_classification_threshold": ...}` or a stored threshold, and is null without one. +- `params={"population": True}` with an empty input frame scores the trained query's whole population and returns it as one DataFrame. It holds all scored rows in memory, so for a large population use `Trainer.predict(...)` and read its `files`. - Each replica loads a snapshot on its first `predict` with that URI and keeps it. `log_model(trainer, snapshot=..., serving_snapshot=...)` logs a model whose replicas load that snapshot at startup and use it when no `snapshot` param is given; `serving_snapshot` takes a URI, or a `GraphSnapshot` or snapshot id stored under the block's artifact root. - Snapshot URIs: `s3://`, `gs://`, `file://` or a local path, read with the ambient credentials. A `/Volumes/...` path off a Databricks cluster, or a stage path, is read through `model_config={"assets_from": connector}` (below), e.g. the connector of a `(connector, "/Volumes/...")` root. - A model with stored assets (word vectors for text columns) on a volume, stage or bucket reads them, and snapshots there, with the identity it is given: `mlflow.pyfunc.load_model(uri, model_config={"assets_from": connector})`, a connector for that store with its credentials; an asset it cannot find raises, naming this option. From 5fb2a8d827478b16e3cb25cabf593c725ccf22e3 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Mon, 2 Nov 2026 20:30:34 +0000 Subject: [PATCH 228/278] Reopen a volume download after a failed read Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_databricks/volume.py | 20 +++++++++------- packages/kumo-databricks/tests/test_volume.py | 24 +++++++++++++++++++ 2 files changed, 36 insertions(+), 8 deletions(-) diff --git a/packages/kumo-databricks/src/kumo_databricks/volume.py b/packages/kumo-databricks/src/kumo_databricks/volume.py index 1f0657de1..e5face782 100644 --- a/packages/kumo-databricks/src/kumo_databricks/volume.py +++ b/packages/kumo-databricks/src/kumo_databricks/volume.py @@ -118,14 +118,18 @@ def read_at(self, offset: int, length: int) -> bytes: return b'' if self._stream is None or offset != self._at: self._open(offset) - with _translated(self._path): - pieces, got = [], 0 - while got < length: - piece = self._stream.read(min(_BUFFER, length - got)) - if not piece: - break - pieces.append(piece) - got += len(piece) + try: + with _translated(self._path): + pieces, got = [], 0 + while got < length: + piece = self._stream.read(min(_BUFFER, length - got)) + if not piece: + break + pieces.append(piece) + got += len(piece) + except BaseException: + self.close() + raise self._at += got return b''.join(pieces) diff --git a/packages/kumo-databricks/tests/test_volume.py b/packages/kumo-databricks/tests/test_volume.py index 6bcdf8bcc..4ae15d24d 100644 --- a/packages/kumo-databricks/tests/test_volume.py +++ b/packages/kumo-databricks/tests/test_volume.py @@ -89,6 +89,30 @@ def test_a_download_that_ignores_the_range_is_refused(): reader.read_at(10, 100) +def test_a_read_after_a_dropped_download_starts_where_asked(): + fs, volume = filesystem({f'{ROOT}/a': BODY}) + do = volume.do + + def dropping(*args, **kwargs): + volume.do = do + response = do(*args, **kwargs) + stream = response['contents'] + + def read(size, read=stream.read): + read(size // 2) + stream.read = read + raise ConnectionError('connection reset') + + stream.read = read + return response + + volume.do = dropping + with fs.open_input_file(f'{ROOT}/a') as reader: + with pytest.raises(OSError, match='ConnectionError'): + reader.read_at(10, 100) + assert reader.read_at(10, 100) == BODY[100:110] + + def test_directories_and_missing_paths(): fs, volume = filesystem({ f'{ROOT}/s/kumo_snapshot.json': b'{}', From a39523659ebd93b4a3b4450d5e610cfad04bd4e8 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Mon, 2 Nov 2026 20:38:22 +0000 Subject: [PATCH 229/278] Drop the pyfunc population param; Trainer.predict scores a population Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/docs/mlops.md | 1 - .../kumo-engine/kumo/engine/_predictor.py | 4 ---- .../kumo/engine/integrations/_mlflow_model.py | 3 +-- .../kumo-engine/test/test_mlflow_model.py | 23 ++++--------------- 4 files changed, 6 insertions(+), 25 deletions(-) diff --git a/packages/kumo-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index 90bf084da..5e51b1224 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -209,7 +209,6 @@ scores = model.predict(pd.DataFrame({"ENTITY": [1, 2, 3]}), params={ ``` - Input rows have an `ENTITY` column and optionally a `TIMESTAMP`; a temporal query needs one of `TIMESTAMP` or `prediction_time`. Scores come back in input order, one per scored row. Binary classification returns `True_PROB`; `TARGET_PRED` needs `params={"binary_classification_threshold": ...}` or a stored threshold, and is null without one. -- `params={"population": True}` with an empty input frame scores the trained query's whole population and returns it as one DataFrame. It holds all scored rows in memory, so for a large population use `Trainer.predict(...)` and read its `files`. - Each replica loads a snapshot on its first `predict` with that URI and keeps it. `log_model(trainer, snapshot=..., serving_snapshot=...)` logs a model whose replicas load that snapshot at startup and use it when no `snapshot` param is given; `serving_snapshot` takes a URI, or a `GraphSnapshot` or snapshot id stored under the block's artifact root. - Snapshot URIs: `s3://`, `gs://`, `file://` or a local path, read with the ambient credentials. A `/Volumes/...` path off a Databricks cluster, or a stage path, is read through `model_config={"assets_from": connector}` (below), e.g. the connector of a `(connector, "/Volumes/...")` root. - A model with stored assets (word vectors for text columns) on a volume, stage or bucket reads them, and snapshots there, with the identity it is given: `mlflow.pyfunc.load_model(uri, model_config={"assets_from": connector})`, a connector for that store with its credentials; an asset it cannot find raises, naming this option. diff --git a/packages/kumo-engine/kumo/engine/_predictor.py b/packages/kumo-engine/kumo/engine/_predictor.py index 6d3be32d4..f002968ea 100644 --- a/packages/kumo-engine/kumo/engine/_predictor.py +++ b/packages/kumo-engine/kumo/engine/_predictor.py @@ -80,10 +80,6 @@ def predict_frame(trainer: Any, snapshot: Any, rows: pd.DataFrame, trainer.scratch_dir.mkdir(parents=True, exist_ok=True) with tempfile.TemporaryDirectory(dir=trainer.scratch_dir) as tmp, \ scoped_scratch_dir(Path(tmp) / 'scratch'): - if len(rows) == 0: - return trainer.predict( - snapshot, prediction_time=None - if time is None else time.to_pydatetime(), **options).frame rows = rows.rename(columns={'ENTITY': 'ENTITY1'}) if time is not None and 'TIMESTAMP' not in rows: rows = rows.assign(TIMESTAMP=time) diff --git a/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py b/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py index 40a9e5e1c..8233a4094 100644 --- a/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py +++ b/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py @@ -43,7 +43,7 @@ def load_context(self, context: Any) -> None: def predict(self, context: Any, model_input: pd.DataFrame, params: Optional[Dict[str, Any]] = None) -> pd.DataFrame: params = params or {} - if len(model_input) == 0 and not params.get('population'): + if len(model_input) == 0: return pd.DataFrame(columns=self._columns) uri = params.get('snapshot') or self._default if not uri: @@ -115,7 +115,6 @@ def signature(snapshot: GraphSnapshot, query: str) -> Tuple[Any, pd.DataFrame]: _schema.ParamSpec('prediction_time', 'string', None), _schema.ParamSpec('binary_classification_threshold', 'double', None), _schema.ParamSpec('num_classes_to_return', 'long', None), - _schema.ParamSpec('population', 'boolean', False), ]) example = pd.DataFrame({ 'ENTITY': { diff --git a/packages/kumo-engine/test/test_mlflow_model.py b/packages/kumo-engine/test/test_mlflow_model.py index e3b51d17a..4f2885104 100644 --- a/packages/kumo-engine/test/test_mlflow_model.py +++ b/packages/kumo-engine/test/test_mlflow_model.py @@ -124,8 +124,7 @@ def other_bundle(path): load_model('models:/m@champion') -def test_logged_model_scores_the_query_population_and_timed_rows( - fitted, tracking_uri, tmp_path, monkeypatch): +def test_logged_model_scores_given_rows(fitted, tracking_uri, tmp_path): from kumo.engine.integrations.mlflow import log_model trainer, table, snapshot = fitted @@ -162,32 +161,20 @@ def test_logged_model_scores_the_query_population_and_timed_rows( empty = pd.DataFrame({'ENTITY': pd.Series([], dtype='int64')}) tables = get_scratch_dir() / 'prediction_tables' before = set(tables.glob('*')) - predict = kumo.Trainer.predict - - def predict_beside_another_call(self, *args, **kwargs): - (tables / 'other-call').mkdir(parents=True, exist_ok=True) - return predict(self, *args, **kwargs) - - monkeypatch.setattr(kumo.Trainer, 'predict', predict_beside_another_call) assert pyfunc.predict(empty, params=params).empty served = pyfunc.unwrap_python_model()._snapshots assert list(served) == [params['snapshot']] - pd.testing.assert_frame_equal( - _sorted( - pyfunc.predict(empty, params={ - **params, 'snapshot': None, 'population': True - })), expected) - for rows, population in ((empty, True), - (pd.DataFrame({'ENTITY': range(50)}), False)): + rows = pd.DataFrame({'ENTITY': range(50)}) + for snapshot_param in (None, params['snapshot']): pd.testing.assert_frame_equal( _sorted( pyfunc.predict(rows, - params={**params, 'population': population})), + params={**params, 'snapshot': snapshot_param})), expected) scratch = pyfunc.unwrap_python_model()._predictor.trainer.scratch_dir assert list(scratch.iterdir()) == [] - assert set(tables.glob('*')) == before | {tables / 'other-call'} + assert set(tables.glob('*')) == before def test_pyfunc_names_a_missing_prediction_time_and_needs_no_threshold( From 503f5781a6611604bc7827b41d40c249f194d4f3 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Mon, 2 Nov 2026 20:42:13 +0000 Subject: [PATCH 230/278] Stub the SageMaker experiment lookup with an experiment, so the test stays offline Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/test/test_mlflow_target.py | 1 + 1 file changed, 1 insertion(+) diff --git a/packages/kumo-engine/test/test_mlflow_target.py b/packages/kumo-engine/test/test_mlflow_target.py index cb5071b4d..5f7122bae 100644 --- a/packages/kumo-engine/test/test_mlflow_target.py +++ b/packages/kumo-engine/test/test_mlflow_target.py @@ -198,6 +198,7 @@ def test_ambient_sagemaker_calls_see_its_uri_and_restore_the_global_one( def lookup(self, name): seen.append((self.tracking_uri, mlflow.get_tracking_uri())) + return SimpleNamespace(experiment_id='1') monkeypatch.setattr(mlflow.MlflowClient, 'get_experiment_by_name', lookup) From 64ea65c1ae78990e5d9e2df11cbb1c1b44ceb7e8 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Mon, 2 Nov 2026 21:34:05 +0000 Subject: [PATCH 231/278] Upload each stored training-table part as it is written Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/pquery/training_table.py | 66 ++++++++++++++++--- .../test/test_training_table_export.py | 30 +++++++++ 2 files changed, 86 insertions(+), 10 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/pquery/training_table.py b/packages/kumo-engine/kumo/engine/pquery/training_table.py index c805684c7..72c7b9da4 100644 --- a/packages/kumo-engine/kumo/engine/pquery/training_table.py +++ b/packages/kumo-engine/kumo/engine/pquery/training_table.py @@ -4,10 +4,13 @@ import hashlib import json +import os import warnings +from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field, replace +from functools import partial from pathlib import Path -from typing import TYPE_CHECKING, Any, Dict, Optional +from typing import TYPE_CHECKING, Any, Dict, Iterator, Optional from kumo_core.artifacts import PROVENANCE_FILE from kumo_core.model_plan import TrainingTableGenerationPlan @@ -104,18 +107,23 @@ def _export_sidecars(self, connector, output_config) -> Dict[str, Any]: from kumo.engine.connector import is_file_destination - manifest = self.staging_dir / PROVENANCE_FILE name = output_config.table_name.rstrip('/') + if (not is_file_destination(connector, output_config.table_name) + or layout.file_type(name) == layout.CSV): + return {} + sidecars = self._sidecars() + return {'sidecars': sidecars} if sidecars else {} + + def _sidecars(self) -> Dict[str, bytes]: + manifest = self.staging_dir / PROVENANCE_FILE try: - if (not is_file_destination(connector, output_config.table_name) - or layout.file_type(name) == layout.CSV - or not manifest.is_file()): + if not manifest.is_file(): return {} - return {'sidecars': {PROVENANCE_FILE: manifest.read_bytes()}} + return {PROVENANCE_FILE: manifest.read_bytes()} except Exception as error: warnings.warn( f"Exporting without {PROVENANCE_FILE}: {error}", - UserWarning, stacklevel=3) + UserWarning, stacklevel=4) return {} @classmethod @@ -134,20 +142,41 @@ def _store(self, settings: Any) -> None: from pyarrow.fs import FileType from kumo.engine.connector import file_store - from kumo.engine.trainer.output_config import TrainingTableOutputConfig path = settings.root.location(f'training_tables/{self.id}/') try: filesystem, root = file_store(settings.root.connector, path) if filesystem.get_file_info( f'{root}/_SUCCESS').type is not FileType.File: - self.export( - TrainingTableOutputConfig(settings.root.connector, path)) + self._upload(filesystem, f'{root}/') except Exception as error: raise RuntimeError( f"Could not store training table {self.id} under the " f"artifact root: {error}") from error + def _upload(self, filesystem: Any, path: str) -> None: + import tempfile + + from kumo_core.artifact_export import export_selected_files + from kumo_core.object_store_sink import _ArrowExportFiles + from kumo_core.training.export import ( + identity, + iter_export_parts, + read_spec, + ) + + from kumo.engine._scratch import get_scratch_dir + + root = get_scratch_dir() + root.mkdir(parents=True, exist_ok=True) + with tempfile.TemporaryDirectory(dir=root) as spill: + export_selected_files( + iter_export_parts(str(self.staging_dir), spill), path, + fs=_ArrowExportFiles(filesystem, None), + identity=identity(read_spec(str(self.staging_dir))), + copy_parquet=partial(_move_each, filesystem), + sidecars=self._sidecars()) + @classmethod def from_source_table( cls, @@ -287,6 +316,23 @@ def from_source_table( source_version=version) +def _move_each(filesystem: Any, parts: Iterator[str], output: str) -> None: + from kumo_core.object_store_sink import upload_file + + def move(part: str) -> None: + upload_file(filesystem, part, output + os.path.basename(part)) + os.remove(part) + + with ThreadPoolExecutor(max_workers=1) as pool: + pending = [] + for part in parts: + pending.append(pool.submit(move, part)) + if len(pending) > 1: + pending.pop(0).result() + for upload in pending: + upload.result() + + def _record_sources(snapshot: GraphSnapshot) -> bool: if snapshot.record_sources is not None: return snapshot.record_sources diff --git a/packages/kumo-engine/test/test_training_table_export.py b/packages/kumo-engine/test/test_training_table_export.py index 5bd5486e3..8914fe598 100644 --- a/packages/kumo-engine/test/test_training_table_export.py +++ b/packages/kumo-engine/test/test_training_table_export.py @@ -417,6 +417,36 @@ def write_files(connector, table_name, files, write_mode, **kwargs): assert not (out / 'rows.csv' / '_provenance.json').exists() +def test_the_store_spills_at_most_two_parts(graph, tmp_path, monkeypatch): + import os + + from kumo_core import object_store_sink + from kumo_core.training import export + + graph_, snapshot, _, _ = graph + stage_parts = export.stage_parts + monkeypatch.setattr( + export, 'stage_parts', lambda spec: { + split: parts * 4 + for split, parts in stage_parts(spec).items() + }) + upload_file = object_store_sink.upload_file + spilled = [] + + def upload(filesystem, source, destination): + spilled.append(len(os.listdir(os.path.dirname(source)))) + upload_file(filesystem, source, destination) + + monkeypatch.setattr(object_store_sink, 'upload_file', upload) + table = kumo.PredictiveQuery( + graph=graph_, query='PREDICT ents.LABEL FOR EACH ents.ID' + ).generate_training_table( + snapshot, mlops=kumo.MLOps(artifact_root=str(tmp_path))) + stored = tmp_path / 'training_tables' / table.id + assert len(spilled) == len(list(stored.glob('part-*.parquet'))) >= 8 + assert max(spilled) <= 2 + + def test_a_mlops_block_stores_each_table_once_under_the_root(graph, tmp_path): graph_, snapshot, _, _ = graph pquery = kumo.PredictiveQuery(graph=graph_, From ea5d22d9e0ddad00bca01d8b6e4d379de54cae34 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Mon, 2 Nov 2026 21:36:00 +0000 Subject: [PATCH 232/278] Pass an export's sidecars through a file-store sink's routing Co-Authored-By: Claude Opus 5.5 --- .../kumo-core/src/kumo_core/object_store_sink.py | 15 +++++++++------ packages/kumo-engine/test/test_filestore_sink.py | 10 ++++++++++ 2 files changed, 19 insertions(+), 6 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/object_store_sink.py b/packages/kumo-core/src/kumo_core/object_store_sink.py index b7e65b523..e416cc975 100644 --- a/packages/kumo-core/src/kumo_core/object_store_sink.py +++ b/packages/kumo-core/src/kumo_core/object_store_sink.py @@ -513,12 +513,15 @@ def write_table(self, connector, table_name, frame, def write_files(self, connector, table_name, files, write_mode=WriteMode.OVERWRITE, project_id=None, *, - export_id=None, extra_columns=None): - sink = FilesystemSink(self) if self.holds_files( - connector, table_name) else self._rest('write_files') - return sink.write_files(connector, table_name, files, write_mode, - project_id, export_id=export_id, - extra_columns=extra_columns) + export_id=None, extra_columns=None, sidecars={}): + if self.holds_files(connector, table_name): + return FilesystemSink(self).write_files( + connector, table_name, files, write_mode, project_id, + export_id=export_id, extra_columns=extra_columns, + sidecars=sidecars) + return self._rest('write_files').write_files( + connector, table_name, files, write_mode, project_id, + export_id=export_id, extra_columns=extra_columns) def publish_model(self, connector, bundle_dir, *, path=None, location=None, metadata={}, metrics={}, project_id=None): diff --git a/packages/kumo-engine/test/test_filestore_sink.py b/packages/kumo-engine/test/test_filestore_sink.py index cb06a2206..449ea0a79 100644 --- a/packages/kumo-engine/test/test_filestore_sink.py +++ b/packages/kumo-engine/test/test_filestore_sink.py @@ -227,3 +227,13 @@ def test_a_file_only_sink_refuses_a_table(connector, tmp_path): with pytest.raises(NotImplementedError, match='outside its files'): connector.sink_backend.publish_model(connector, str(tmp_path), path='DB.SCHEMA.T') + + +def test_a_file_store_export_keeps_its_sidecars(tmp_path, connector): + part = tmp_path / 'part-0.parquet' + pd.DataFrame({'x': range(3)}).to_parquet(part) + connector.sink_backend.write_files( + connector, '@stage/tt/', [str(part)], + sidecars={'_provenance.json': b'{}'}) + assert connector.sink_backend.store.objects[ + '@stage/tt/_provenance.json'] == b'{}' From 086c97fc141fc268cfdf4869b30c8dc83e8a21e7 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Mon, 2 Nov 2026 21:57:45 +0000 Subject: [PATCH 233/278] Keep unsized volume exports as plain objects up to a part CSV exports open the stream without a size, so anything over 256 MiB rolled over to .chunks parts that Spark cannot read as one file. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-databricks/src/kumo_databricks/volume.py | 2 +- packages/kumo-databricks/tests/test_volume.py | 9 +++++++++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/packages/kumo-databricks/src/kumo_databricks/volume.py b/packages/kumo-databricks/src/kumo_databricks/volume.py index e5face782..6d85765d9 100644 --- a/packages/kumo-databricks/src/kumo_databricks/volume.py +++ b/packages/kumo-databricks/src/kumo_databricks/volume.py @@ -51,7 +51,7 @@ def _translated(path: str) -> Iterator[None]: class VolumeHandler(ChunkedFileHandler): part_bytes = 4 << 30 # under the Files API's 5 GiB upload cap - plain_bytes = 256 << 20 + plain_bytes = part_bytes def __init__(self, client: WorkspaceClient) -> None: self._client = client diff --git a/packages/kumo-databricks/tests/test_volume.py b/packages/kumo-databricks/tests/test_volume.py index 4ae15d24d..71de4a7eb 100644 --- a/packages/kumo-databricks/tests/test_volume.py +++ b/packages/kumo-databricks/tests/test_volume.py @@ -62,6 +62,15 @@ def test_sized_writes_stay_plain_up_to_a_part(): and path.endswith('/_SUCCESS') for path in volume.objects) +def test_unsized_writes_stay_plain_through_a_part(): + assert VolumeHandler.plain_bytes == VolumeHandler.part_bytes == 4 << 30 + fs, volume = filesystem(part_bytes=1000) + write(fs, f'{ROOT}/export.csv', BODY[:1000]) + assert volume.objects[f'{ROOT}/export.csv'] == BODY[:1000] + write(fs, f'{ROOT}/over.csv', BODY[:1001]) + assert f'{ROOT}/over.csv' not in volume.objects + + def test_a_failed_upload_raises_without_the_signed_url(): from kumo_core.object_store_sink import open_sized From 2bc33514138eca4797ec63fbb84becaf7c342469 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Mon, 2 Nov 2026 21:57:45 +0000 Subject: [PATCH 234/278] Document that volume exports need READ VOLUME too The export preflight lists the destination before writing. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-databricks/README.md | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/packages/kumo-databricks/README.md b/packages/kumo-databricks/README.md index 2874ff510..8de659f9d 100644 --- a/packages/kumo-databricks/README.md +++ b/packages/kumo-databricks/README.md @@ -91,10 +91,11 @@ parts) is stored as one object up to 4 GiB, so Spark and other tools read it directly. A larger one streams into `.chunks//00000`, `00001`, ... (4 GiB each, under the Files API's 5 GiB upload cap) with an empty `_SUCCESS` written last, and Kumo reads the newest committed generation back as -one file. A file written without a known size stays one object up to 256 MiB, +one file. A file written without a known size stays one object up to 4 GiB, staged on local disk under the engine scratch directory, and streams into parts above that. The identity needs `USE CATALOG`, `USE SCHEMA`, and `READ VOLUME` -(loads) and `WRITE VOLUME` (exports) on the volume; no warehouse is used. +(loads and exports) and `WRITE VOLUME` (exports) on the volume; no warehouse +is used. ## Model registry From c55e04a86e9dcedf98e61c7ee7b9e8575de11b18 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Mon, 2 Nov 2026 22:15:08 +0000 Subject: [PATCH 235/278] Name a fit's datasets by what they are and label each with its role Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/docs/mlops.md | 4 +- packages/kumo-engine/kumo/engine/_manifest.py | 23 +++++--- .../kumo/engine/integrations/mlflow.py | 4 +- .../kumo/engine/trainer/trainer.py | 3 +- packages/kumo-engine/test/test_artifacts.py | 53 ++++++++++++++----- .../kumo-engine/test/test_data_summaries.py | 6 +-- .../kumo-engine/test/test_mlflow_tracking.py | 2 +- .../test/test_tracking_contract.py | 4 +- 8 files changed, 69 insertions(+), 30 deletions(-) diff --git a/packages/kumo-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index 5e51b1224..d07544f40 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -302,7 +302,7 @@ A text column `infer_metadata()` reads as categorical (short values from a small ### Custom trackers -Other systems plug in through `kumo.engine.tracking.Tracking`, 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 attempt after each of its events, and `finish(record, error)` with the final record (`None` when tracking stopped early). Optional members: `prepare()`, `tags`, `check(registry)` (when `kumo.MLOps` is built), `log_model(bundle_dir, manifest, snapshot, registry)` (logs and registers, returning the version) or `connector` (a connector whose model registry `registry` registers in), `load_model(location)` and `register(result, location)` (for loading and `result.export(cfg)`), `platform`/`run_id`/`run_url` (recorded as `result.tracking`). A tracker that raises is not called again for that fit and its run ends failed. Source tables arrive in `finish` as `Dataset` summaries in `record.summaries`, each with its `name`, `source` (`connector`, `location`, `version`, `version_kind`), `schema`, `profile` and `context`; a tracker's `tags` are added except `kumo.*` and param names. Skip summary types your tracker does not handle; new ones may be added. +Other systems plug in through `kumo.engine.tracking.Tracking`, 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 attempt after each of its events, and `finish(record, error)` with the final record (`None` when tracking stopped early). Optional members: `prepare()`, `tags`, `check(registry)` (when `kumo.MLOps` is built), `log_model(bundle_dir, manifest, snapshot, registry)` (logs and registers, returning the version) or `connector` (a connector whose model registry `registry` registers in), `load_model(location)` and `register(result, location)` (for loading and `result.export(cfg)`), `platform`/`run_id`/`run_url` (recorded as `result.tracking`). A tracker that raises is not called again for that fit and its run ends failed. The fit's data arrive in `finish` as `Dataset` summaries in `record.summaries`, each with its `name`, `source` (`connector`, `location`, `version`, `version_kind`), `schema`, `profile` and `context` (its role); a tracker's `tags` are added except `kumo.*` and param names. Skip summary types your tracker does not handle; new ones may be added. ### Progress events @@ -313,4 +313,4 @@ Other systems plug in through `kumo.engine.tracking.Tracking`, which receives re - Tracker calls run on a background thread; `fit` waits at most 30 s for tracking to finish, and up to an hour for model logging and registration (600 s MLflow HTTP timeout during the upload). During a fit MLflow HTTP calls use 2 retries and a 30 s timeout unless `MLFLOW_HTTP_REQUEST_MAX_RETRIES` / `MLFLOW_HTTP_REQUEST_TIMEOUT` are set. Only global rank 0 tracks. - MLflow rejects `@` in metric names: the trainer's `test_map@10` is logged as `test/map_at_10`. Params MLflow rejects are counted in `kumo.dropped_params`. - Snowflake keeps 200 metrics and 1000 params per run; the rest are counted in `kumo.dropped_metrics` / `kumo.dropped_params`. Snowflake has no tags (they are params of the same name), no child runs and no dataset inputs. A taken run name gets `_2`, `_3` and so on; an active run elsewhere in the process is never resumed. -- The fit's source tables are MLflow run inputs: Databricks tables as Delta tables at the snapshot's version (so Unity Catalog links them to the registered model), file sources by URI. +- A fit's data are MLflow run inputs, each with its role as the context: the training table (`training examples`), the graph snapshot (`graph context`) and the snapshot's tables (`source table`). Databricks tables are Delta tables at the snapshot's version (so Unity Catalog links them to the registered model), files are URIs. diff --git a/packages/kumo-engine/kumo/engine/_manifest.py b/packages/kumo-engine/kumo/engine/_manifest.py index d9985f78c..0ff6182d7 100644 --- a/packages/kumo-engine/kumo/engine/_manifest.py +++ b/packages/kumo-engine/kumo/engine/_manifest.py @@ -172,13 +172,18 @@ class Manifest(Record): EMBED_OMITS = () IS_DATASET: ClassVar[bool] = False URI_TAG: ClassVar[Optional[str]] = None + NOUN: ClassVar[Optional[str]] = None + + @property + def title(self) -> str: + short = self.created_at and ( + f"{self.created_at[:10]}-{re.split('[@-]', self.id)[-1][:6]}") + return f"{self.NOUN or self.kind.replace('_', ' ')} {short or self.id}" def as_artifact(self, key: str) -> Artifact: return Artifact(self.kind, self.id, self.location, self.IS_DATASET, self.URI_TAG, version=self.id, version_kind='exact', - label=self.created_at and - f"{self.kind.replace('_', '-')}-{self.created_at[:10]}-" - f"{re.split('[@-]', self.id)[-1][:6]}") + label=self.title) def artifacts(manifest: Manifest) -> List[Artifact]: @@ -212,7 +217,11 @@ def _walk(value: Any, key: str, root: Manifest) -> Iterator[Artifact]: Details = Mapping[str, Tuple[Mapping[str, Any], Mapping[str, Any]]] -def artifact_summaries(manifest: Manifest, context: str, +FIT_ROLES = {'snapshot': 'graph context', 'training_table': + 'training examples', 'source_table': 'source table'} + + +def artifact_summaries(manifest: Manifest, roles: Mapping[str, str], details: Optional[Details] = None) -> List[Summary]: r"""Location tags, the locations table and datasets of a manifest.""" located = [item for item in artifacts(manifest) if item.location] @@ -235,8 +244,9 @@ def artifact_summaries(manifest: Manifest, context: str, 'location': location, 'version': item.version, 'version_kind': item.version_kind} schema, profile = (details or {}).get(item.id, ({}, {})) - summaries.append(Dataset(item.label or item.id, source, schema, - profile, item.context or context)) + summaries.append(Dataset( + item.label or item.id, source, schema, profile, item.context + or roles.get(item.kind, item.kind.replace('_', ' ')))) return summaries @@ -308,6 +318,7 @@ class SnapshotManifest(Manifest): kind: str = 'snapshot' IS_DATASET: ClassVar[bool] = True + NOUN: ClassVar[Optional[str]] = 'graph snapshot' parents: Optional[SnapshotParents] = None record_sources: Optional[bool] = omit_none() diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index 6fc8be3e6..161b640ae 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -729,7 +729,7 @@ def log_model( none. target: as in `MlflowTracking`. """ - from kumo.engine._manifest import artifact_summaries + from kumo.engine._manifest import FIT_ROLES, artifact_summaries from kumo.engine._mlops import snapshot_uri from kumo.engine.trainer._provenance import bundle_manifest from kumo.engine.trainer._run_spec import registry_metrics @@ -754,7 +754,7 @@ def log_model( tracking._experiment_for(client)).info.run_id experiment_id = client.get_run(run_id).info.experiment_id summaries = [] if result.manifest is None else artifact_summaries( - result.manifest, 'training') + result.manifest, FIT_ROLES) try: tracking._log_inputs( run_id, [s for s in summaries if isinstance(s, Dataset)]) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index d9d394a58..2ff73cb9f 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -75,6 +75,7 @@ ) from kumo.engine._events import Event, EventEmitter, query_sha256 from kumo.engine._manifest import ( + FIT_ROLES, Dependency, ModelManifest, TrackingRef, @@ -849,7 +850,7 @@ def fit( *snapshot_summaries(Path(snapshot.index_path), snapshot.graph_config), *([] if result.manifest is None else artifact_summaries( - result.manifest, 'training', + result.manifest, FIT_ROLES, {**table_details(Path(snapshot.index_path), snapshot.id), **training_table_details(train_table)})), diff --git a/packages/kumo-engine/test/test_artifacts.py b/packages/kumo-engine/test/test_artifacts.py index ed5fe584d..bad38920f 100644 --- a/packages/kumo-engine/test/test_artifacts.py +++ b/packages/kumo-engine/test/test_artifacts.py @@ -11,6 +11,7 @@ import kumo.engine as kumo from kumo.engine._deps import Dependencies, FileAsset from kumo.engine._manifest import ( + FIT_ROLES, Artifact, Dependency, Manifest, @@ -21,6 +22,7 @@ SnapshotParents, SourceTableEntry, TrainingTableManifest, + TrainingTableParents, artifact_summaries, artifacts, read_manifest, @@ -61,15 +63,17 @@ def test_a_new_kind_names_its_artifacts_without_other_changes(): assert artifacts(report) == [ Artifact('report', 'r', 's3://root/reports/r', True, - version='r', version_kind='exact'), + version='r', version_kind='exact', label='report r'), Artifact('snapshot', 'g@d', 's3://root/snapshots/g/d', True, - version='g@d', version_kind='exact'), + version='g@d', version_kind='exact', + label='graph snapshot g@d'), Artifact('source_table', 'orders', 's3://data/orders', True, connector='S3', version='7', version_kind='exact'), Artifact('model', 'm', 's3://root/models/m', False, - 'kumo.bundle_uri', version='m', version_kind='exact'), + 'kumo.bundle_uri', version='m', version_kind='exact', + label='model m'), Artifact('snapshot', 'p', None, True, version='p', - version_kind='exact'), + version_kind='exact', label='graph snapshot p'), ] @@ -108,7 +112,7 @@ def test_a_new_kind_and_asset_reach_tags_and_datasets(tmp_path): assert dependency.location == f'{root}/assets/dummy/{dependency.sha256}' assert ReportManifest.from_dict(report.to_dict()) == report - assert artifact_summaries(report, 'training', { + assert artifact_summaries(report, FIT_ROLES, { 'orders': ({'ID': {'stype': 'ID', 'dtype': 'int64'}}, {'rows': 3}) }) == [ Tag('kumo.bundle_uri', f'{root}/models/m'), @@ -118,15 +122,38 @@ def test_a_new_kind_and_asset_reach_tags_and_datasets(tmp_path): ('source_table', 'orders', 's3://data/orders'), ('model', 'm', f'{root}/models/m'), ('asset', dependency.sha256, dependency.location)), brief=True), - Dataset('g@d', {'connector': 'LOCAL_FILE', - 'location': f'{root}/snapshots/g/d', - 'version': 'g@d', 'version_kind': 'exact'}, {}, {}, - 'training'), + Dataset('graph snapshot g@d', { + 'connector': 'LOCAL_FILE', 'location': f'{root}/snapshots/g/d', + 'version': 'g@d', 'version_kind': 'exact'}, {}, {}, + 'graph context'), Dataset('orders', {'connector': 'S3', 'location': 's3://data/orders', 'version': '7', 'version_kind': 'exact'}, {'ID': {'stype': 'ID', 'dtype': 'int64'}}, {'rows': 3}, - 'training'), - Dataset('r', {'connector': 'LOCAL_FILE', - 'location': f'{root}/reports/r', 'version': 'r', - 'version_kind': 'exact'}, {}, {}, 'training'), + 'source table'), + Dataset('report r', {'connector': 'LOCAL_FILE', + 'location': f'{root}/reports/r', 'version': 'r', + 'version_kind': 'exact'}, {}, {}, 'report'), + ] + + +def test_a_fits_datasets_are_named_and_labelled_by_role(): + users = SourceTableEntry(connector='S3', location='s3://data/users', + version='3', version_kind='exact') + snapshot = SnapshotManifest( + id='g@5ac245', created_at='2026-11-01T09:00:00+00:00', + location='s3://root/snapshots/g/5ac245', + parents=SnapshotParents(source_tables={'users': users})) + model = ModelManifest(id='m', parents=ModelParents( + training_table=TrainingTableManifest( + id='tt-fadf04', created_at='2026-11-01T10:00:00+00:00', + location='s3://root/training_tables/tt-fadf04', + parents=TrainingTableParents(snapshot=snapshot)))) + + assert [(s.name, s.context, s.source['version']) + for s in artifact_summaries(model, FIT_ROLES) + if isinstance(s, Dataset)] == [ + ('graph snapshot 2026-11-01-5ac245', 'graph context', 'g@5ac245'), + ('training table 2026-11-01-fadf04', 'training examples', + 'tt-fadf04'), + ('users', 'source table', '3'), ] diff --git a/packages/kumo-engine/test/test_data_summaries.py b/packages/kumo-engine/test/test_data_summaries.py index a69d52787..0aec16d69 100644 --- a/packages/kumo-engine/test/test_data_summaries.py +++ b/packages/kumo-engine/test/test_data_summaries.py @@ -15,7 +15,7 @@ split_summaries, table_details, ) -from kumo.engine._manifest import artifact_summaries +from kumo.engine._manifest import FIT_ROLES, artifact_summaries from kumo.engine._scratch import get_scratch_dir from kumo.engine.pquery import PredictiveQuery @@ -99,7 +99,7 @@ def test_summaries_match_the_snapshot_and_training_table(scratch): users['AGE'].max()) assert age['mean'] == pytest.approx(users['AGE'].mean()) datasets = {s.name: s for s in artifact_summaries( - snapshot.manifest, 'training', + snapshot.manifest, FIT_ROLES, details=table_details(Path(snapshot.index_path))) if isinstance(s, Dataset)} assert set(datasets) == {'users', 'orders'} @@ -110,7 +110,7 @@ def test_summaries_match_the_snapshot_and_training_table(scratch): assert datasets['orders'].profile['rows'] == 6000 assert datasets['users'].source['location'] == str( (data / 'users.parquet').resolve()) - assert datasets['users'].context == 'training' + assert datasets['users'].context == 'source table' def test_data_warnings_name_each_table_that_breaks_a_check(tmp_path): diff --git a/packages/kumo-engine/test/test_mlflow_tracking.py b/packages/kumo-engine/test/test_mlflow_tracking.py index 9da9bb23a..413ea54ab 100644 --- a/packages/kumo-engine/test/test_mlflow_tracking.py +++ b/packages/kumo-engine/test/test_mlflow_tracking.py @@ -385,7 +385,7 @@ def test_source_tables_are_logged_as_native_inputs( monkeypatch.setattr(DeltaDatasetSource, '_lookup_table_id', lambda self, name: 'table-id') replay(MlflowTracking(log_system_metrics=False), summaries=[ - Dataset(name, source, {}, {}, 'training') + Dataset(name, source, {}, {}, 'source table') for name, source in INPUTS.items()]) run, = [run for run in _runs(mlflow.MlflowClient()) diff --git a/packages/kumo-engine/test/test_tracking_contract.py b/packages/kumo-engine/test/test_tracking_contract.py index 0ef99f5b5..0b24970a8 100644 --- a/packages/kumo-engine/test/test_tracking_contract.py +++ b/packages/kumo-engine/test/test_tracking_contract.py @@ -433,7 +433,7 @@ def test_adapter_records_a_single_trial_fit_epoch_by_epoch_without_trials( Scalar('test/brier', 0.25), Table('graph/col_stats', ('table', 'rows'), (('ents', 400),)), Dataset('ents', INPUTS['ents'], {'ID': {'stype': 'id', 'dtype': 'int64'}}, - {'rows': 400}, 'training'), + {'rows': 400}, 'source table'), ) @@ -463,7 +463,7 @@ def test_adapter_logs_each_summary_type(adapter, mlflow_adapter): assert (dataset.dataset.name, dataset.dataset.digest, dataset.tags[0].value) == ('ents', hashlib.sha256(json.dumps( ['ents', INPUTS['ents'], SUMMARIES[-1].schema], - sort_keys=True).encode()).hexdigest()[:8], 'training') + sort_keys=True).encode()).hexdigest()[:8], 'source table') assert json.loads(dataset.dataset.schema) == { 'mlflow_colspec': [{'name': 'ID', 'type': 'long'}]} assert json.loads(dataset.dataset.profile) == {'rows': 400} From 167384d46417ed40a545a252b7a140ab4e416599 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Mon, 2 Nov 2026 23:59:41 +0000 Subject: [PATCH 236/278] Split snapshot loading into allocate, verify and shared read helpers Co-Authored-By: Claude Opus 5.5 --- .../kumo-engine/kumo/engine/_snapshot_io.py | 113 ++++++++++-------- 1 file changed, 62 insertions(+), 51 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_snapshot_io.py b/packages/kumo-engine/kumo/engine/_snapshot_io.py index b0f70d88d..18ba8bbfb 100644 --- a/packages/kumo-engine/kumo/engine/_snapshot_io.py +++ b/packages/kumo-engine/kumo/engine/_snapshot_io.py @@ -102,7 +102,6 @@ def _transfer(connector: Any, path: str, items: Sequence[Any], return [] stop = threading.Event() with copies(items) as workers: - # After `copies`: Arrow's IO pool size sets the connection pool. client = _Client(connector, path) logger.info('Snapshot %s of %s: %d files, %d bytes, %d workers', action, path, len(items), sum(sizes), workers) @@ -132,7 +131,6 @@ def _retried(client: _Client, work: Callable[..., Any], item: Any, if attempt + 1 == _ATTEMPTS: raise logger.warning('Retrying a snapshot transfer after: %s', error) - # Arrow keeps the credentials a client was built with. client.renew(current) stop.wait(2**attempt) attempt += 1 @@ -152,6 +150,18 @@ def _read_json(filesystem: FileSystem, path: str) -> Optional[Any]: return None +def _read_bytes(filesystem: FileSystem, path: str) -> bytes: + with filesystem.open_input_stream(path) as reader: + return reader.readall() + + +def _listed_sizes(index_manifest: bytes) -> Dict[str, int]: + return { + entry['rel_path']: int(entry['size_bytes']) + for entry in json.loads(index_manifest)['files'] + } + + def _open_output(filesystem: FileSystem, path: str, size: int) -> Any: from kumo_core.object_store_sink import open_sized @@ -191,7 +201,6 @@ def export_location(snapshot: Any, path: str) -> str: def store(snapshot: Any, settings: Any) -> None: - r"""Exports ``snapshot`` under the root unless it is there.""" try: connector = settings.root.connector path = settings.root.location('snapshots') @@ -223,10 +232,7 @@ def export(snapshot: Any, connector: Any, path: str) -> str: 'reused it', location, snapshot.id) return location index_manifest = (source / _INDEX_MANIFEST).read_bytes() - sizes = { - entry['rel_path']: int(entry['size_bytes']) - for entry in json.loads(index_manifest)['files'] - } + sizes = _listed_sizes(index_manifest) generation = uuid.uuid4().hex filesystem, root = _location(connector, location) base = f'{root}/{generation}' @@ -283,15 +289,14 @@ def _holds_copy(connector: Any, location: str, snapshot_id: str, return False filesystem, _ = _location(connector, location) remote = file_sizes(filesystem, root) - listed = json.loads((index / _INDEX_MANIFEST).read_bytes())['files'] + listed = _listed_sizes((index / _INDEX_MANIFEST).read_bytes()) if (_DONE not in remote or f'index/{_INDEX_MANIFEST}' not in remote) or any( - remote.get(f"index/{entry['rel_path']}") != int( - entry['size_bytes']) for entry in listed): + remote.get(f'index/{name}') != size + for name, size in listed.items()): return False - with filesystem.open_input_stream( - f'{root}/{PROVENANCE_FILE}') as reader: - return reader.readall() == snapshot_manifest + return _read_bytes(filesystem, + f'{root}/{PROVENANCE_FILE}') == snapshot_manifest except Exception: return False @@ -414,9 +419,7 @@ def _index_files( from kumo_core.sources.layout import confined filesystem, _ = _location(connector, path) - with filesystem.open_input_stream( - f'{root}/index/{_INDEX_MANIFEST}') as reader: - index_manifest = reader.readall() + index_manifest = _read_bytes(filesystem, f'{root}/index/{_INDEX_MANIFEST}') if _fingerprint(index_manifest) != manifest['index_manifest']: raise _corrupt(path, f'{_INDEX_MANIFEST} is not the exported one') try: @@ -427,9 +430,7 @@ def _index_files( if len({name for name, _ in files}) != len(files): raise _corrupt(path, f'{_INDEX_MANIFEST} lists a file twice') try: - with filesystem.open_input_stream( - f'{root}/{PROVENANCE_FILE}') as reader: - snapshot_manifest = reader.readall() + snapshot_manifest = _read_bytes(filesystem, f'{root}/{PROVENANCE_FILE}') except FileNotFoundError: snapshot_manifest = None return index_manifest, files, snapshot_manifest @@ -437,11 +438,6 @@ def _index_files( def load_into(connector: Any, path: str, root: str, manifest: dict, target: Path) -> None: - from diskgraph import ( - DiskgraphCorruptionError, - snapshot_range_checksums, - verify_snapshot_integrity, - ) from kumo_core.diskgraph.snapshot import ( DiskgraphIngestResult, write_diskgraph_success_marker, @@ -456,20 +452,7 @@ def load_into(connector: Any, path: str, root: str, manifest: dict, discard_destination(str(target)) discard_destination(str(staging)) try: - staging.mkdir(parents=True) - needed = sum(size for _, size in files) - disk = os.statvfs(staging) - if disk.f_bavail * disk.f_frsize < needed: - raise OSError(f'Loading this snapshot needs {needed} free bytes ' - f'in {staging}.') - for name, size in files: - (staging / name).parent.mkdir(parents=True, exist_ok=True) - fd = os.open(staging / name, os.O_WRONLY | os.O_CREAT | os.O_EXCL) - try: - if size: - os.posix_fallocate(fd, 0, size) - finally: - os.close(fd) + _allocate(staging, files) ranges = sorted(((name, offset, min(span, size - offset)) for name, size in files for offset in range(0, size, span)), @@ -482,19 +465,7 @@ def load_into(connector: Any, path: str, root: str, manifest: dict, (staging / _INDEX_MANIFEST).write_bytes(index_manifest) write_diskgraph_success_marker(str(staging), DiskgraphIngestResult.for_index(target)) - try: - verify_snapshot_integrity(str(staging)) - loaded = snapshot_range_checksums(str(staging), span) - except DiskgraphCorruptionError as error: - raise _corrupt(path, str(error)) from error - for name, crcs in sorted(loaded.items()): - expected = manifest['crc32'].get(name) - if crcs != expected: - offset = next((index * span for index, crc in enumerate(crcs) - if not expected or index >= len(expected) - or crc != expected[index]), 0) - raise _corrupt( - path, f'{name} failed its checksum at offset {offset}') + _verify(path, staging, manifest) _check_readable(staging, manifest) (staging / _EXPORT_IDENTITY).write_text(export_identity(manifest)) if snapshot_manifest is not None: @@ -509,6 +480,46 @@ def load_into(connector: Any, path: str, root: str, manifest: dict, raise +def _allocate(staging: Path, files: List[Tuple[str, int]]) -> None: + staging.mkdir(parents=True) + needed = sum(size for _, size in files) + disk = os.statvfs(staging) + if disk.f_bavail * disk.f_frsize < needed: + raise OSError(f'Loading this snapshot needs {needed} free bytes ' + f'in {staging}.') + for name, size in files: + (staging / name).parent.mkdir(parents=True, exist_ok=True) + fd = os.open(staging / name, os.O_WRONLY | os.O_CREAT | os.O_EXCL) + try: + if size: + os.posix_fallocate(fd, 0, size) + finally: + os.close(fd) + + +def _verify(path: str, staging: Path, manifest: dict) -> None: + from diskgraph import ( + DiskgraphCorruptionError, + snapshot_range_checksums, + verify_snapshot_integrity, + ) + + span = manifest['range_bytes'] + try: + verify_snapshot_integrity(str(staging)) + loaded = snapshot_range_checksums(str(staging), span) + except DiskgraphCorruptionError as error: + raise _corrupt(path, str(error)) from error + for name, crcs in sorted(loaded.items()): + expected = manifest['crc32'].get(name) + if crcs != expected: + offset = next((index * span for index, crc in enumerate(crcs) + if not expected or index >= len(expected) + or crc != expected[index]), 0) + raise _corrupt(path, + f'{name} failed its checksum at offset {offset}') + + def _check_readable(index: Path, manifest: dict) -> None: from diskgraph import SamplerEngine From 38091c22f9a07aea53fc7a7d313fa848632c11e9 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 3 Nov 2026 00:01:39 +0000 Subject: [PATCH 237/278] Let an unusable MLflow server's own error say why, not blame credentials Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/integrations/mlflow.py | 4 ++-- packages/kumo-engine/test/test_mlflow_target.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index 161b640ae..315ef6acb 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -174,8 +174,8 @@ def check(self, registry: Optional[ModelRegistryLocation] = None) -> None: except ValueError: raise except Exception as error: - raise RuntimeError(f'MLflow at {uri} is not usable with these ' - f'credentials: {full_error(error)}') from error + raise RuntimeError( + f'MLflow at {uri} is not usable: {full_error(error)}') from error if not missing: return with _limited_http(), self.target.scope(): diff --git a/packages/kumo-engine/test/test_mlflow_target.py b/packages/kumo-engine/test/test_mlflow_target.py index 5f7122bae..1694e14ea 100644 --- a/packages/kumo-engine/test/test_mlflow_target.py +++ b/packages/kumo-engine/test/test_mlflow_target.py @@ -240,8 +240,8 @@ def test_an_unreachable_server_fails_when_mlops_is_built( monkeypatch.setenv('MLFLOW_HTTP_REQUEST_MAX_RETRIES', '0') tracking = MlflowTracking(target='http://127.0.0.1:9', experiment_name='e') - with pytest.raises(RuntimeError, match='MLflow at http://127.0.0.1:9 ' - 'is not usable with these credentials'): + with pytest.raises(RuntimeError, + match='MLflow at http://127.0.0.1:9 is not usable: '): kumo.MLOps(str(tmp_path), tracking=tracking) with pytest.raises(ValueError, match='needs experiment_name'): kumo.MLOps(str(tmp_path), From b7c7d6b82f011bbe6ef80c075d8ff906e933bdd5 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 3 Nov 2026 00:01:53 +0000 Subject: [PATCH 238/278] Share MLflow's metric batching between run and system metrics Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/integrations/mlflow.py | 45 ++++++++----------- 1 file changed, 19 insertions(+), 26 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index 315ef6acb..bc04e534f 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -85,6 +85,14 @@ _SAGEMAKER_MODEL = re.compile(r'[a-zA-Z0-9](-*[a-zA-Z0-9]){0,56}') _UC_FORBIDDEN = re.compile(r'[ /\x00-\x1f\x7f]') _HINTED: set = set() +_BATCH_METRICS = 1000 +_SYSTEM_METRICS_FAILURES = 6 +_COLUMN_TYPES = {'int64': 'long', 'int32': 'integer', 'float64': 'double', + 'float32': 'float', 'bool': 'boolean', 'date': 'datetime', + 'time': 'datetime'} +_Reading = Tuple[Mapping[str, float], int, int] # values, second, time +_FLUSH_SECONDS = 10.0 +_MIN_READINGS = 2 class MlflowTracking(TrackingBase): @@ -221,9 +229,8 @@ def load_model(self, location: ModelRegistryLocation, with self.target.scope(): version = self.target.client().get_model_version_by_alias( name, alias).version - uri = f'models:/{name}/{version}' - trainer = load_model(uri, target=self.target, **kwargs) - return trainer + return load_model(f'models:/{name}/{version}', target=self.target, + **kwargs) def register(self, result: Any, location: ModelRegistryLocation) -> PublishedModel: @@ -262,9 +269,8 @@ def register(self, result: Any, if fit is not None: self._log_metrics(fit.run_id, metrics, model_id=info.model_id) - published = _register(self.target, client, info, location, - pointers, held) - return published + return _register(self.target, client, info, location, pointers, + held) def _logged_model(self, fit: Any) -> Any: try: @@ -452,11 +458,14 @@ def log_metrics(self, metrics: Mapping[str, float], step: int) -> None: def _log_readings(self, client: Any, run_id: str, readings: Sequence['_Reading']) -> None: entities = self._mlflow.entities - metrics = [ + self._log_batches(client, run_id, [ entities.Metric(metric_key(name), float(value), timestamp, step) for values, step, timestamp in readings for name, value in values.items() - ] + ]) + + def _log_batches(self, client: Any, run_id: str, + metrics: Sequence[Any]) -> None: with self.target.scope(): for start in range(0, len(metrics), _BATCH_METRICS): client.log_batch( @@ -568,15 +577,11 @@ def _log_metrics(self, run_id: Optional[str], return entities = self._mlflow.entities timestamp = int(time.time() * 1000) - batch = [ + self._log_batches(self._client, run_id, [ entities.Metric(metric_key(key), float(value), timestamp, step, model_id=model_id) for key, value in metrics.items() - ] - with self.target.scope(): - for start in range(0, len(batch), _BATCH_METRICS): - self._client.log_batch( - run_id, metrics=batch[start:start + _BATCH_METRICS]) + ]) def _digest(name: str, source: Mapping[str, Optional[str]], @@ -590,18 +595,6 @@ def _millis(moment: Optional[datetime]) -> Optional[int]: return None if moment is None else int(moment.timestamp() * 1000) -_BATCH_METRICS = 1000 -_SYSTEM_METRICS_FAILURES = 6 -_COLUMN_TYPES = {'int64': 'long', 'int32': 'integer', 'float64': 'double', - 'float32': 'float', 'bool': 'boolean', 'date': 'datetime', - 'time': 'datetime'} - - -_Reading = Tuple[Mapping[str, float], int, int] # values, second, time -_FLUSH_SECONDS = 10.0 -_MIN_READINGS = 2 - - class _SystemMetrics: def __init__(self, log: Callable[[List[_Reading]], None], interval: Optional[float] = None) -> None: From e754d3c4012929c3ab09ffa1216dd6e2a1ec586b Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 3 Nov 2026 00:02:03 +0000 Subject: [PATCH 239/278] Give Snowflake's carrier-model arguments only the bundle they use Co-Authored-By: Claude Opus 5.5 --- packages/kumo-snowflake/src/kumo_snowflake/registry.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/packages/kumo-snowflake/src/kumo_snowflake/registry.py b/packages/kumo-snowflake/src/kumo_snowflake/registry.py index 9ac4d98c6..32d23e589 100644 --- a/packages/kumo-snowflake/src/kumo_snowflake/registry.py +++ b/packages/kumo-snowflake/src/kumo_snowflake/registry.py @@ -37,7 +37,6 @@ import math import re import shutil -import tempfile from contextlib import contextmanager from pathlib import Path from typing import ( @@ -187,8 +186,7 @@ def _carrier_signature() -> Dict[str, "model_signature.ModelSignature"]: } -def _model_args(bundle: Path, metadata: Mapping[str, str], - connector: SourceConnector, workdir: str) -> Dict[str, Any]: +def _model_args(bundle: Path) -> Dict[str, Any]: r"""``log_model`` arguments for the bundle's carrier model.""" return { 'model': _bundle_model(str(bundle)), @@ -397,9 +395,8 @@ def publish_model( f"has no {_SPEC_FILE}, which an export writes into every one.") record = {**_finite(metrics), **_version_record(metadata)} - with tempfile.TemporaryDirectory() as workdir, \ - self._registry(connector) as registry: - args = _model_args(bundle, metadata, connector, workdir) + with self._registry(connector) as registry: + args = _model_args(bundle) name = location.version or _next_version(registry, location.model_name) for attempt in range(3): From dfa3f768f9abdd3bf489f5f4c52f718b7622fd08 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 3 Nov 2026 00:02:05 +0000 Subject: [PATCH 240/278] Drop run-name helpers nothing in this change calls Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_run_names.py | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_run_names.py b/packages/kumo-engine/kumo/engine/_run_names.py index 875a36d5c..c4c8442ab 100644 --- a/packages/kumo-engine/kumo/engine/_run_names.py +++ b/packages/kumo-engine/kumo/engine/_run_names.py @@ -45,16 +45,6 @@ def fit_name(model: str, started: datetime) -> str: return run_name('fit', f'{model}-{started:%Y%m%d-%H%M%S}') -def is_dated_fit(name: str) -> bool: - r"""Whether a fit run is still named by its start, not a version.""" - return re.search(r'-\d{8}-\d{6}$', name) is not None - - -def registered_fit_name(model_name: str, version: str) -> str: - return run_name('fit', versioned(model_label(model_name, None, None), - version)) - - def trial_name(trial: int, attempt: int, fit: str) -> str: return f'trial-{trial}-{fit}' + (f'-attempt-{attempt}' if attempt > 1 else '') From f49809a15b29cbf1befaf28506d7fec79fb558af Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 3 Nov 2026 00:02:29 +0000 Subject: [PATCH 241/278] Return early from a publish that sets no alias Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/trainer/trainer.py | 37 ++++++++++--------- 1 file changed, 19 insertions(+), 18 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 2ff73cb9f..56a1b4bec 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -534,7 +534,7 @@ def _publish(bundle_dir: Path, bundle_id: Optional[str], metadata.update(registry_pointers(bundle_id)) location = export_config.location alias = location.alias if location else None - if location and alias: + if alias: location = replace(location, alias=None) published = export_config.connector.sink_backend.publish_model( export_config.connector, @@ -544,23 +544,24 @@ def _publish(bundle_dir: Path, bundle_id: Optional[str], metadata=metadata, metrics=metrics, ) - if location and alias: - try: - aliasing = export_config.connector.sink_backend - if not isinstance(aliasing, ModelAliasingSink): - raise NotImplementedError( - f'{type(export_config.connector).__name__} cannot set ' - f'model aliases') - aliasing.set_model_alias( - cast(SourceConnector, export_config.connector), - replace(location, version=published.version), alias) - except Exception as error: - if warn: - warn_failure( - f"Published {published.name} version " - f"{published.version} without alias {alias}: " - f"{full_error(error)}", error, 4) - return published, error + if not alias: + return published, None + try: + aliasing = export_config.connector.sink_backend + if not isinstance(aliasing, ModelAliasingSink): + raise NotImplementedError( + f'{type(export_config.connector).__name__} cannot set ' + f'model aliases') + aliasing.set_model_alias( + cast(SourceConnector, export_config.connector), + replace(location, version=published.version), alias) + except Exception as error: + if warn: + warn_failure( + f"Published {published.name} version " + f"{published.version} without alias {alias}: " + f"{full_error(error)}", error, 4) + return published, error return published, None From 7946f659e7d47af8831049665517f16fc3b399d4 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 3 Nov 2026 00:03:00 +0000 Subject: [PATCH 242/278] Say once in the docs that a held bundle reuses its registered version Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/docs/mlops.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/kumo-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index d07544f40..5520328a8 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -144,8 +144,8 @@ table = kumo.TrainingTable.load_from_connector(cfg, path=f"training_tables/{tabl - `result.registered_model` is the version `MLOps(registry=)` registered; `result.published_model` the one `fit(export_config=...)` published. - `result.tracking` is the tracker's run: `TrackingRef(platform="databricks", run_id="...", url="https://.../ml/experiments/.../runs/...")`. The model's `_provenance.json` records the same. - `fit(scoring_defaults=kumo.ScoringDefaults(...))` stores prediction defaults in the bundle before it is stored and registered: `binary_classification_threshold` (binary and multilabel classification) and `num_classes_to_return` (multiclass, multilabel, link prediction). `predict` and the MLflow pyfunc apply them when a call omits them; explicit arguments and a query's `TOP k` win. A binary model always returns `True_PROB`; with no threshold passed or stored, its `TARGET_PRED` is null. `result.with_scoring_defaults(...)` retunes a fitted model as a new bundle whose manifest names the original in `parents.derived_from`. -- `result.export(cfg)` (or `trainer.export(cfg)` for a loaded model) registers a model at `cfg.registry` through its tracker. On MLflow and Unity Catalog, registering a bundle the registry already holds under that name returns its version instead of adding one (found by its bundle id tag); the Snowflake Model Registry adds a new version each time. -- If registration fails, `fit` warns with the full error and returns with `registered_model=None` (with `on_error="raise"`, it raises), and the fit run ends failed with `kumo.error`. On MLflow and Unity Catalog, a bundle the registry already holds under the name returns that version (logged at info) and the alias moves to it. +- `result.export(cfg)` (or `trainer.export(cfg)` for a loaded model) registers a model at `cfg.registry` through its tracker. On MLflow and Unity Catalog, registering a bundle the registry already holds under that name returns its version instead of adding one (found by its bundle id tag, logged at info) and moves the alias to it; the Snowflake Model Registry adds a new version each time. +- If registration fails, `fit` warns with the full error and returns with `registered_model=None` (with `on_error="raise"`, it raises), and the fit run ends failed with `kumo.error`. **Model names** follow the registry's rules, checked when `kumo.MLOps` is built: From 05274762db6c36e0047a9d53f7c67f016a09d34d Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 3 Nov 2026 00:03:11 +0000 Subject: [PATCH 243/278] Inline the fit's duration and wrap its registration call Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/trainer/trainer.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 56a1b4bec..f34b1af54 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -833,14 +833,14 @@ def fit( }) session.finish(None, error) raise - seconds = time.monotonic() - fit_start events.emit('fit_ended', metrics={ - **result.metrics, 'seconds': seconds + **result.metrics, 'seconds': time.monotonic() - fit_start }, params={'status': 'succeeded'}) registered, error = _store_and_register( result.bundle_dir, bundle_manifest(result.bundle_dir, result.bundle_id), settings, - session=session, snapshot=snapshot, metrics=registry_metrics(result), warn=warn, + session=session, snapshot=snapshot, + metrics=registry_metrics(result), warn=warn, stored=(deps.errors or [None])[0]) if registered is not None: result = replace(result, registered_model=registered) From 2e54fec628338aa036a47a9e449e5d408e637c1c Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 3 Nov 2026 00:03:38 +0000 Subject: [PATCH 244/278] Inline asset helpers that have one caller Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_deps.py | 25 +++++++---------------- 1 file changed, 7 insertions(+), 18 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_deps.py b/packages/kumo-engine/kumo/engine/_deps.py index 7855af7c6..03230de81 100644 --- a/packages/kumo-engine/kumo/engine/_deps.py +++ b/packages/kumo-engine/kumo/engine/_deps.py @@ -248,23 +248,16 @@ def reads(connector: Any, location: str) -> bool: and is_file_destination(connector, location)) -def stored_roots(manifest: Optional[ModelManifest], - accepts: Callable[[str], bool]) -> List[str]: - r"""The asset roots of ``manifest``'s stored locations ``accepts``.""" - return sorted({ - record.location.rsplit('/', 2)[0] - for record in Dependencies.of(manifest).records - if record.location and accepts(record.location) - }) - - def connector_roots(manifest: Optional[ModelManifest], connector: Any) -> List[Tuple[Any, str]]: r"""The asset roots of ``manifest`` that ``connector`` reads.""" if connector is None: return [] - return [(connector, root) for root in stored_roots( - manifest, lambda location: reads(connector, location))] + return [(connector, root) for root in sorted({ + record.location.rsplit('/', 2)[0] + for record in Dependencies.of(manifest).records + if record.location and reads(connector, record.location) + })] @contextmanager @@ -278,11 +271,6 @@ def assets_through(connector: Any, bundle: Union[str, Path], yield -def asset_cache() -> Path: - base = os.environ.get('XDG_CACHE_HOME') or Path.home() / '.cache' - return Path(base) / 'kumo' / 'assets' - - def _store(asset: Asset, sha256: str, settings: Any) -> Optional[str]: from kumo_core.object_store_sink import _prepared @@ -335,7 +323,8 @@ def _fetch(connector: Any, path: str, sha256: str) -> Optional[Path]: return None if isinstance(filesystem, LocalFileSystem): return Path(root) - target = asset_cache() / sha256 + cache = os.environ.get('XDG_CACHE_HOME') or Path.home() / '.cache' + target = Path(cache) / 'kumo' / 'assets' / sha256 if (target / PROVENANCE_FILE).is_file(): return target target.parent.mkdir(parents=True, exist_ok=True) From 3f863672d1ee17eadb26199a2057d8348f2635a5 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 3 Nov 2026 00:00:07 +0000 Subject: [PATCH 245/278] Split a source table's change status into its current source and place Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/graph.py | 53 ++++++++++++----------- 1 file changed, 27 insertions(+), 26 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index 3f1b40c3c..b85d9200c 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -338,41 +338,42 @@ def _snapshot_id(graph_key: str, data_version: str) -> str: return f'{graph_key}@{data_version[:16]}' -_COMPARABLE = ('exact', 'timestamp', 'listing') +_COMPARABLE = frozenset({'exact', 'timestamp', 'listing'}) -def _source_change(table: Table, recorded: Optional[SourceTableEntry], - record_sources: bool) -> Dict[str, Any]: +def _current_source(table: Table, + record_sources: bool) -> Optional[SourceTableEntry]: from kumoapi.source_table import CustomSQLSourceTable - before = None if recorded is None else recorded.version - current, same_source = None, False - comparable = (recorded is not None - and recorded.version_kind in _COMPARABLE) + if isinstance(table.definition.source_table, CustomSQLSourceTable): + return None + connector = table.connector try: - source = table.definition.source_table - if not isinstance(source, CustomSQLSourceTable): - connector = table.connector - version = connector.source_backend.get_data_version( - connector, table.table_name) - now = source_table(connector, table.table_name, None, version, - record_sources=record_sources) - current = now.version - comparable = comparable and now.version_kind in _COMPARABLE - same_source = recorded is not None and ( - recorded.connector, recorded.location, - recorded.location_fingerprint) == (now.connector, now.location, - now.location_fingerprint) + version = connector.source_backend.get_data_version( + connector, table.table_name) except Exception: - current = None - if not (same_source and comparable and before is not None - and current is not None): - status = 'unknown' - else: - status = 'unchanged' if before == current else 'changed' + return None + return source_table(connector, table.table_name, None, version, + record_sources=record_sources) + + +def _source_change(table: Table, recorded: Optional[SourceTableEntry], + record_sources: bool) -> Dict[str, Any]: + now = _current_source(table, record_sources) + before = None if recorded is None else recorded.version + current = None if now is None else now.version + known = (None not in (recorded, now, before, current) + and {recorded.version_kind, now.version_kind} <= _COMPARABLE + and _place(recorded) == _place(now)) + status = ('unknown' if not known else + 'unchanged' if before == current else 'changed') return {'status': status, 'recorded': before, 'current': current} +def _place(entry: SourceTableEntry) -> tuple: + return entry.connector, entry.location, entry.location_fingerprint + + @dataclass(frozen=True) class GraphSnapshot: r"""The materialized artifact produced by `Graph.snapshot()`: an on-disk From c1f093ed73b7f81d482c3c74f908935e2a9ed1a4 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 3 Nov 2026 00:00:42 +0000 Subject: [PATCH 246/278] Drop docstrings that restate private helpers' names Co-Authored-By: Claude Opus 5.5 --- packages/kumo-bigquery/src/kumo_bigquery/session.py | 1 - packages/kumo-core/src/kumo_core/artifacts.py | 2 -- packages/kumo-core/src/kumo_core/io/filesystem.py | 1 - packages/kumo-core/src/kumo_core/object_store_sink.py | 3 --- .../kumo-databricks/src/kumo_databricks/session.py | 1 - packages/kumo-engine/kumo/engine/_manifest.py | 11 ----------- packages/kumo-engine/kumo/engine/_mlops.py | 7 ------- packages/kumo-engine/kumo/engine/_provenance.py | 2 -- .../kumo-engine/kumo/engine/connector/__init__.py | 3 +-- packages/kumo-engine/kumo/engine/graph.py | 1 - .../kumo-engine/kumo/engine/trainer/_provenance.py | 4 ---- packages/kumo-engine/kumo/engine/trainer/trainer.py | 4 ---- 12 files changed, 1 insertion(+), 39 deletions(-) diff --git a/packages/kumo-bigquery/src/kumo_bigquery/session.py b/packages/kumo-bigquery/src/kumo_bigquery/session.py index 9e5283ab6..62ee6be6c 100644 --- a/packages/kumo-bigquery/src/kumo_bigquery/session.py +++ b/packages/kumo-bigquery/src/kumo_bigquery/session.py @@ -280,7 +280,6 @@ def _credentials(self, connector: SourceConnector) -> Credentials: def table_location(self, connector: SourceConnector, table_name: str) -> str: - """Fully qualified name of ``table_name`` in the connector's target.""" return _name(connector, table_name) def _client(self, connector: SourceConnector) -> Any: diff --git a/packages/kumo-core/src/kumo_core/artifacts.py b/packages/kumo-core/src/kumo_core/artifacts.py index f3b984386..134463d8d 100644 --- a/packages/kumo-core/src/kumo_core/artifacts.py +++ b/packages/kumo-core/src/kumo_core/artifacts.py @@ -345,7 +345,6 @@ def verify_model_bundle(bundle: Path) -> str: def bundle_id_from_manifest(bundle: Path) -> str: - """The bundle id from the manifest digests.""" manifest_path = Path(bundle) / VitalArtifactPublicationPolicy.MANIFEST_NAME manifest = json.loads(manifest_path.read_text(encoding='utf-8')) return _bundle_id((str(entry['relative_path']), str(entry['sha256'])) @@ -360,7 +359,6 @@ def _bundle_id(pairs: Iterable[tuple[str, str]]) -> str: def add_bundle_file(bundle_dir: Path, role: str, relative_path: str) -> None: - """List a written bundle file in the publication manifest.""" bundle_dir = Path(bundle_dir) path = bundle_dir / relative_path manifest_path = bundle_dir / VitalArtifactPublicationPolicy.MANIFEST_NAME diff --git a/packages/kumo-core/src/kumo_core/io/filesystem.py b/packages/kumo-core/src/kumo_core/io/filesystem.py index c8c227148..54f7b2256 100644 --- a/packages/kumo-core/src/kumo_core/io/filesystem.py +++ b/packages/kumo-core/src/kumo_core/io/filesystem.py @@ -97,7 +97,6 @@ def fsync_directory(directory: str | Path) -> None: def atomic_write_bytes(path: str | Path, payload: bytes) -> None: - """Replace a local file so readers see the old or the new bytes.""" path = Path(path) temporary: Path | None = None try: diff --git a/packages/kumo-core/src/kumo_core/object_store_sink.py b/packages/kumo-core/src/kumo_core/object_store_sink.py index e416cc975..b21e40168 100644 --- a/packages/kumo-core/src/kumo_core/object_store_sink.py +++ b/packages/kumo-core/src/kumo_core/object_store_sink.py @@ -67,7 +67,6 @@ def upload_bundle(filesystem: FileSystem, destination: str, def holds_bundle(filesystem: FileSystem, destination: str, bundle_dir: str) -> bool: - """Whether ``destination`` holds the same files, sizes and provenance.""" from kumo_core.artifacts import PROVENANCE_FILE root = Path(bundle_dir) @@ -84,7 +83,6 @@ def holds_bundle(filesystem: FileSystem, destination: str, def file_sizes(filesystem: FileSystem, directory: str) -> dict[str, int]: - """The size of each file under ``directory``, by relative path.""" from pyarrow.fs import FileSelector, FileType return { @@ -382,7 +380,6 @@ def _prepared(filesystem: FileSystem, directory: str) -> None: def holds_export(filesystem: FileSystem, path: str, export_id: str) -> bool: - """Whether the directory ``path`` holds a complete export ``export_id``.""" from kumo_core.artifact_export import EXPORT_MARKER_NAME from pyarrow.fs import FileType diff --git a/packages/kumo-databricks/src/kumo_databricks/session.py b/packages/kumo-databricks/src/kumo_databricks/session.py index b274cd234..579f648e0 100644 --- a/packages/kumo-databricks/src/kumo_databricks/session.py +++ b/packages/kumo-databricks/src/kumo_databricks/session.py @@ -196,7 +196,6 @@ def _principal(self, connector: SourceConnector) -> Principal: def table_location(self, connector: SourceConnector, table_name: str) -> str: - """Fully qualified name of ``table_name`` in the connector's target.""" return self._qualified(connector, self._bare(connector, table_name)) @staticmethod diff --git a/packages/kumo-engine/kumo/engine/_manifest.py b/packages/kumo-engine/kumo/engine/_manifest.py index 2fb4d7d05..a8df4c800 100644 --- a/packages/kumo-engine/kumo/engine/_manifest.py +++ b/packages/kumo-engine/kumo/engine/_manifest.py @@ -1,8 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Shared JSON manifest format for engine artifacts.""" - import json import types import warnings @@ -36,14 +34,11 @@ def omit_none() -> Any: - """A field left out of the JSON while it is None.""" return field(default=None, metadata={_OMIT_NONE: True}) @dataclass(frozen=True) class Artifact: - """An artifact a manifest names.""" - kind: str id: str location: Optional[str] @@ -155,8 +150,6 @@ def as_artifact(self, key: str) -> Artifact: @dataclass(frozen=True, kw_only=True) class Manifest(Record): - """An artifact's record.""" - kind: str = '' schema_version: int = SCHEMA_VERSION id: str @@ -175,7 +168,6 @@ def as_artifact(self, key: str) -> Artifact: def artifacts(manifest: Manifest) -> List[Artifact]: - """``manifest`` and the artifacts it names, once each.""" found: Dict[Any, Artifact] = {} for artifact in _walk(manifest, manifest.kind, manifest): found.setdefault((artifact.kind, artifact.id), artifact) @@ -201,7 +193,6 @@ def _walk(value: Any, key: str, root: Manifest) -> Iterator[Artifact]: def stamp() -> Dict[str, Any]: - """Envelope fields for a manifest written now.""" from kumo.engine._provenance import environment return dict( @@ -225,12 +216,10 @@ def _kinds() -> Dict[str, Type['Manifest']]: def stub(kind: str, id: str) -> 'Manifest': - """A manifest knowing only ``kind`` and ``id``.""" return _kinds()[kind](id=id, reconstructed=True) def embedded(manifest: M) -> M: - """``manifest`` minus the fields a child's copy omits.""" return replace(manifest, **{ name: None for name in manifest.EMBED_OMITS diff --git a/packages/kumo-engine/kumo/engine/_mlops.py b/packages/kumo-engine/kumo/engine/_mlops.py index d1e1f14cc..9d551dd3a 100644 --- a/packages/kumo-engine/kumo/engine/_mlops.py +++ b/packages/kumo-engine/kumo/engine/_mlops.py @@ -1,8 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -r"""The MLOps switch: :class:`MLOps` settings and the :func:`mlops` block.""" - import contextlib import contextvars import os @@ -164,13 +162,11 @@ def mlops(settings: MLOps) -> Iterator[MLOps]: def resolve(settings: Optional[MLOps] = None) -> Optional[MLOps]: - r"""The ``mlops=`` argument, else the enclosing block's.""" return settings if settings is not None else _CURRENT.get() def raise_on_error(settings: Optional[MLOps], error: Optional[BaseException], call: str, stacklevel: int = 2) -> None: - r"""Raises or warns about ``error``, as ``settings.on_error`` says.""" if error is None or settings is None: return message = f'{call} finished, but MLOps failed: {full_error(error)}' @@ -180,13 +176,11 @@ def raise_on_error(settings: Optional[MLOps], error: Optional[BaseException], def full_error(error: BaseException) -> str: - r"""``Type: message``, not shortened.""" return f'{type(error).__name__}: {error}' def warn_failure(message: str, error: BaseException, stacklevel: int = 2, also: Optional[BaseException] = None) -> None: - r"""Warns ``message`` once per ``error`` (and ``also``).""" if getattr(error, '_kumo_warned', False): return warnings.warn(message, UserWarning, stacklevel=stacklevel + 1) @@ -198,7 +192,6 @@ def warn_failure(message: str, error: BaseException, stacklevel: int = 2, def on_root(source: Any, path: str) -> Tuple[Any, str]: - r"""``path`` under an `MLOps` root, else ``source`` as given.""" if isinstance(source, MLOps): return source.root.connector, source.root.location(path.lstrip('/')) return source, path diff --git a/packages/kumo-engine/kumo/engine/_provenance.py b/packages/kumo-engine/kumo/engine/_provenance.py index 2aea72b08..ebb0e7d6e 100644 --- a/packages/kumo-engine/kumo/engine/_provenance.py +++ b/packages/kumo-engine/kumo/engine/_provenance.py @@ -1,8 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Runtime versions and source tables recorded in manifests.""" - import hashlib import platform import subprocess diff --git a/packages/kumo-engine/kumo/engine/connector/__init__.py b/packages/kumo-engine/kumo/engine/connector/__init__.py index 41530ffbd..703ef9a1c 100644 --- a/packages/kumo-engine/kumo/engine/connector/__init__.py +++ b/packages/kumo-engine/kumo/engine/connector/__init__.py @@ -42,8 +42,7 @@ def connector_for_root(uri: str) -> Connector: - r"""A connector rooted at ``uri``: ``s3://``, ``gs://``, ``file://`` or a - local path.""" + r"""A connector rooted at ``uri``.""" from urllib.parse import urlparse from urllib.request import url2pathname diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index b85d9200c..e2e684210 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -226,7 +226,6 @@ def _schema_tables(config: Mapping[str, Any]) -> Dict[str, Dict[str, Any]]: def schema_changes(before: Mapping[str, Any], after: Mapping[str, Any]) -> List[str]: - r"""Schema differences between two graph configs.""" old, new = _schema_tables(before), _schema_tables(after) changes = [ f"table {name} {'removed' if name in old else 'added'}" diff --git a/packages/kumo-engine/kumo/engine/trainer/_provenance.py b/packages/kumo-engine/kumo/engine/trainer/_provenance.py index 5d8b0fbdb..4132d0196 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_provenance.py +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -1,8 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""A model bundle's ``_provenance.json``; excluded from the bundle id.""" - from dataclasses import dataclass, replace from datetime import datetime, timezone from pathlib import Path @@ -58,7 +56,6 @@ def write(self, bundle_dir: Path, bundle_id: str) -> None: def embedded_model(manifest: ModelManifest) -> ModelManifest: - r"""``manifest`` for embedding, its own model parents cut to stubs.""" copy = embedded(manifest) parents = copy.parents if parents is None: @@ -73,7 +70,6 @@ def cut(model: Optional[ModelManifest]) -> Optional[ModelManifest]: def bundle_manifest(bundle_dir: Path, bundle_id: Optional[str]) -> Optional[ModelManifest]: - r"""The manifest of the bundle at ``bundle_dir``, or a stub.""" manifest = read_provenance(bundle_dir) if manifest is not None: return manifest diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 27ff90ba6..97f8de2c4 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -244,7 +244,6 @@ def _find_holdout(log_dir: str) -> tuple[Optional[Path], Optional[bool]]: def _bundle_holdout(result: TrainingResult) -> None: - r"""Copy the holdout into the bundle.""" from kumo_core.artifacts import add_bundle_file if result.holdout_path is None: @@ -279,7 +278,6 @@ def _listed_holdout(bundle_dir: Path) -> tuple[Optional[Path], Optional[bool]]: def _store_model(bundle_dir: Path, manifest: Optional[ModelManifest], settings: Any) -> None: - r"""Copies the bundle under the root unless it is there.""" from kumo_core.object_store_sink import holds_bundle, upload_bundle from pyarrow.fs import FileSelector @@ -363,7 +361,6 @@ def _publish(bundle_dir: Path, metrics: Mapping[str, float], def _warm_start_plan( public_plan: ModelPlan, warm_start: Optional['Trainer'] ) -> Tuple[CoreModelPlan, ModelPlan, Optional[Dict[str, str]], Optional[str]]: - r"""The model plan with the warm start's inherited, and its artifacts.""" from kumo_core.automl.runtime import inherit_warm_start_model_plan from kumo_core.io.json import load_json @@ -703,7 +700,6 @@ def _train( self, materialized: Any, run_dir: Path, run_id: str, warm_start_artifacts: Optional[Mapping[str, str]] ) -> Tuple[TrainingResult, Optional[str]]: - r"""The trained model and its fit mode.""" from kumo_core.training.orchestration import ( CallableTrainingOptions, ResolvedTrainingInput, From b5c706d02ec0362b9903bb2d0d689036c8e7d7c7 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 3 Nov 2026 00:00:42 +0000 Subject: [PATCH 247/278] Shorten the range-checksum doc comments Co-Authored-By: Claude Opus 5.5 --- .../diskgraph/diskgraph-ingest/src/snapshot_mirror.rs | 8 ++------ packages/diskgraph/diskgraph-pybind/src/lib.rs | 3 +-- 2 files changed, 3 insertions(+), 8 deletions(-) diff --git a/packages/diskgraph/diskgraph-ingest/src/snapshot_mirror.rs b/packages/diskgraph/diskgraph-ingest/src/snapshot_mirror.rs index bfba23cfb..4e90e9704 100644 --- a/packages/diskgraph/diskgraph-ingest/src/snapshot_mirror.rs +++ b/packages/diskgraph/diskgraph-ingest/src/snapshot_mirror.rs @@ -489,8 +489,7 @@ pub(crate) fn validate_manifest_rel_path(rel: &str) -> Result<(), &'static str> Ok(()) } -/// Read and validate `_MANIFEST.json` at `output_dir`: schema version and -/// every rel_path's confinement to the snapshot root. +/// `_MANIFEST.json` at `output_dir`, schema and paths checked. fn read_manifest(output_dir: &std::path::Path) -> std::result::Result { let manifest_path = output_dir.join(MANIFEST_FILENAME); let manifest_bytes = match std::fs::read(&manifest_path) { @@ -592,10 +591,7 @@ pub fn verify_snapshot_integrity(output_dir: &std::path::Path) -> std::result::R Ok(()) } -/// CRC-32 of every `range_bytes` span of each file `_MANIFEST.json` lists, -/// and of `_MANIFEST.json` itself, keyed by rel_path (an empty file has one -/// zero entry). Lets a copy of the snapshot be checked for content as well -/// as size; the manifest format is unchanged. +/// CRC-32 per `range_bytes` span of each listed file and the manifest. pub fn snapshot_range_checksums( output_dir: &std::path::Path, range_bytes: u64, diff --git a/packages/diskgraph/diskgraph-pybind/src/lib.rs b/packages/diskgraph/diskgraph-pybind/src/lib.rs index a47daef4d..caa73c2b2 100644 --- a/packages/diskgraph/diskgraph-pybind/src/lib.rs +++ b/packages/diskgraph/diskgraph-pybind/src/lib.rs @@ -1553,8 +1553,7 @@ fn verify_snapshot_integrity(py: Python<'_>, output_dir: &str) -> PyResult<()> { .map_err(|e| DiskgraphCorruptionError::new_err(e.to_string())) } -/// CRC-32 of every `range_bytes` span of each file in the snapshot's -/// `_MANIFEST.json` and of the manifest itself, keyed by relative path. +/// CRC-32 per `range_bytes` span of each listed file and the manifest. #[pyfunction] fn snapshot_range_checksums( py: Python<'_>, From d2d7e1a041f33b7e01c142143223885fdc158bec Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 3 Nov 2026 00:01:05 +0000 Subject: [PATCH 248/278] Name the local and remote write checks of an artifact root Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_mlops.py | 38 +++++++++++++--------- 1 file changed, 23 insertions(+), 15 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_mlops.py b/packages/kumo-engine/kumo/engine/_mlops.py index 9d551dd3a..515633a79 100644 --- a/packages/kumo-engine/kumo/engine/_mlops.py +++ b/packages/kumo-engine/kumo/engine/_mlops.py @@ -122,22 +122,9 @@ def _require_writable(root: _Root) -> None: filesystem, probe = store try: if filesystem.type_name == 'local': - existing = posixpath.dirname(probe) - while not os.path.exists(existing): - existing = posixpath.dirname(existing) - if not os.access(existing, os.W_OK | os.X_OK): - raise PermissionError(f'{existing} is not writable') - return - try: - filesystem.open_output_stream(probe).close() - except FileNotFoundError: - filesystem.create_dir(posixpath.dirname(probe), recursive=True) - filesystem.open_output_stream(probe).close() - delete = getattr(root.connector.sink_backend, 'delete_file', None) - if delete is None: - filesystem.delete_file(probe) + _require_local_writable(probe) else: - delete(root.connector, probe) + _write_and_delete(filesystem, root.connector, probe) except Exception as error: raise PermissionError( f'kumo.MLOps cannot write under artifact_root {root.uri("")!r} ' @@ -145,6 +132,27 @@ def _require_writable(root: _Root) -> None: f'rights or pick another root.') from error +def _require_local_writable(probe: str) -> None: + existing = posixpath.dirname(probe) + while not os.path.exists(existing): + existing = posixpath.dirname(existing) + if not os.access(existing, os.W_OK | os.X_OK): + raise PermissionError(f'{existing} is not writable') + + +def _write_and_delete(filesystem: Any, connector: Any, probe: str) -> None: + try: + filesystem.open_output_stream(probe).close() + except FileNotFoundError: + filesystem.create_dir(posixpath.dirname(probe), recursive=True) + filesystem.open_output_stream(probe).close() + delete = getattr(connector.sink_backend, 'delete_file', None) + if delete is None: + filesystem.delete_file(probe) + else: + delete(connector, probe) + + _CURRENT: contextvars.ContextVar[Optional[MLOps]] = contextvars.ContextVar( 'kumo_mlops', default=None) From 6bea711770080139631007753e3c6cb937b8eb29 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 3 Nov 2026 00:01:05 +0000 Subject: [PATCH 249/278] Read a training table export's name once when choosing sidecars Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/pquery/training_table.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/pquery/training_table.py b/packages/kumo-engine/kumo/engine/pquery/training_table.py index 72c7b9da4..133319d0d 100644 --- a/packages/kumo-engine/kumo/engine/pquery/training_table.py +++ b/packages/kumo-engine/kumo/engine/pquery/training_table.py @@ -107,9 +107,9 @@ def _export_sidecars(self, connector, output_config) -> Dict[str, Any]: from kumo.engine.connector import is_file_destination - name = output_config.table_name.rstrip('/') - if (not is_file_destination(connector, output_config.table_name) - or layout.file_type(name) == layout.CSV): + name = output_config.table_name + if (not is_file_destination(connector, name) + or layout.file_type(name.rstrip('/')) == layout.CSV): return {} sidecars = self._sidecars() return {'sidecars': sidecars} if sidecars else {} From 24c1431fffea42aa1a728959322c6429813de1a1 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 3 Nov 2026 00:01:40 +0000 Subject: [PATCH 250/278] Share the label query and table generation across training table tests Co-Authored-By: Claude Opus 5.5 --- .../test/test_training_table_export.py | 66 ++++++++----------- 1 file changed, 27 insertions(+), 39 deletions(-) diff --git a/packages/kumo-engine/test/test_training_table_export.py b/packages/kumo-engine/test/test_training_table_export.py index 8914fe598..acbc23e0a 100644 --- a/packages/kumo-engine/test/test_training_table_export.py +++ b/packages/kumo-engine/test/test_training_table_export.py @@ -3,10 +3,12 @@ # SPDX-License-Identifier: Apache-2.0 import json +import os from collections import Counter from dataclasses import replace from functools import partial from types import SimpleNamespace +from unittest.mock import Mock import numpy as np import pandas as pd @@ -25,6 +27,7 @@ SPLITS = {'train': 0, 'val': 1, 'test': 2} FREQ = '__kumo_internal_frequency_column__' +LABEL = 'PREDICT ents.LABEL FOR EACH ents.ID' @pytest.fixture(scope='module') @@ -96,6 +99,12 @@ def _rows(table, columns): for row in zip(*values)) +def _generate(graph, **kwargs): + graph, snapshot, _, _ = graph + pquery = kumo.PredictiveQuery(graph=graph, query=LABEL) + return pquery, pquery.generate_training_table(snapshot, **kwargs) + + def _export(graph, query, name): graph, snapshot, connector, out = graph pquery = kumo.PredictiveQuery(graph=graph, query=query) @@ -107,7 +116,7 @@ def _export(graph, query, name): @pytest.mark.parametrize('name, query', [ ('temporal', 'PREDICT COUNT(events.*, 0, 7, days) > 0 FOR EACH ents.ID'), - ('static', 'PREDICT ents.LABEL FOR EACH ents.ID'), + ('static', LABEL), ]) def test_node_export_roundtrip(graph, name, query): _, snapshot, connector, out = graph @@ -307,14 +316,13 @@ def test_manifest_and_export_sidecar(graph, tmp_path): from kumo_core.training.export import export_train_table _, snapshot, connector, out = graph - query = 'PREDICT ents.LABEL FOR EACH ents.ID' - pquery, table, _ = _export(graph, query, 'lineage') + pquery, table, _ = _export(graph, LABEL, 'lineage') manifest = table.manifest assert manifest.id == table.id and table.id.startswith('tt-') assert table.id != _spec(table)['tt_id'] assert manifest.parents == TrainingTableParents( snapshot=snapshot.manifest) - assert manifest.query == query + assert manifest.query == LABEL assert manifest.extra == {} assert manifest.split_seed == 0 again = pquery.generate_training_table(snapshot) @@ -346,8 +354,7 @@ def test_manifest_and_export_sidecar(graph, tmp_path): def test_custom_table_follows_the_snapshots_record_sources( graph, monkeypatch): graph_, snapshot, connector, _ = graph - pquery, _, _ = _export(graph, 'PREDICT ents.LABEL FOR EACH ents.ID', - 'unlocated') + pquery, _, _ = _export(graph, LABEL, 'unlocated') def location(snapshot): table = kumo.TrainingTable.from_source_table( @@ -368,25 +375,19 @@ def test_manifest_write_failure_keeps_generation_and_export( graph, monkeypatch): from kumo.engine.pquery import training_table - _, _, connector, out = graph - - def fail(*args, **kwargs): - raise OSError('read-only scratch') - - monkeypatch.setattr(training_table, 'write_manifest', fail) + out = graph[3] + monkeypatch.setattr(training_table, 'write_manifest', + Mock(side_effect=OSError('read-only scratch'))) with pytest.warns(UserWarning, match='read-only scratch'): - _, table, exported = _export( - graph, 'PREDICT ents.LABEL FOR EACH ents.ID', 'no_manifest') + _, table, exported = _export(graph, LABEL, 'no_manifest') assert table.manifest is None assert not (out / 'no_manifest' / '_provenance.json').exists() assert exported.num_rows == sum(_seeds(table)[1].values()) def test_unreadable_manifest_still_exports(graph): - graph_, snapshot, connector, out = graph - table = kumo.PredictiveQuery( - graph=graph_, query='PREDICT ents.LABEL FOR EACH ents.ID' - ).generate_training_table(snapshot) + _, _, connector, out = graph + _, table = _generate(graph) manifest = table.staging_dir / '_provenance.json' manifest.chmod(0) try: @@ -399,31 +400,23 @@ def test_unreadable_manifest_still_exports(graph): def test_no_sidecar_for_warehouse_or_csv_exports(graph): - graph_, snapshot, connector, out = graph - table = kumo.PredictiveQuery( - graph=graph_, query='PREDICT ents.LABEL FOR EACH ents.ID' - ).generate_training_table(snapshot) - calls = [] - - def write_files(connector, table_name, files, write_mode, **kwargs): - calls.append(kwargs) - + _, _, connector, out = graph + _, table = _generate(graph) + write_files = Mock() warehouse = SimpleNamespace( source_type='SNOWFLAKE', sink_backend=SimpleNamespace(write_files=write_files)) table.export(TrainingTableOutputConfig(warehouse, 'DB.SCHEMA.TT')) - assert calls == [{'export_id': f'{table.tt_id}:tt-export-v1'}] + assert [call.kwargs for call in write_files.call_args_list] == [{ + 'export_id': f'{table.tt_id}:tt-export-v1'}] table.export(TrainingTableOutputConfig(connector, 'rows.csv/')) assert not (out / 'rows.csv' / '_provenance.json').exists() def test_the_store_spills_at_most_two_parts(graph, tmp_path, monkeypatch): - import os - from kumo_core import object_store_sink from kumo_core.training import export - graph_, snapshot, _, _ = graph stage_parts = export.stage_parts monkeypatch.setattr( export, 'stage_parts', lambda spec: { @@ -438,23 +431,18 @@ def upload(filesystem, source, destination): upload_file(filesystem, source, destination) monkeypatch.setattr(object_store_sink, 'upload_file', upload) - table = kumo.PredictiveQuery( - graph=graph_, query='PREDICT ents.LABEL FOR EACH ents.ID' - ).generate_training_table( - snapshot, mlops=kumo.MLOps(artifact_root=str(tmp_path))) + _, table = _generate(graph, mlops=kumo.MLOps(artifact_root=str(tmp_path))) stored = tmp_path / 'training_tables' / table.id assert len(spilled) == len(list(stored.glob('part-*.parquet'))) >= 8 assert max(spilled) <= 2 def test_a_mlops_block_stores_each_table_once_under_the_root(graph, tmp_path): - graph_, snapshot, _, _ = graph - pquery = kumo.PredictiveQuery(graph=graph_, - query='PREDICT ents.LABEL FOR EACH ents.ID') + snapshot = graph[1] root = tmp_path / 'artifacts' cfg = kumo.MLOps(artifact_root=str(root)) with kumo.mlops(cfg): - table = pquery.generate_training_table(snapshot) + pquery, table = _generate(graph) stored = root / 'training_tables' / table.id parts = {p.name: p.stat().st_mtime_ns for p in stored.iterdir()} pquery.generate_training_table(snapshot) From 9a9057ca7bc54ab14eb0d488fe366e16c2601c2e Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 3 Nov 2026 00:02:11 +0000 Subject: [PATCH 251/278] Export snapshots in tests through one helper Co-Authored-By: Claude Opus 5.5 --- .../kumo-engine/test/test_snapshot_export.py | 80 +++++++++---------- 1 file changed, 40 insertions(+), 40 deletions(-) diff --git a/packages/kumo-engine/test/test_snapshot_export.py b/packages/kumo-engine/test/test_snapshot_export.py index 29aab0f49..fdb3b7f5f 100644 --- a/packages/kumo-engine/test/test_snapshot_export.py +++ b/packages/kumo-engine/test/test_snapshot_export.py @@ -12,12 +12,15 @@ import zlib from contextlib import contextmanager from dataclasses import replace +from functools import partial from pathlib import Path from types import SimpleNamespace from unittest.mock import Mock import pandas as pd import pytest +from kumoapi.graph import Edge +from kumoapi.typing import Stype from pyarrow.fs import LocalFileSystem, SubTreeFileSystem import kumo.engine as kumo @@ -52,8 +55,12 @@ def _index_files(path): } +def _put(snapshot, connector, path='one'): + return snapshot.export(SnapshotOutputConfig(connector, path)) + + def _export(snapshot, exports, tmp_path): - location = snapshot.export(SnapshotOutputConfig(exports, 'one')) + location = _put(snapshot, exports) return location, tmp_path / 'exports' / location @@ -88,10 +95,10 @@ def test_export_load_round_trip_relocates_and_is_reused( monkeypatch.setattr(snapshot_io, '_CHUNK', 1024) monkeypatch.setattr(snapshot_io, '_RANGE_BYTES', 4096) snapshot = graph.snapshot() - location = snapshot.export(SnapshotOutputConfig(exports, 'one/')) + location = _put(snapshot, exports, 'one/') graph_id, data_version = snapshot.id.split('@') assert location == f'one/{graph_id}/{data_version}' - assert snapshot.export(SnapshotOutputConfig(exports, 'one')) == location + assert _put(snapshot, exports) == location exported = tmp_path / 'exports' / location recorded = (snapshot.index_path / '_provenance.json').read_bytes() assert (_generation(exported) / '_provenance.json').read_bytes() == recorded @@ -106,8 +113,7 @@ def test_export_load_round_trip_relocates_and_is_reused( _generation(exported) / 'index') assert (loaded.index_path / '_provenance.json').read_bytes() == recorded - reexported = tmp_path / 'exports' / loaded.export( - SnapshotOutputConfig(exports, 'two')) + reexported = tmp_path / 'exports' / _put(loaded, exports, 'two') assert (_generation(reexported) / '_provenance.json').read_bytes() == recorded @@ -161,13 +167,13 @@ def interrupted(client, path, size): patch.setattr(snapshot_io, '_open_output', interrupted) patch.setattr(snapshot_io, '_ATTEMPTS', 1) with pytest.raises(OSError, match='connection reset'): - rebuilt.export(SnapshotOutputConfig(exports, 'one')) + _put(rebuilt, exports) assert _generation(exported) == first scratch('gpu') assert _index_files(_load(exports, location).index_path) == _index_files( snapshot.index_path) - assert rebuilt.export(SnapshotOutputConfig(exports, 'one')) == location + assert _put(rebuilt, exports) == location assert _generation(exported) != first and not first.exists() @@ -179,12 +185,12 @@ def test_reexporting_an_identical_snapshot_reuses_the_export( monkeypatch.setattr(snapshot_io, '_transfer', Mock(side_effect=AssertionError('uploaded again'))) with caplog.at_level('INFO', logger=snapshot_io.__name__): - assert snapshot.export(SnapshotOutputConfig(exports, 'one')) == location + assert _put(snapshot, exports) == location assert _generation(exported) == first and 'reused it' in caplog.text (first / '_DONE').unlink() with pytest.raises(AssertionError, match='uploaded again'): - snapshot.export(SnapshotOutputConfig(exports, 'one')) + _put(snapshot, exports) def test_concurrent_exports_of_one_snapshot_both_succeed( @@ -195,16 +201,16 @@ def test_concurrent_exports_of_one_snapshot_both_succeed( def racing(*args): if not raced: raced.append(None) - raced[0] = snapshot.export(SnapshotOutputConfig(exports, 'one')) + raced[0] = _put(snapshot, exports) return transfer(*args) monkeypatch.setattr(snapshot_io, '_transfer', racing) - assert raced == [snapshot.export(SnapshotOutputConfig(exports, 'one'))] + assert raced == [_put(snapshot, exports)] assert _load(exports, raced[0]).id == snapshot.id def _rebuilt(snapshot, tmp_path, name): - r"""Same id and sizes, one digit changed in ``col_stats.json``.""" + r"""Same id and sizes, other bytes.""" index = (tmp_path / name / snapshot.index_path.parent.name / snapshot.index_path.name) shutil.copytree(snapshot.index_path, index) @@ -235,14 +241,14 @@ def held(*args): def export_first(): try: - first.export(SnapshotOutputConfig(exports, 'one')) + _put(first, exports) except Exception as error: errors.append(error) worker = threading.Thread(target=export_first) worker.start() assert paused.wait(30) - location = second.export(SnapshotOutputConfig(exports, 'one')) + location = _put(second, exports) resumed.set() worker.join(30) assert not errors @@ -259,13 +265,13 @@ def test_a_load_follows_an_export_that_commits_while_it_downloads( tmp_path, scratch, graph, exports, monkeypatch): first = graph.snapshot() second = _rebuilt(first, tmp_path, 'second') - location = first.export(SnapshotOutputConfig(exports, 'one')) + location = _put(first, exports) load_into, calls = snapshot_io.load_into, [] def overtaken(*args): calls.append(args[2]) if len(calls) == 1: - second.export(SnapshotOutputConfig(exports, 'one')) + _put(second, exports) return load_into(*args) monkeypatch.setattr(snapshot_io, 'load_into', overtaken) @@ -279,8 +285,8 @@ def test_a_cached_load_is_reused_only_for_the_same_export( tmp_path, scratch, graph, exports): first = graph.snapshot() second = _rebuilt(first, tmp_path, 'second') - one = first.export(SnapshotOutputConfig(exports, 'one')) - two = second.export(SnapshotOutputConfig(exports, 'two')) + one = _put(first, exports) + two = _put(second, exports, 'two') assert _index_files(first.index_path) != _index_files(second.index_path) scratch('gpu') @@ -362,7 +368,7 @@ def expiring(client, path, size): return open_output(client, path, size) monkeypatch.setattr(snapshot_io, '_open_output', expiring) - graph.snapshot().export(SnapshotOutputConfig(exports, 'one')) + _put(graph.snapshot(), exports) assert len(clients) == 2 and clients[0] is not clients[1] @@ -373,7 +379,7 @@ def test_reexport_after_redaction_publishes_the_redacted_manifest( first = _generation(exported) redacted = graph.snapshot(record_sources=False) assert redacted.id == full.id - assert redacted.export(SnapshotOutputConfig(exports, 'one')) == location + assert _put(redacted, exports) == location committed = _generation(exported) assert committed != first assert (committed / '_provenance.json').read_bytes() == ( @@ -395,7 +401,7 @@ def test_reexport_replaces_an_export_that_is_not_an_intact_copy( stats.write_bytes(stats.read_bytes()[:3]) (first / 'index/index_metadata.json').unlink() _edit_manifest(exported, lambda m: m.update(snapshot_id='other')) - assert snapshot.export(SnapshotOutputConfig(exports, 'one')) == location + assert _put(snapshot, exports) == location assert not first.exists() scratch('empty') assert _index_files(_load(exports, location).index_path) == _index_files( @@ -519,7 +525,7 @@ def test_snapshot_manifest_travels_unchanged_or_is_reconstructed( def test_warehouse_connectors_are_not_supported(scratch, graph): connector = SimpleNamespace(sink_backend=object()) with pytest.raises(NotImplementedError, match='file connector'): - graph.snapshot().export(SnapshotOutputConfig(connector, 'x')) + _put(graph.snapshot(), connector, 'x') with pytest.raises(NotImplementedError, match='file connector'): _load(connector, 'x') @@ -555,8 +561,7 @@ def __init__(self, *args, **kwargs): monkeypatch.setattr('kumo_gcs.backend.fs.GcsFileSystem', Bucket) snapshot = graph.snapshot() connector = kumo.S3Connector() if scheme == 's3' else kumo.GCSConnector() - location = snapshot.export( - SnapshotOutputConfig(connector, f'{scheme}://bucket/snapshots/')) + location = _put(snapshot, connector, f'{scheme}://bucket/snapshots/') graph_id, data_version = snapshot.id.split('@') assert location == f'{scheme}://bucket/snapshots/{graph_id}/{data_version}' scratch('gpu') @@ -629,11 +634,6 @@ def test_source_changes_compares_directory_listings(tmp_path, scratch): @pytest.mark.parametrize('record_sources', [True, False]) def test_source_changes_notices_a_rebound_source(tmp_path, scratch, graph, dated, record_sources): - import shutil - - from kumoapi.graph import Edge - from kumoapi.typing import Stype - snapshot = graph.snapshot(record_sources=record_sources) shutil.copytree(dated, tmp_path / 'other') connector = kumo.FileConnector(root_dir=str(tmp_path / 'other')) @@ -737,7 +737,7 @@ def _remove(self, path, *, tree): path.rstrip('/'))) connector = SimpleNamespace(sink_backend=sink) snapshot = graph.snapshot() - location = snapshot.export(SnapshotOutputConfig(connector, '@stage/s')) + location = _put(snapshot, connector, '@stage/s') assert any('.chunks/' in name for name in store.objects) scratch('gpu') loaded = _load(connector, location) @@ -753,7 +753,7 @@ def _remove(self, path, *, tree): assert (loaded.index_path / info.path[len(root) + 1:] ).read_bytes() == reader.read() with pytest.raises(NotImplementedError, match='S3, GCS or file connector'): - snapshot.export(SnapshotOutputConfig(connector, 'DB.SCHEMA')) + _put(snapshot, connector, 'DB.SCHEMA') def test_a_mlops_block_stores_each_snapshot_once(tmp_path, scratch, graph, @@ -787,15 +787,15 @@ def test_loading_a_missing_export_names_its_full_path(tmp_path, exports): _load(exports, 'snapshots/graph/version') -@pytest.mark.parametrize('on_error', ['warn', 'raise']) +@pytest.mark.parametrize('on_error, outcome', [ + ('warn', partial(pytest.warns, UserWarning, match='ACCESS_DENIED')), + ('raise', + partial(pytest.raises, RuntimeError, match='Could not store snapshot')), +]) def test_a_failed_store_follows_on_error(tmp_path, scratch, graph, - monkeypatch, on_error): + monkeypatch, on_error, outcome): monkeypatch.setattr(snapshot_io, 'export', Mock(side_effect=OSError('ACCESS_DENIED'))) - cfg = kumo.MLOps(str(tmp_path / 'root'), on_error=on_error) - if on_error == 'raise': - with pytest.raises(RuntimeError, match='Could not store snapshot'): - graph.snapshot(mlops=cfg) - else: - with pytest.warns(UserWarning, match='ACCESS_DENIED'): - graph.snapshot(mlops=cfg) + with outcome(): + graph.snapshot(mlops=kumo.MLOps(str(tmp_path / 'root'), + on_error=on_error)) From 6eac84edc9269442449988b3daca889a9040fa69 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 3 Nov 2026 00:02:29 +0000 Subject: [PATCH 252/278] Default the stub-fit helper's trainer and drop repeated model assertions Co-Authored-By: Claude Opus 5.5 --- .../test/test_bundle_provenance.py | 26 +++++++------------ 1 file changed, 9 insertions(+), 17 deletions(-) diff --git a/packages/kumo-engine/test/test_bundle_provenance.py b/packages/kumo-engine/test/test_bundle_provenance.py index 7b0f6bd28..09f3f5567 100644 --- a/packages/kumo-engine/test/test_bundle_provenance.py +++ b/packages/kumo-engine/test/test_bundle_provenance.py @@ -30,9 +30,7 @@ def test_fit_records_a_self_contained_manifest(fitted): result = trainer.result manifest = result.manifest assert manifest.kind == 'model' - assert re.fullmatch('[0-9a-f]{64}', trainer.result.id) - assert trainer.result.id == result.id == manifest.id - assert trainer.result.manifest == manifest + assert re.fullmatch('[0-9a-f]{64}', result.id) and result.id == manifest.id table = manifest.parents.training_table assert table.kind == 'training_table' and table.id.startswith('tt-') assert table.parents.snapshot == snapshot.manifest @@ -103,7 +101,8 @@ def test_each_fit_returns_a_new_result(fitted): assert second is not first and trainer.result is second -def _finish_stub_fit(bundle, trainer, warm_start=None, settings=None): +def _finish_stub_fit(bundle, trainer=None, warm_start=None, settings=None): + trainer = trainer or kumo.Trainer(ModelPlan()) (bundle / 'model').mkdir(parents=True) (bundle / 'model' / 'config.json').write_text('{}') (bundle / MANIFEST_FILE).write_text(json.dumps({'artifacts': []})) @@ -121,20 +120,16 @@ def _finish_stub_fit(bundle, trainer, warm_start=None, settings=None): def test_missing_parents_are_stubs_and_retrain_chains_stay_flat(tmp_path): - first = _finish_stub_fit(tmp_path / 'one', kumo.Trainer(ModelPlan())) + first = _finish_stub_fit(tmp_path / 'one') assert first.manifest.parents.training_table.to_dict() == { 'kind': 'training_table', 'schema_version': 1, 'id': 'tt-1', 'created_at': None, 'environment': None, 'reconstructed': True, 'parents': None, 'query': None, 'generation_plan': None, 'task_type': None, 'split_seed': None} - second = _finish_stub_fit(tmp_path / 'two', kumo.Trainer(ModelPlan()), - warm_start=first.manifest) - third = _finish_stub_fit(tmp_path / 'three', kumo.Trainer(ModelPlan()), - warm_start=second.manifest) + second = _finish_stub_fit(tmp_path / 'two', warm_start=first.manifest) + third = _finish_stub_fit(tmp_path / 'three', warm_start=second.manifest) parent = third.manifest.parents.warm_start assert parent == embedded_model(second.manifest) - assert parent.parents.warm_start == ModelManifest(id=first.id, - reconstructed=True) assert parent.parents == ModelParents( training_table=second.manifest.parents.training_table, warm_start=ModelManifest(id=first.id, reconstructed=True)) @@ -142,21 +137,18 @@ def test_missing_parents_are_stubs_and_retrain_chains_stay_flat(tmp_path): def test_a_mlops_block_stores_each_model_once(tmp_path): cfg = kumo.MLOps(artifact_root=str(tmp_path / 'artifacts')) - result = _finish_stub_fit(tmp_path / 'bundle', kumo.Trainer(ModelPlan()), - settings=cfg) + result = _finish_stub_fit(tmp_path / 'bundle', settings=cfg) stored = tmp_path / 'artifacts' / 'models' / result.id assert kumo.Trainer.load(stored).result.id == result.id loaded = kumo.Trainer.load_from_connector(cfg, path=f'models/{result.id}') assert loaded.result.id == result.id assert loaded.result.snapshot_id == result.snapshot_id (stored / 'model' / 'config.json').write_text('{"damaged": 1}') - again = _finish_stub_fit(tmp_path / 'again', kumo.Trainer(ModelPlan()), - settings=cfg) + again = _finish_stub_fit(tmp_path / 'again', settings=cfg) assert again.id == result.id assert (stored / 'model' / 'config.json').read_text() == '{}' off = kumo.MLOps(artifact_root=str(tmp_path / 'off'), store={'snapshot'}) - _finish_stub_fit(tmp_path / 'off-bundle', kumo.Trainer(ModelPlan()), - settings=off) + _finish_stub_fit(tmp_path / 'off-bundle', settings=off) assert not (tmp_path / 'off').exists() From 7fdbf800762530bc4389649b96420d72bfbbf50d Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 3 Nov 2026 00:02:52 +0000 Subject: [PATCH 253/278] Score through one helper in the query-from-snapshot tests Co-Authored-By: Claude Opus 5.5 --- .../test/test_query_from_snapshot.py | 59 ++++++++----------- 1 file changed, 23 insertions(+), 36 deletions(-) diff --git a/packages/kumo-engine/test/test_query_from_snapshot.py b/packages/kumo-engine/test/test_query_from_snapshot.py index 8afcbc7ef..9dfd3ce51 100644 --- a/packages/kumo-engine/test/test_query_from_snapshot.py +++ b/packages/kumo-engine/test/test_query_from_snapshot.py @@ -6,50 +6,46 @@ import pandas as pd import pytest +from kumoapi.model_plan import ModelPlan import kumo.engine as kumo from kumo.engine._scratch import get_scratch_dir +from kumo.engine.pquery.predictive_query import PredictiveQuery -def _scores(trainer, table, snapshot): - frame = trainer.predict(table, snapshot=snapshot, - binary_classification_threshold=0.5).frame +def _scores(model, *args, **kwargs): + frame = model.predict(*args, binary_classification_threshold=0.5, + **kwargs).frame return frame.sort_values('ENTITY').reset_index(drop=True) def test_predict_scores_a_snapshot_with_the_saved_query(fitted, tmp_path): trainer, graph_table, snapshot = fitted - expected = _scores(trainer, graph_table, snapshot) - - def scores(**options): - return trainer.predict( - snapshot, binary_classification_threshold=0.5, - **options).frame.sort_values('ENTITY').reset_index(drop=True) - - pd.testing.assert_frame_equal(scores(), expected) + expected = _scores(trainer, graph_table, snapshot=snapshot) pd.DataFrame({ 'ENTITY1': range(50), 'TIMESTAMP': [pd.Timestamp(graph_table.timestamps[0])] * 50, }).to_parquet(tmp_path / 'rows.parquet') rows = kumo.FileConnector(root_dir=str(tmp_path))['rows.parquet'] - pd.testing.assert_frame_equal(scores(entities=rows), expected) - pd.testing.assert_frame_equal(scores(prediction_table=graph_table), expected) - legacy = trainer.predict(prediction_table=graph_table, snapshot=snapshot, - binary_classification_threshold=0.5).frame - pd.testing.assert_frame_equal( - legacy.sort_values('ENTITY').reset_index(drop=True), expected) - with pytest.raises(ValueError, match='not both'): - scores(entities=rows, prediction_time=graph_table.timestamps[0]) other = replace(snapshot, id='other@' + snapshot.id.split('@')[1]) - scored = trainer.predict(other, binary_classification_threshold=0.5).frame - pd.testing.assert_frame_equal( - scored.sort_values('ENTITY').reset_index(drop=True), expected) + for args, options in [ + ((snapshot, ), {}), + ((snapshot, ), {'entities': rows}), + ((snapshot, ), {'prediction_table': graph_table}), + ((), {'prediction_table': graph_table, 'snapshot': snapshot}), + ((other, ), {}), + ]: + pd.testing.assert_frame_equal(_scores(trainer, *args, **options), + expected) + with pytest.raises(ValueError, match='not both'): + _scores(trainer, snapshot, entities=rows, + prediction_time=graph_table.timestamps[0]) def test_exported_snapshot_scores_in_a_fresh_scratch_without_a_graph( fitted, tmp_path): trainer, graph_table, snapshot = fitted - expected = _scores(trainer, graph_table, snapshot) + expected = _scores(trainer, graph_table, snapshot=snapshot) exports = kumo.FileConnector(root_dir=str(tmp_path / 'exports')) location = snapshot.export(kumo.SnapshotOutputConfig(exports, 'snapshots')) bundle = trainer._training_result.bundle_dir @@ -60,11 +56,8 @@ def test_exported_snapshot_scores_in_a_fresh_scratch_without_a_graph( assert loaded.id == snapshot.id assert loaded.index_path.is_relative_to(tmp_path / 'fresh') model = kumo.Trainer.load(bundle) - scored = model.result.predict( - loaded, binary_classification_threshold=0.5).frame - pd.testing.assert_frame_equal( - scored.sort_values('ENTITY').reset_index(drop=True), expected, - atol=1e-5) + pd.testing.assert_frame_equal(_scores(model.result, loaded), expected, + atol=1e-5) finally: kumo.configure(scratch_dir=previous) @@ -75,8 +68,6 @@ def test_exported_snapshot_scores_in_a_fresh_scratch_without_a_graph( 'prediction_time': '2024-01-01' }]) def test_a_prediction_table_takes_no_entities_or_time(scope): - from kumoapi.model_plan import ModelPlan - with pytest.raises(ValueError, match='not both'): kumo.Trainer(ModelPlan()).predict(snapshot=object(), prediction_table=object(), **scope) @@ -94,10 +85,6 @@ def test_an_argument_given_twice_is_refused(fitted, call): def test_each_fit_result_scores_with_its_own_model(fitted): - from kumoapi.model_plan import ModelPlan - - from kumo.engine.pquery.predictive_query import PredictiveQuery - fixture_trainer, prediction_table, snapshot = fitted table = PredictiveQuery._from_snapshot( snapshot, @@ -112,11 +99,11 @@ def test_each_fit_result_scores_with_its_own_model(fitted): assert first.bundle_id != second.bundle_id scores = [ - _scores(result, prediction_table, snapshot) + _scores(result, prediction_table, snapshot=snapshot) for result in (first, second) ] assert not scores[0].equals(scores[1]) for result, score in zip((first, second), scores): own = kumo.Trainer.load(result.bundle_dir) pd.testing.assert_frame_equal( - score, _scores(own, prediction_table, snapshot)) + score, _scores(own, prediction_table, snapshot=snapshot)) From b571af13cde11624f85291e34807f467a480f449 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 3 Nov 2026 00:03:28 +0000 Subject: [PATCH 254/278] Warn about an unreadable manifest from one handler Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_manifest.py | 29 ++++++++----------- 1 file changed, 12 insertions(+), 17 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_manifest.py b/packages/kumo-engine/kumo/engine/_manifest.py index a8df4c800..0834e4014 100644 --- a/packages/kumo-engine/kumo/engine/_manifest.py +++ b/packages/kumo-engine/kumo/engine/_manifest.py @@ -306,25 +306,20 @@ def write_manifest(path: Path, manifest: Manifest) -> None: def read_manifest(path: Path, kind: Type[M]) -> Optional[M]: try: manifest = json.loads(path.read_text(encoding='utf-8')) + version = (manifest.get('schema_version') + if isinstance(manifest, dict) else None) + if type(version) is not int or not 1 <= version <= SCHEMA_VERSION: + warnings.warn( + f"Ignoring {path}: manifest schema_version {version!r} is not " + f"supported by this kumo-engine (up to {SCHEMA_VERSION}).", + UserWarning, + stacklevel=3, + ) + return None + return kind.from_dict(manifest) except FileNotFoundError: return None - except (OSError, ValueError) as error: - warnings.warn(f"Ignoring unreadable manifest {path}: {error}", - UserWarning, stacklevel=3) - return None - version = (manifest.get('schema_version') - if isinstance(manifest, dict) else None) - if type(version) is not int or not 1 <= version <= SCHEMA_VERSION: - warnings.warn( - f"Ignoring {path}: manifest schema_version {version!r} is not " - f"supported by this kumo-engine (up to {SCHEMA_VERSION}).", - UserWarning, - stacklevel=3, - ) - return None - try: - return kind.from_dict(manifest) - except (AttributeError, TypeError, ValueError) as error: + except (OSError, ValueError, AttributeError, TypeError) as error: warnings.warn(f"Ignoring unreadable manifest {path}: {error}", UserWarning, stacklevel=3) return None From 8cbfdcbe1551bb1e78d54da56639dcd86cafd5f3 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 3 Nov 2026 00:04:06 +0000 Subject: [PATCH 255/278] Inline the fit-mode lookup into the holdout metrics Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/trainer/_run_spec.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/_run_spec.py b/packages/kumo-engine/kumo/engine/trainer/_run_spec.py index e0fde9478..bb044bf1f 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_run_spec.py +++ b/packages/kumo-engine/kumo/engine/trainer/_run_spec.py @@ -74,7 +74,8 @@ def holdout_metrics(result: Any) -> Dict[str, float]: return logged_names({ name: value for name, value in result.metrics.items() if not name.startswith(('train_', 'val_')) - }, refit_full=_fit_mode(result) == 'refit_full') + }, refit_full=getattr(getattr(result, 'manifest', None), 'fit_mode', + None) == 'refit_full') def registry_metrics(result: Any) -> Dict[str, float]: @@ -87,10 +88,6 @@ def registry_metrics(result: Any) -> Dict[str, float]: }.items()} -def _fit_mode(result: Any) -> Optional[str]: - return getattr(getattr(result, 'manifest', None), 'fit_mode', None) - - def _registered(model: Any) -> Mapping[str, Any]: if model is None: return {} From e7067c690a50ca97039bf2b806cf926972e46230 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 3 Nov 2026 00:04:31 +0000 Subject: [PATCH 256/278] Format Snowflake's schema-check error with full_error Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/integrations/snowflake.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/packages/kumo-engine/kumo/engine/integrations/snowflake.py b/packages/kumo-engine/kumo/engine/integrations/snowflake.py index 777a53b33..11196a299 100644 --- a/packages/kumo-engine/kumo/engine/integrations/snowflake.py +++ b/packages/kumo-engine/kumo/engine/integrations/snowflake.py @@ -17,6 +17,7 @@ from kumo_core.summaries import Summary from kumo.engine._events import Event +from kumo.engine._mlops import full_error from kumo.engine._run_names import record_name from kumo.engine._summary_tables import route, table_json from kumo.engine._tracking import ( @@ -118,7 +119,7 @@ def check(self, registry: Any = None) -> None: except Exception as error: raise RuntimeError( f'Snowflake Experiment Tracking cannot use {database}: ' - f'{type(error).__name__}: {error}') from error + f'{full_error(error)}') from error if not rows: raise ValueError( f'Snowflake schema {database}.{schema} does not exist or this ' From adad4f8d865c622339255557c3bc9fb6012b4975 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 3 Nov 2026 00:02:25 +0000 Subject: [PATCH 257/278] Trim volume and sink docstrings and redundant path stripping Co-Authored-By: Claude Opus 5.5 --- packages/kumo-databricks/src/kumo_databricks/sink.py | 2 -- packages/kumo-databricks/src/kumo_databricks/volume.py | 8 +------- 2 files changed, 1 insertion(+), 9 deletions(-) diff --git a/packages/kumo-databricks/src/kumo_databricks/sink.py b/packages/kumo-databricks/src/kumo_databricks/sink.py index e29027771..0b82aca01 100644 --- a/packages/kumo-databricks/src/kumo_databricks/sink.py +++ b/packages/kumo-databricks/src/kumo_databricks/sink.py @@ -155,12 +155,10 @@ class DatabricksSinkBackend(FileStoreRouting, _DatabricksTables): 'DatabricksSinkBackend(identity=identity)})') def holds_files(self, connector: SourceConnector, path: str) -> bool: - """Volume paths hold files; anything else is a table.""" return is_volume_path(path) def filesystem(self, connector: SourceConnector, path: str) -> tuple[FileSystem, str]: - """A volume path's filesystem, as the connector's identity.""" if not self.holds_files(connector, path): raise ValueError( f"{path!r} is not a Unity Catalog volume path; Databricks " diff --git a/packages/kumo-databricks/src/kumo_databricks/volume.py b/packages/kumo-databricks/src/kumo_databricks/volume.py index 6d85765d9..55e486107 100644 --- a/packages/kumo-databricks/src/kumo_databricks/volume.py +++ b/packages/kumo-databricks/src/kumo_databricks/volume.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Unity Catalog volumes as a pyarrow filesystem, over the Files API.""" +"""Unity Catalog volumes as a pyarrow filesystem.""" from __future__ import annotations @@ -36,7 +36,6 @@ def volume_filesystem(client: WorkspaceClient) -> FileSystem: @contextmanager def _translated(path: str) -> Iterator[None]: - """OS errors, unchained; non-SDK text may hold a signed URL.""" try: yield except NotFound as error: @@ -60,7 +59,6 @@ def get_type_name(self) -> str: return 'databricks-volume' def _size(self, path: str) -> Optional[int]: - path = path.rstrip('/') with _translated(path): try: return int(self._client.files.get_metadata(path).content_length) @@ -68,7 +66,6 @@ def _size(self, path: str) -> Optional[int]: return None def _list(self, directory: str) -> Dict[str, int]: - directory = directory.rstrip('/') with _translated(directory): try: entries = list( @@ -95,7 +92,6 @@ def _part(self, path: str, size: int) -> _Part: return _Part(self._client, path) def _remove(self, path: str, *, tree: bool) -> None: - path = path.rstrip('/') with _translated(path): if tree: clear_volume(self._client, path) @@ -107,7 +103,6 @@ def _remove(self, path: str, *, tree: bool) -> None: class _Part: - """Serves consecutive reads from one open-ended ranged download.""" def __init__(self, client: WorkspaceClient, path: str) -> None: self._client, self._path = client, path self._stream = None @@ -185,7 +180,6 @@ def readinto(self, buffer) -> int: class _Upload(io.RawIOBase): - """One ``files.upload`` fed from writes; ``close`` commits it.""" def __init__(self, client: WorkspaceClient, path: str) -> None: self._path = path self._feed = _Feed() From 50ddd18c5d7fff05013165ab59588786139281f0 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 3 Nov 2026 00:15:50 +0000 Subject: [PATCH 258/278] Tighten volume README and hoist test imports Co-Authored-By: Claude Opus 5.5 --- packages/kumo-databricks/README.md | 19 +++++++++---------- packages/kumo-databricks/tests/test_volume.py | 10 +++------- 2 files changed, 12 insertions(+), 17 deletions(-) diff --git a/packages/kumo-databricks/README.md b/packages/kumo-databricks/README.md index 8de659f9d..bfa1b9c4a 100644 --- a/packages/kumo-databricks/README.md +++ b/packages/kumo-databricks/README.md @@ -86,16 +86,15 @@ trainer.predict(prediction_table, snapshot=snapshot, output_config=PredictionOut For an MLOps artifact root on a volume, see `packages/kumo-engine/docs/mlops.md`. -A file whose size is known when it is written (bundle files, exported Parquet -parts) is stored as one object up to 4 GiB, so Spark and other tools read it -directly. A larger one streams into `.chunks//00000`, -`00001`, ... (4 GiB each, under the Files API's 5 GiB upload cap) with an empty -`_SUCCESS` written last, and Kumo reads the newest committed generation back as -one file. A file written without a known size stays one object up to 4 GiB, -staged on local disk under the engine scratch directory, and streams into parts -above that. The identity needs `USE CATALOG`, `USE SCHEMA`, and `READ VOLUME` -(loads and exports) and `WRITE VOLUME` (exports) on the volume; no warehouse -is used. +A file written with a known size (bundle files, exported Parquet parts) is one +object up to 4 GiB, readable directly by Spark and other tools. A larger one +streams into `.chunks//00000`, `00001`, ... (4 GiB each, +under the Files API's 5 GiB upload cap) with an empty `_SUCCESS` written last; +Kumo reads the newest committed generation back as one file. A file written +without a known size is staged on local disk under the engine scratch +directory and stored the same way. The identity needs `USE CATALOG`, +`USE SCHEMA`, and `READ VOLUME` (loads and exports) and `WRITE VOLUME` +(exports) on the volume; no warehouse is used. ## Model registry diff --git a/packages/kumo-databricks/tests/test_volume.py b/packages/kumo-databricks/tests/test_volume.py index 71de4a7eb..266403c43 100644 --- a/packages/kumo-databricks/tests/test_volume.py +++ b/packages/kumo-databricks/tests/test_volume.py @@ -9,8 +9,11 @@ import pytest from databricks.sdk.errors import PermissionDenied from fakes import FakeVolume, FakeWorkspaceClient +from kumo_core.artifact_export import ArtifactDestinationConflict +from kumo_core.object_store_sink import open_sized from kumo_databricks.sink import DatabricksSinkBackend from kumo_databricks.volume import VolumeHandler +from kumoapi.jobs import WriteMode from pyarrow.fs import FileType, PyFileSystem ROOT = '/Volumes/main/analytics/artifacts' @@ -47,8 +50,6 @@ def test_reads_cross_part_boundaries(): def test_sized_writes_stay_plain_up_to_a_part(): - from kumo_core.object_store_sink import open_sized - fs, volume = filesystem(part_bytes=1000) fs.handler.plain_bytes = 100 for name, body in (('plain', BODY[:1000]), ('parts', BODY[:1500])): @@ -72,8 +73,6 @@ def test_unsized_writes_stay_plain_through_a_part(): def test_a_failed_upload_raises_without_the_signed_url(): - from kumo_core.object_store_sink import open_sized - fs, volume = filesystem() signed = 'https://bucket.example/part?X-Amz-Signature=secret' @@ -195,9 +194,6 @@ def test_predictions_and_frames_written_to_a_volume_path(tmp_path): def test_an_occupied_volume_directory_is_refused_before_scoring(tmp_path): - from kumo_core.artifact_export import ArtifactDestinationConflict - from kumoapi.jobs import WriteMode - sink, volume = volume_sink() target = f'{ROOT}/predictions/' sink.preflight_write_files(SimpleNamespace(), target) From 12faf65392144706d6d4198c08e2e9520c535885 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 3 Nov 2026 00:21:35 +0000 Subject: [PATCH 259/278] Route FileStoreRouting's file paths through one helper Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_core/object_store_sink.py | 44 ++++++++----------- 1 file changed, 19 insertions(+), 25 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/object_store_sink.py b/packages/kumo-core/src/kumo_core/object_store_sink.py index d7433cf3d..84c06c08e 100644 --- a/packages/kumo-core/src/kumo_core/object_store_sink.py +++ b/packages/kumo-core/src/kumo_core/object_store_sink.py @@ -433,8 +433,10 @@ def _uri(self, path: str) -> str: class FileStoreRouting: - """Sends ``holds_files`` paths to :class:`FilesystemSink`, the rest on.""" - def _rest(self, method): + """Routes ``holds_files`` paths to :class:`FilesystemSink`.""" + def _route(self, connector, path, method): + if path is not None and self.holds_files(connector, path): + return FilesystemSink(self) rest = super() if not hasattr(rest, method): raise NotImplementedError( @@ -444,39 +446,31 @@ def _rest(self, method): def write_table(self, connector, table_name, frame, write_mode=WriteMode.OVERWRITE, project_id=None, *, export_id=None): - sink = FilesystemSink(self) if self.holds_files( - connector, table_name) else self._rest('write_table') - return sink.write_table(connector, table_name, frame, write_mode, - project_id, export_id=export_id) + return self._route(connector, table_name, 'write_table').write_table( + connector, table_name, frame, write_mode, project_id, + export_id=export_id) def write_files(self, connector, table_name, files, write_mode=WriteMode.OVERWRITE, project_id=None, *, export_id=None, extra_columns=None): - sink = FilesystemSink(self) if self.holds_files( - connector, table_name) else self._rest('write_files') - return sink.write_files(connector, table_name, files, write_mode, - project_id, export_id=export_id, - extra_columns=extra_columns) + return self._route(connector, table_name, 'write_files').write_files( + connector, table_name, files, write_mode, project_id, + export_id=export_id, extra_columns=extra_columns) def publish_model(self, connector, bundle_dir, *, path=None, location=None, metadata={}, metrics={}, project_id=None): - if location is None and path is not None and self.holds_files( - connector, path): - return FilesystemSink(self).publish_model( - connector, bundle_dir, path=path, metadata=metadata) - return self._rest('publish_model').publish_model( - connector, bundle_dir, path=path, location=location, - metadata=metadata, metrics=metrics, project_id=project_id) + sink = self._route(connector, path if location is None else None, + 'publish_model') + return sink.publish_model(connector, bundle_dir, path=path, + location=location, metadata=metadata, + metrics=metrics, project_id=project_id) def fetch_model(self, connector, local_dir, *, path=None, location=None, project_id=None): - if location is None and path is not None and self.holds_files( - connector, path): - return FilesystemSink(self).fetch_model(connector, local_dir, - path=path) - return self._rest('fetch_model').fetch_model( - connector, local_dir, path=path, location=location, - project_id=project_id) + sink = self._route(connector, path if location is None else None, + 'fetch_model') + return sink.fetch_model(connector, local_dir, path=path, + location=location, project_id=project_id) def preflight_write_files(self, connector, table_name, write_mode=None): if self.holds_files(connector, table_name): From bde829bbc273e58f14dcf3af85b43ad84eede036 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 3 Nov 2026 00:21:50 +0000 Subject: [PATCH 260/278] Raise the occupied-directory error from one place Co-Authored-By: Claude Opus 5.5 --- packages/kumo-core/src/kumo_core/object_store_sink.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/object_store_sink.py b/packages/kumo-core/src/kumo_core/object_store_sink.py index 84c06c08e..24002b756 100644 --- a/packages/kumo-core/src/kumo_core/object_store_sink.py +++ b/packages/kumo-core/src/kumo_core/object_store_sink.py @@ -156,7 +156,7 @@ def _conflict(self, path: str) -> ValueError: return ArtifactDestinationConflict( f"Output directory '{self._uri(path.rstrip('/'))}/' already " - f"contains data from another export. Choose an empty destination " + "contains data from another export. Choose an empty destination " "path.") def publish_model( @@ -291,10 +291,7 @@ def _occupied(self, filesystem: FileSystem, directory: str, with filesystem.open_input_stream(path) as reader: if reader.readall().decode('utf-8') == marker: return True - raise ValueError( - f"Output directory '{self._uri(directory)}/' already " - f"contains data from another export. Choose an empty destination " - f"path.") + raise self._conflict(directory) def _write_one(filesystem: FileSystem, frame: pd.DataFrame, path: str, From 07b2caf3b5f82ea4bc118bb1ffc9f97992a8585c Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 3 Nov 2026 00:21:50 +0000 Subject: [PATCH 261/278] Drop docstrings that restate the file-store names Co-Authored-By: Claude Opus 5.5 --- packages/kumo-core/src/kumo_core/object_store_sink.py | 4 ---- packages/kumo-engine/kumo/engine/connector/__init__.py | 1 - 2 files changed, 5 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/object_store_sink.py b/packages/kumo-core/src/kumo_core/object_store_sink.py index 24002b756..35cda0b6e 100644 --- a/packages/kumo-core/src/kumo_core/object_store_sink.py +++ b/packages/kumo-core/src/kumo_core/object_store_sink.py @@ -33,7 +33,6 @@ def _chunked(filesystem: FileSystem) -> bool: def open_sized(filesystem: FileSystem, path: str, size: int) -> Any: - """``open_output_stream`` that tells a chunked filesystem the size.""" if _chunked(filesystem): from kumo_core.chunked_fs import SIZE_KEY @@ -51,7 +50,6 @@ def upload_file(filesystem: FileSystem, source: str, def upload_bundle(filesystem: FileSystem, destination: str, bundle_dir: str) -> None: - """Copy ``bundle_dir`` under ``destination``, provenance last.""" from kumo_core.artifacts import PROVENANCE_FILE root = Path(bundle_dir) @@ -140,7 +138,6 @@ def write_files(self, connector, table_name, files, return self._uri(output) def preflight_write_files(self, connector, table_name, write_mode=None): - """Refuse an occupied output directory before scoring.""" del write_mode from pyarrow.fs import FileSelector @@ -416,7 +413,6 @@ def convert_csv(self, source, destination): class FilesystemSink(ObjectStoreSinkBackend): - """The object-store file writers over any sink's ``filesystem()``.""" def __init__(self, sink: FileStoreSink) -> None: self._sink = sink self.provider = type(sink).__name__ diff --git a/packages/kumo-engine/kumo/engine/connector/__init__.py b/packages/kumo-engine/kumo/engine/connector/__init__.py index c2bdeea22..34b55e4f8 100644 --- a/packages/kumo-engine/kumo/engine/connector/__init__.py +++ b/packages/kumo-engine/kumo/engine/connector/__init__.py @@ -51,7 +51,6 @@ def _file_sink(connector: Connector, path: str) -> Optional[FileStoreSink]: def is_file_destination(connector: Connector, path: str) -> bool: - r"""Whether ``path`` is a file location rather than a table.""" return _file_sink(connector, path) is not None From 48a8552a9f92cc9225fdc1f45a4debf8c82f331e Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 3 Nov 2026 00:22:56 +0000 Subject: [PATCH 262/278] Share the streaming store fixture and the one-form check in the chunked tests Co-Authored-By: Claude Opus 5.5 --- packages/kumo-core/tests/test_chunked_fs.py | 104 ++++++++++---------- 1 file changed, 53 insertions(+), 51 deletions(-) diff --git a/packages/kumo-core/tests/test_chunked_fs.py b/packages/kumo-core/tests/test_chunked_fs.py index 2eea963eb..801904a39 100644 --- a/packages/kumo-core/tests/test_chunked_fs.py +++ b/packages/kumo-core/tests/test_chunked_fs.py @@ -2,12 +2,15 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import gc import io import os +import threading import pytest from kumo_core import chunked_fs from kumo_core.chunked_fs import ChunkedFileHandler +from kumo_core.object_store_sink import open_sized from pyarrow.fs import FileSelector, FileType, PyFileSystem PART = 16 @@ -60,12 +63,40 @@ def _remove(self, path, *, tree): del self.objects[name] +class _Upload(io.BytesIO): + def __init__(self, handler, path): + super().__init__() + self.handler, self.path = handler, path + + def close(self): + self.handler.objects[self.path] = self.getvalue() + self.handler.streamed.append(self.path) + super().close() + + +class StreamingHandler(MemoryHandler): + plain_bytes = 6 + + def __init__(self): + super().__init__() + self.streamed = [] + + def _stream(self, path): + return _Upload(self, path) + + @pytest.fixture def store(): handler = MemoryHandler() return handler, PyFileSystem(handler) +@pytest.fixture +def streaming(): + handler = StreamingHandler() + return handler, PyFileSystem(handler) + + def _write(fs, path, data): with fs.open_output_stream(path) as writer: for start in range(0, len(data), 5): @@ -77,6 +108,12 @@ def _read(fs, path): return reader.readall() +def _is_plain(handler): + plain = 'root/a.bin' in handler.objects + assert any('.chunks/' in name for name in handler.objects) != plain + return plain + + def _generations(handler): parts = {} for name in sorted(handler.objects): @@ -129,9 +166,7 @@ def test_a_write_removes_the_other_form(store, first, second): _write(fs, 'root/a.bin', b'1' * first) _write(fs, 'root/a.bin', b'2' * second) assert _read(fs, 'root/a.bin') == b'2' * second - plain = 'root/a.bin' in handler.objects - assert plain == (second <= PART) - assert any('.chunks/' in name for name in handler.objects) != plain + assert _is_plain(handler) == (second <= PART) @pytest.mark.parametrize('sized', [False, True]) @@ -139,8 +174,6 @@ def test_a_write_removes_the_other_form(store, first, second): @pytest.mark.parametrize('old', [3, 2 * PART]) def test_a_failed_chunked_write_leaves_the_old_file( monkeypatch, streams, sized, old): - from kumo_core.object_store_sink import open_sized - handler = StreamingHandler() if streams else MemoryHandler() fs = PyFileSystem(handler) _write(fs, 'root/a.bin', b'o' * old) @@ -157,18 +190,19 @@ def second_part_fails(*args): with (open_sized(fs, 'root/a.bin', 3 * PART) if sized else fs.open_output_stream('root/a.bin')) as writer: writer.write(b'n' * (3 * PART)) - assert _read(PyFileSystem(handler), 'root/a.bin') == b'o' * old + assert _read(fs, 'root/a.bin') == b'o' * old monkeypatch.setattr(handler, upload, original) _write(fs, 'root/a.bin', b'n' * (3 * PART)) - assert _read(PyFileSystem(handler), 'root/a.bin') == b'n' * (3 * PART) + assert _read(fs, 'root/a.bin') == b'n' * (3 * PART) assert list(_generations(handler).values()) == [ ['00000', '00001', '00002', '_SUCCESS'] ] -def test_a_failed_plain_stream_is_aborted_not_committed(monkeypatch): - from kumo_core.object_store_sink import open_sized +def test_a_failed_plain_stream_is_aborted_not_committed( + monkeypatch, streaming): + handler, fs = streaming def fail(self, data): raise OSError('upload failed') @@ -177,12 +211,11 @@ def abort(self): aborted.append(self.path) io.BytesIO.close(self) - handler = StreamingHandler() aborted = [] monkeypatch.setattr(_Upload, 'write', fail) monkeypatch.setattr(_Upload, 'abort', abort, raising=False) with pytest.raises(OSError, match='upload failed'): - with open_sized(PyFileSystem(handler), 'root/a.bin', 3) as writer: + with open_sized(fs, 'root/a.bin', 3) as writer: writer.write(b'abc') assert aborted == ['root/a.bin'] assert handler.objects == {} @@ -264,32 +297,10 @@ def test_parts_spool_under_the_spool_dir_and_are_removed( assert not os.listdir(tmp_path / 'own') -class _Upload(io.BytesIO): - def __init__(self, handler, path): - super().__init__() - self.handler, self.path = handler, path - - def close(self): - self.handler.objects[self.path] = self.getvalue() - self.handler.streamed.append(self.path) - super().close() - - -class StreamingHandler(MemoryHandler): - plain_bytes = 6 - - def __init__(self): - super().__init__() - self.streamed = [] - - def _stream(self, path): - return _Upload(self, path) - - @pytest.mark.parametrize('size', [0, 6, 7, PART, 2 * PART + 5, 3 * PART]) -def test_a_streaming_handler_streams_parts_and_spools_only_small_files(size): - handler = StreamingHandler() - fs = PyFileSystem(handler) +def test_a_streaming_handler_streams_parts_and_spools_only_small_files( + streaming, size): + handler, fs = streaming _write(fs, 'root/a.bin', b'x' * (2 * PART)) data = os.urandom(size) _write(fs, 'root/a.bin', data) @@ -326,9 +337,7 @@ def test_opening_a_file_again_reuses_its_layout(store): def test_layout_cache_survives_concurrent_opens_and_writes(store): - import threading - - handler, fs = store + _, fs = store for i in range(20): _write(fs, f'root/{i}.bin', b'x' * (PART + 1)) errors = [] @@ -350,11 +359,8 @@ def work(i): @pytest.mark.parametrize('size', [7, PART, PART + 1, 3 * PART]) -def test_a_known_size_streams_without_spooling(size): - from kumo_core.object_store_sink import open_sized - - handler = StreamingHandler() - fs = PyFileSystem(handler) +def test_a_known_size_streams_without_spooling(streaming, size): + handler, fs = streaming _write(fs, 'root/a.bin', b'x' * (2 * PART)) handler.spooled.clear() data = os.urandom(size) @@ -362,9 +368,8 @@ def test_a_known_size_streams_without_spooling(size): for start in range(0, size, 5): writer.write(data[start:start + 5]) assert _read(fs, 'root/a.bin') == data - plain = 'root/a.bin' in handler.objects + plain = _is_plain(handler) assert plain == (size <= PART) - assert any('.chunks/' in name for name in handler.objects) != plain # Only the empty _SUCCESS marker of a chunked file is put from disk. assert len(handler.spooled) == (0 if plain else 1) @@ -383,12 +388,9 @@ def test_reads_spanning_parts_fill_the_callers_buffer(store): @pytest.mark.parametrize('size', [1, 3 * PART]) def test_a_writer_that_fails_to_open_keeps_the_old_file( - monkeypatch, tmp_path, size): - import gc - - handler = StreamingHandler() + monkeypatch, tmp_path, streaming, size): + handler, fs = streaming handler.spool_dir = str(tmp_path) - fs = PyFileSystem(handler) _write(fs, 'root/a.bin', b'preserve-me') def fail(path): From dafc5434a3f654d2960f4aa81721433515817b5f Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 3 Nov 2026 00:22:56 +0000 Subject: [PATCH 263/278] Merge the two occupied-directory preflight tests Co-Authored-By: Claude Opus 5.5 --- .../kumo-engine/test/test_filestore_sink.py | 29 ++++++------------- 1 file changed, 9 insertions(+), 20 deletions(-) diff --git a/packages/kumo-engine/test/test_filestore_sink.py b/packages/kumo-engine/test/test_filestore_sink.py index f8dd18e95..2685785ad 100644 --- a/packages/kumo-engine/test/test_filestore_sink.py +++ b/packages/kumo-engine/test/test_filestore_sink.py @@ -109,6 +109,10 @@ def test_output_configs_take_a_sinks_file_paths(connector): PredictionOutputConfig(connector, output_paths={'predictions': 'T'}) PredictionOutputConfig(connector, PredictionArtifactType.PREDICTIONS, 'PREDS') + tables = SimpleNamespace(sink_backend=object()) + with pytest.raises(NotImplementedError, match='requires'): + PredictionOutputConfig(tables, PredictionArtifactType.PREDICTIONS, + 'DB.S.T').writer() def test_bundles_upload_provenance_last(tmp_path, connector): @@ -152,10 +156,6 @@ def test_predictions_export_through_a_filesystem_only_sink(fitted, connector): PredictionOutputConfig(connector, output_paths={'predictions': '@stage/again/'})) assert any(name.startswith('@stage/again/') for name in store.objects) - tables = SimpleNamespace(sink_backend=object()) - with pytest.raises(NotImplementedError, match='requires'): - PredictionOutputConfig(tables, PredictionArtifactType.PREDICTIONS, - 'DB.S.T').writer() def test_parquet_writes_stay_plain_up_to_one_part(connector): @@ -176,26 +176,15 @@ def test_parquet_writes_stay_plain_up_to_one_part(connector): def test_an_occupied_directory_is_refused_by_its_full_location( tmp_path, connector): - files = connector.sink_backend + sink = connector.sink_backend part = tmp_path / 'part-0.parquet' pd.DataFrame({'x': range(3)}).to_parquet(part) - files.write_files(connector, '@stage/out/', [str(part)]) - - files.preflight_write_files(connector, '@stage/new/') - with pytest.raises(ArtifactDestinationConflict, match="'@stage/out/'"): - files.preflight_write_files(connector, '@stage/out/') - - -def test_an_occupied_directory_is_refused_before_scoring(connector): - sink = connector.sink_backend sink.preflight_write_files(connector, '@stage/out/') - sink.store.objects['@stage/out/old.parquet'] = b'x' - for mode in (None, WriteMode.APPEND): - with pytest.raises(ArtifactDestinationConflict, match='already'): + sink.write_files(connector, '@stage/out/', [str(part)]) + for mode in (None, WriteMode.APPEND, WriteMode.OVERWRITE): + with pytest.raises(ArtifactDestinationConflict, + match="'@stage/out/' already"): sink.preflight_write_files(connector, '@stage/out/', mode) - with pytest.raises(ArtifactDestinationConflict, match='already'): - sink.preflight_write_files(connector, '@stage/out/', - WriteMode.OVERWRITE) def test_a_file_only_sink_refuses_a_table(connector, tmp_path): From 6036fd286343b200cf56549265271496a0692a98 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 3 Nov 2026 00:26:16 +0000 Subject: [PATCH 264/278] Say in the docs that MLflow's default store is sqlite:///mlflow.db Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/docs/mlops.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/kumo-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index 5520328a8..83f4f7def 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -12,7 +12,7 @@ Each block assumes a `graph` and a `pquery` (see the README). It trains, stores ```python import kumo.engine as kumo -from kumo.engine.integrations.mlflow import MlflowTracking # MLflow's own tracking URI, else ./mlruns +from kumo.engine.integrations.mlflow import MlflowTracking # where MLflow's own setting points; by default sqlite:///mlflow.db cfg = kumo.MLOps("artifacts", tracking=MlflowTracking("churn"), registry=kumo.ModelRegistryLocation("users-churn", alias="candidate")) @@ -183,7 +183,7 @@ Jobs that load `alias="champion"` pick up the new version on their next run. ### Trackers `MlflowTracking(experiment_name=None, run_name=None, log_system_metrics=True, tags=None, log_model=False, *, target=None)`: -- With no `target`, it records where MLflow's own configuration points (`MLFLOW_TRACKING_URI`, `mlflow.set_tracking_uri`; automatic on Databricks). `target` names another server with its credentials: a `DatabricksConnector`, a `SageMakerMLflow(arn)`, or a tracking URI. With a `target`, `experiment_name` is required to record runs. +- With no `target`, it records where MLflow's own setting points (`MLFLOW_TRACKING_URI`, `mlflow.set_tracking_uri`; automatic on Databricks); MLflow's default is a local `sqlite:///mlflow.db`. `target` names another server with its credentials: a `DatabricksConnector`, a `SageMakerMLflow(arn)`, or a tracking URI. With a `target`, `experiment_name` is required to record runs. - `log_model=True` logs the model on the fit run without registering it; `MLOps(registry=)` implies it. - Building `kumo.MLOps` creates a missing experiment; if MLflow refuses (on Databricks, the experiment's folder must exist), it raises and says so. From 7fe756a89d8ce675979c4be77524272b5a846b04 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 3 Nov 2026 00:26:31 +0000 Subject: [PATCH 265/278] Say plainly when MLflow refuses a local folder store Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/integrations/mlflow.py | 14 +++++++++++--- .../kumo-engine/test/test_mlflow_target.py | 19 +++++++++++++++++++ 2 files changed, 30 insertions(+), 3 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index bc04e534f..56cc046dc 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -33,7 +33,7 @@ Tuple, cast, ) -from urllib.parse import quote +from urllib.parse import quote, urlparse from kumo_core.instrumentation import format_error from kumo_core.summaries import Dataset, Summary, Table @@ -182,8 +182,7 @@ def check(self, registry: Optional[ModelRegistryLocation] = None) -> None: except ValueError: raise except Exception as error: - raise RuntimeError( - f'MLflow at {uri} is not usable: {full_error(error)}') from error + raise _unusable(uri, error) from error if not missing: return with _limited_http(), self.target.scope(): @@ -691,6 +690,15 @@ def check_model_name(name: str, registry_uri: str) -> None: f"'main.ml.users_churn'.") +def _unusable(uri: str, error: Exception) -> RuntimeError: + if (urlparse(uri).scheme in ('', 'file') + and 'MLFLOW_ALLOW_FILE_STORE' in str(error)): + return RuntimeError( + f'MLflow refused the local folder store at {uri}: use its default ' + f'sqlite:///mlflow.db or another tracking backend.') + return RuntimeError(f'MLflow at {uri} is not usable: {full_error(error)}') + + def _sagemaker_ui_hint(tracking_uri: str) -> Optional[str]: match = SAGEMAKER_ARN.fullmatch(tracking_uri or '') if match is None: diff --git a/packages/kumo-engine/test/test_mlflow_target.py b/packages/kumo-engine/test/test_mlflow_target.py index 1694e14ea..280e5ce6d 100644 --- a/packages/kumo-engine/test/test_mlflow_target.py +++ b/packages/kumo-engine/test/test_mlflow_target.py @@ -8,6 +8,7 @@ from types import SimpleNamespace import pytest +from packaging.version import Version import kumo.engine as kumo from kumo.engine.connector import ( @@ -233,6 +234,24 @@ def test_pin_mlflow_reads_its_run_link_switch_per_call(): assert MLFLOW_SUPPRESS_PRINTING_URL_TO_STDOUT.name == NO_LINKS +@pytest.mark.skipif(Version(mlflow.__version__) < Version('3.17'), + reason='this MLflow still accepts folder stores') +@pytest.mark.parametrize('store', ['file', 'path']) +def test_a_refused_folder_store_says_to_use_sqlite(tmp_path, monkeypatch, + store): + from kumo.engine.integrations.mlflow import MlflowTracking + + monkeypatch.delenv('MLFLOW_ALLOW_FILE_STORE', raising=False) + uri = (tmp_path / 'mlruns').as_uri() if store == 'file' else str( + tmp_path / 'mlruns') + tracking = MlflowTracking(target=uri, experiment_name='e') + with pytest.raises(RuntimeError) as raised: + kumo.MLOps(str(tmp_path / 'root'), tracking=tracking) + assert str(raised.value) == ( + f'MLflow refused the local folder store at {uri}: use its default ' + f'sqlite:///mlflow.db or another tracking backend.') + + def test_an_unreachable_server_fails_when_mlops_is_built( tmp_path, monkeypatch): from kumo.engine.integrations.mlflow import MlflowTracking From d78d12bfaf774025a406f7d0d4ffe36836f0dc6a Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 3 Nov 2026 00:33:14 +0000 Subject: [PATCH 266/278] Delete manifest, warning and model helpers nothing in this PR uses Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_manifest.py | 69 +------------------ packages/kumo-engine/kumo/engine/_mlops.py | 19 +---- .../kumo/engine/pquery/training_table.py | 3 +- .../kumo/engine/trainer/_provenance.py | 10 ++- .../kumo/engine/trainer/trainer.py | 7 +- packages/kumo-engine/test/test_artifacts.py | 62 ----------------- 6 files changed, 10 insertions(+), 160 deletions(-) delete mode 100644 packages/kumo-engine/test/test_artifacts.py diff --git a/packages/kumo-engine/kumo/engine/_manifest.py b/packages/kumo-engine/kumo/engine/_manifest.py index 0834e4014..93391ce99 100644 --- a/packages/kumo-engine/kumo/engine/_manifest.py +++ b/packages/kumo-engine/kumo/engine/_manifest.py @@ -5,14 +5,12 @@ import types import warnings from collections.abc import Mapping as AbcMapping -from dataclasses import dataclass, field, fields, replace +from dataclasses import dataclass, field, fields from datetime import datetime, timezone from pathlib import Path from typing import ( Any, - ClassVar, Dict, - Iterator, List, Mapping, Optional, @@ -37,19 +35,6 @@ def omit_none() -> Any: return field(default=None, metadata={_OMIT_NONE: True}) -@dataclass(frozen=True) -class Artifact: - kind: str - id: str - location: Optional[str] - is_dataset: bool - tag: Optional[str] = None - connector: Optional[str] = None - version: Optional[str] = None - version_kind: Optional[str] = None - context: Optional[str] = None - - @dataclass(frozen=True, kw_only=True) class Record: """A JSON object that keeps unknown keys in ``extra``.""" @@ -57,9 +42,6 @@ class Record: extra: Mapping[str, Any] = field(default_factory=dict, repr=False, compare=False) - def as_artifact(self, key: str) -> Optional[Artifact]: - return None - def to_dict(self) -> Dict[str, Any]: data = dict(self.extra) for item in fields(self): @@ -139,11 +121,6 @@ class SourceTableEntry(Record): dialect: Optional[str] = omit_none() namespace: Optional[List[Any]] = omit_none() - def as_artifact(self, key: str) -> Artifact: - return Artifact('source_table', key, self.location, True, - connector=self.connector, version=self.version, - version_kind=self.version_kind) - KINDS = ('snapshot', 'training_table', 'model') @@ -156,40 +133,6 @@ class Manifest(Record): created_at: Optional[str] = None environment: Optional[Environment] = None reconstructed: Optional[bool] = omit_none() - location: Optional[str] = omit_none() - - EMBED_OMITS = () - IS_DATASET: ClassVar[bool] = False - URI_TAG: ClassVar[Optional[str]] = None - - def as_artifact(self, key: str) -> Artifact: - return Artifact(self.kind, self.id, self.location, self.IS_DATASET, - self.URI_TAG, version=self.id, version_kind='exact') - - -def artifacts(manifest: Manifest) -> List[Artifact]: - found: Dict[Any, Artifact] = {} - for artifact in _walk(manifest, manifest.kind, manifest): - found.setdefault((artifact.kind, artifact.id), artifact) - return list(found.values()) - - -def _walk(value: Any, key: str, root: Manifest) -> Iterator[Artifact]: - if isinstance(value, Record): - artifact = value.as_artifact(key) - if artifact is not None: - yield artifact - if isinstance(value, Manifest) and not value.IS_DATASET and \ - value is not root: - return - for item in fields(value): - yield from _walk(getattr(value, item.name), item.name, root) - elif isinstance(value, AbcMapping): - for name, item in value.items(): - yield from _walk(item, str(name), root) - elif isinstance(value, (list, tuple)): - for item in value: - yield from _walk(item, key, root) def stamp() -> Dict[str, Any]: @@ -219,13 +162,6 @@ def stub(kind: str, id: str) -> 'Manifest': return _kinds()[kind](id=id, reconstructed=True) -def embedded(manifest: M) -> M: - return replace(manifest, **{ - name: None - for name in manifest.EMBED_OMITS - }) - - def _parent(hint: Type[M], value: Any) -> M: try: return hint.from_dict(value) @@ -246,7 +182,6 @@ class SnapshotManifest(Manifest): """The sources and engine build of a snapshot.""" kind: str = 'snapshot' - IS_DATASET: ClassVar[bool] = True parents: Optional[SnapshotParents] = None record_sources: Optional[bool] = omit_none() @@ -262,7 +197,6 @@ class TrainingTableManifest(Manifest): """The snapshot, query and rows of a training table.""" kind: str = 'training_table' - IS_DATASET: ClassVar[bool] = True parents: Optional[TrainingTableParents] = None query: Optional[str] = None generation_plan: Optional[Dict[str, Any]] = None @@ -288,7 +222,6 @@ class ModelManifest(Manifest): """How a model bundle was fitted; ``id`` is the bundle id.""" kind: str = 'model' - URI_TAG: ClassVar[Optional[str]] = 'kumo.bundle_uri' parents: Optional[ModelParents] = None fit: Optional[FitTiming] = None fit_mode: Optional[str] = None diff --git a/packages/kumo-engine/kumo/engine/_mlops.py b/packages/kumo-engine/kumo/engine/_mlops.py index 515633a79..3c8b60fbb 100644 --- a/packages/kumo-engine/kumo/engine/_mlops.py +++ b/packages/kumo-engine/kumo/engine/_mlops.py @@ -177,26 +177,11 @@ def raise_on_error(settings: Optional[MLOps], error: Optional[BaseException], call: str, stacklevel: int = 2) -> None: if error is None or settings is None: return - message = f'{call} finished, but MLOps failed: {full_error(error)}' + message = (f'{call} finished, but MLOps failed: ' + f'{type(error).__name__}: {error}') if settings.on_error == 'raise': raise RuntimeError(message) from error - warn_failure(message, error, stacklevel + 1) - - -def full_error(error: BaseException) -> str: - return f'{type(error).__name__}: {error}' - - -def warn_failure(message: str, error: BaseException, stacklevel: int = 2, - also: Optional[BaseException] = None) -> None: - if getattr(error, '_kumo_warned', False): - return warnings.warn(message, UserWarning, stacklevel=stacklevel + 1) - for reported in (error, also): - try: - setattr(reported, '_kumo_warned', True) - except AttributeError: - pass def on_root(source: Any, path: str) -> Tuple[Any, str]: diff --git a/packages/kumo-engine/kumo/engine/pquery/training_table.py b/packages/kumo-engine/kumo/engine/pquery/training_table.py index 133319d0d..73e838bee 100644 --- a/packages/kumo-engine/kumo/engine/pquery/training_table.py +++ b/packages/kumo-engine/kumo/engine/pquery/training_table.py @@ -24,7 +24,6 @@ SourceTableEntry, TrainingTableManifest, TrainingTableParents, - embedded, read_manifest, stamp, stub, @@ -396,7 +395,7 @@ def _record_manifest(table: TrainingTable, snapshot: GraphSnapshot, TrainingTableManifest( id=table.id, parents=TrainingTableParents( - snapshot=embedded(parent) if parent is not None else + snapshot=parent if parent is not None else stub('snapshot', snapshot.id), source_table=source_table), query=table.query, diff --git a/packages/kumo-engine/kumo/engine/trainer/_provenance.py b/packages/kumo-engine/kumo/engine/trainer/_provenance.py index 4132d0196..9a92656b0 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_provenance.py +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -10,7 +10,6 @@ FitTiming, ModelManifest, ModelParents, - embedded, read_manifest, stamp, stub, @@ -38,7 +37,7 @@ def write(self, bundle_dir: Path, bundle_id: str) -> None: ModelManifest( id=bundle_id, parents=ModelParents( - training_table=embedded(table) if table is not None else + training_table=table if table is not None else stub('training_table', train_table.id or train_table.tt_id), warm_start=None if self.warm_start is None else @@ -56,15 +55,14 @@ def write(self, bundle_dir: Path, bundle_id: str) -> None: def embedded_model(manifest: ModelManifest) -> ModelManifest: - copy = embedded(manifest) - parents = copy.parents + parents = manifest.parents if parents is None: - return copy + return manifest def cut(model: Optional[ModelManifest]) -> Optional[ModelManifest]: return None if model is None else stub('model', model.id) - return replace(copy, parents=replace( + return replace(manifest, parents=replace( parents, warm_start=cut(parents.warm_start))) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 97f8de2c4..f29f4aae3 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -131,12 +131,9 @@ def snapshot_id(self) -> Optional[str]: def predict(self, *args: Any, **kwargs: Any) -> 'PredictionResult': """Score with this model, as :meth:`Trainer.predict`.""" - return self._model('predict').predict(*args, **kwargs) - - def _model(self, call: str) -> 'Trainer': trainer = self._trainer if trainer is None: - raise RuntimeError(f'This result has no model to {call} with; ' + raise RuntimeError('This result has no model to predict with; ' 'use Trainer.load on its bundle_dir.') current = trainer._training_result if current is None or current.bundle_dir != self.bundle_dir: @@ -144,7 +141,7 @@ def _model(self, call: str) -> 'Trainer': trainer._embedding_models, accelerator=trainer._accelerator) self._trainer = trainer - return trainer + return trainer.predict(*args, **kwargs) def holdout_df(self) -> pd.DataFrame: r"""Read the holdout predictions as a pandas DataFrame. diff --git a/packages/kumo-engine/test/test_artifacts.py b/packages/kumo-engine/test/test_artifacts.py deleted file mode 100644 index d6cbb3cb2..000000000 --- a/packages/kumo-engine/test/test_artifacts.py +++ /dev/null @@ -1,62 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. -# All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -from dataclasses import dataclass -from typing import List, Optional - -from kumo.engine._manifest import ( - Artifact, - Manifest, - ModelManifest, - ModelParents, - Record, - SnapshotManifest, - SnapshotParents, - SourceTableEntry, - TrainingTableManifest, - artifacts, -) - - -@dataclass(frozen=True, kw_only=True) -class ReportParents(Record): - snapshot: SnapshotManifest - model: ModelManifest - - -@dataclass(frozen=True, kw_only=True) -class ReportManifest(Manifest): - kind: str = 'report' - parents: ReportParents - pages: Optional[List[SnapshotManifest]] = None - - IS_DATASET = True - - -def test_a_new_kind_names_its_artifacts_without_other_changes(): - orders = SourceTableEntry(connector='S3', location='s3://data/orders', - version='7', version_kind='exact') - snapshot = SnapshotManifest( - id='g@d', location='s3://root/snapshots/g/d', - parents=SnapshotParents(source_tables={'orders': orders})) - model = ModelManifest( - id='m', location='s3://root/models/m', parents=ModelParents( - training_table=TrainingTableManifest(id='t'))) - report = ReportManifest( - id='r', location='s3://root/reports/r', - parents=ReportParents(snapshot=snapshot, model=model), - pages=[snapshot, SnapshotManifest(id='p', reconstructed=True)]) - - assert artifacts(report) == [ - Artifact('report', 'r', 's3://root/reports/r', True, - version='r', version_kind='exact'), - Artifact('snapshot', 'g@d', 's3://root/snapshots/g/d', True, - version='g@d', version_kind='exact'), - Artifact('source_table', 'orders', 's3://data/orders', True, - connector='S3', version='7', version_kind='exact'), - Artifact('model', 'm', 's3://root/models/m', False, - 'kumo.bundle_uri', version='m', version_kind='exact'), - Artifact('snapshot', 'p', None, True, version='p', - version_kind='exact'), - ] From d4b8ec44ea2c78505f1e6211217fb8cf5329d059 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 3 Nov 2026 00:11:34 +0000 Subject: [PATCH 267/278] Share the tracking URI, fit data and quick plan fixtures across tracking tests Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/test/conftest.py | 50 ++++++++++++ packages/kumo-engine/test/test_fit_events.py | 44 ++-------- .../kumo-engine/test/test_mlflow_model.py | 50 ++++-------- .../kumo-engine/test/test_mlflow_target.py | 9 --- .../kumo-engine/test/test_mlflow_tracking.py | 80 +++++-------------- .../test/test_snowflake_tracking.py | 11 +-- .../test/test_tracking_contract.py | 28 ++----- 7 files changed, 97 insertions(+), 175 deletions(-) diff --git a/packages/kumo-engine/test/conftest.py b/packages/kumo-engine/test/conftest.py index 4b60addea..28e0791fb 100644 --- a/packages/kumo-engine/test/conftest.py +++ b/packages/kumo-engine/test/conftest.py @@ -164,3 +164,53 @@ def fitted(tmp_path_factory): yield trainer, pquery.generate_prediction_table(snapshot), snapshot finally: kumo.configure(scratch_dir=previous) + + +@pytest.fixture +def tracking_uri(tmp_path, monkeypatch): + mlflow = pytest.importorskip('mlflow') + monkeypatch.chdir(tmp_path) + previous = mlflow.get_tracking_uri() + uri = f'sqlite:///{tmp_path / "mlflow.db"}' + mlflow.set_tracking_uri(uri) + try: + yield uri + finally: + while mlflow.active_run() is not None: + mlflow.end_run() + mlflow.set_tracking_uri(previous) + + +@pytest.fixture +def table_and_snapshot(tmp_path): + previous_scratch = get_scratch_dir() + kumo.configure(scratch_dir=tmp_path / 'scratch') + try: + rng = np.random.default_rng(0) + (tmp_path / 'data').mkdir() + pd.DataFrame({ + 'ID': np.arange(400), + 'Y': rng.random(400), + 'F1': rng.normal(size=400), + 'amount ($)': rng.normal(size=400), + }).to_parquet(tmp_path / 'data' / 'ents.parquet') + connector = kumo.FileConnector(root_dir=str(tmp_path / 'data')) + ents = kumo.Table.from_source_table(connector['ents.parquet'], + primary_key='ID').infer_metadata() + graph = kumo.Graph(tables={'ents': ents}, edges=[]) + snapshot = graph.snapshot() + query = kumo.PredictiveQuery(graph=graph, + query='PREDICT ents.Y FOR EACH ents.ID') + yield query.generate_training_table(snapshot), snapshot + finally: + kumo.configure(scratch_dir=previous_scratch) + + +@pytest.fixture +def quick_plan(): + plan = ModelPlan() + plan.training_job.num_experiments = 1 + plan.training_job.disable_explain = True + plan.optimization.max_epochs = 2 + plan.optimization.max_steps_per_epoch = 2 + return plan diff --git a/packages/kumo-engine/test/test_fit_events.py b/packages/kumo-engine/test/test_fit_events.py index 8f43b35f0..ed70a3a36 100644 --- a/packages/kumo-engine/test/test_fit_events.py +++ b/packages/kumo-engine/test/test_fit_events.py @@ -6,8 +6,6 @@ import logging import threading -import numpy as np -import pandas as pd import pytest from kumoapi.model_plan import ModelPlan @@ -16,32 +14,8 @@ from kumo.engine.tracking import Event -@pytest.fixture -def table_and_snapshot(tmp_path): - previous_scratch = get_scratch_dir() - kumo.configure(scratch_dir=tmp_path / 'scratch') - try: - rng = np.random.default_rng(0) - (tmp_path / 'data').mkdir() - pd.DataFrame({ - 'ID': np.arange(400), - 'Y': rng.random(400), - 'F1': rng.normal(size=400), - }).to_parquet(tmp_path / 'data' / 'ents.parquet') - connector = kumo.FileConnector(root_dir=str(tmp_path / 'data')) - ents = kumo.Table.from_source_table(connector['ents.parquet'], - primary_key='ID').infer_metadata() - graph = kumo.Graph(tables={'ents': ents}, edges=[]) - snapshot = graph.snapshot() - query = kumo.PredictiveQuery(graph=graph, - query='PREDICT ents.Y FOR EACH ents.ID') - yield query.generate_training_table(snapshot), snapshot - finally: - kumo.configure(scratch_dir=previous_scratch) - - -def test_fit_reports_progress_events(table_and_snapshot, monkeypatch, - caplog): +def test_fit_reports_progress_events(table_and_snapshot, quick_plan, + monkeypatch, caplog): from kumo_core.automl.optuna_searcher import TPESearcher update = TPESearcher.update @@ -55,11 +29,8 @@ def fail_second_trial(self, *args, **kwargs): monkeypatch.setattr(TPESearcher, '_num_updates', 0, raising=False) monkeypatch.setattr(TPESearcher, 'update', fail_second_trial) table, snapshot = table_and_snapshot - plan = ModelPlan() + plan = quick_plan plan.training_job.num_experiments = 2 - plan.training_job.disable_explain = True - plan.optimization.max_epochs = 2 - plan.optimization.max_steps_per_epoch = 2 events = [] broken_calls = [] @@ -124,7 +95,7 @@ def fail(*args, **kwargs): @pytest.mark.parametrize('on_error', ['warn', 'raise']) def test_an_unfinished_tracker_is_reported_by_the_error_policy( - table_and_snapshot, tmp_path, monkeypatch, on_error): + table_and_snapshot, quick_plan, tmp_path, monkeypatch, on_error): import kumo.engine._tracking as tracking monkeypatch.setattr(tracking, 'FINISH_SECONDS', 0.5) @@ -141,12 +112,7 @@ def finish(self, record, error): release.wait(30) table, snapshot = table_and_snapshot - plan = ModelPlan() - plan.training_job.num_experiments = 1 - plan.training_job.disable_explain = True - plan.optimization.max_epochs = 2 - plan.optimization.max_steps_per_epoch = 2 - trainer = kumo.Trainer(plan) + trainer = kumo.Trainer(quick_plan) cfg = kumo.MLOps(str(tmp_path / 'root'), tracking=Blocked(), store=False, on_error=on_error) try: diff --git a/packages/kumo-engine/test/test_mlflow_model.py b/packages/kumo-engine/test/test_mlflow_model.py index 4f2885104..bd5956651 100644 --- a/packages/kumo-engine/test/test_mlflow_model.py +++ b/packages/kumo-engine/test/test_mlflow_model.py @@ -10,13 +10,8 @@ import pandas as pd import pytest from kumo_core.artifacts import PROVENANCE_FILE -from kumoapi.model_plan import ModelPlan from kumoapi.plugin import PublishedModel -from test_mlflow_tracking import ( # noqa: F401 - _runs, - table_and_snapshot, - tracking_uri, -) +from test_mlflow_tracking import _fit_run import kumo.engine as kumo from kumo.engine._scratch import get_scratch_dir @@ -24,15 +19,7 @@ from kumo.engine.trainer._run_spec import registry_metrics mlflow = pytest.importorskip('mlflow') - - -def _plan(): - plan = ModelPlan() - plan.training_job.num_experiments = 1 - plan.training_job.disable_explain = True - plan.optimization.max_epochs = 2 - plan.optimization.max_steps_per_epoch = 2 - return plan +pytestmark = pytest.mark.usefixtures('tracking_uri') def _sorted(frame): @@ -40,24 +27,19 @@ def _sorted(frame): def test_fit_logs_registers_and_serves_the_model(table_and_snapshot, - tmp_path, caplog, - monkeypatch): + quick_plan, tmp_path, caplog): from kumo.engine.integrations.mlflow import MlflowTracking, load_model table, snapshot = table_and_snapshot - plan = _plan() cfg = kumo.MLOps(str(tmp_path / 'root'), tracking=MlflowTracking(log_system_metrics=False), registry=kumo.ModelRegistryLocation('m', alias='champion')) - trainer = kumo.Trainer(plan) + trainer = kumo.Trainer(quick_plan) result = trainer.fit(table, snapshot=snapshot, mlops=cfg) assert 'MLflow tracking failed' not in caplog.text - parent, = [ - run for run in _runs(mlflow.MlflowClient()) - if 'mlflow.parentRunId' not in run.data.tags - ] + parent = _fit_run() model, = mlflow.search_logged_models(output_format='list') assert (model.name, model.source_run_id) == ('m', parent.info.run_id) version = mlflow.MlflowClient().get_model_version_by_alias('m', 'champion') @@ -124,7 +106,7 @@ def other_bundle(path): load_model('models:/m@champion') -def test_logged_model_scores_given_rows(fitted, tracking_uri, tmp_path): +def test_logged_model_scores_given_rows(fitted, tmp_path): from kumo.engine.integrations.mlflow import log_model trainer, table, snapshot = fitted @@ -244,8 +226,7 @@ def test_scores_pair_with_their_rows_when_scoring_reorders(): ('mlflow', kumo.ModelRegistryLocation('m', alias='SageMakerLifecycle-p')), ('mlflow', kumo.ModelRegistryLocation('m', version='3')), ]) -def test_invalid_mlflow_registry_is_refused(tracking_uri, tmp_path, tracking, - registry): +def test_invalid_mlflow_registry_is_refused(tmp_path, tracking, registry): from kumo.engine.integrations.mlflow import MlflowTracking if tracking == 'mlflow': @@ -342,7 +323,7 @@ def finish(self, record, error): def test_model_upload_gets_a_longer_http_timeout_and_a_late_one_is_dropped( - tracking_uri, monkeypatch, caplog): + monkeypatch, caplog): import os from test_tracking_contract import CONTEXT @@ -428,7 +409,7 @@ def finish(self, record, error): @pytest.mark.parametrize('on_error', ['warn', 'raise']) def test_a_registration_failure_is_tagged_on_the_fit_run_and_can_raise( - table_and_snapshot, tmp_path, monkeypatch, on_error): + table_and_snapshot, quick_plan, tmp_path, monkeypatch, on_error): from mlflow.exceptions import MlflowException from mlflow.tracking import MlflowClient @@ -447,7 +428,7 @@ def missing_catalog(*args, **kwargs): tracking=MlflowTracking(log_system_metrics=False), registry=kumo.ModelRegistryLocation('nope.s.m'), on_error=on_error) - trainer = kumo.Trainer(_plan()) + trainer = kumo.Trainer(quick_plan) cause = f'MlflowException: {detail}' with warnings.catch_warnings(record=True) as caught: warnings.simplefilter('always') @@ -464,8 +445,7 @@ def missing_catalog(*args, **kwargs): if on_error == 'warn' else []) assert result.id and result.registered_model is None - run, = [run for run in _runs(mlflow.MlflowClient()) - if 'mlflow.parentRunId' not in run.data.tags] + run = _fit_run() assert run.info.status == 'FAILED' assert run.data.tags['kumo.error'].startswith('MlflowException: ' 'CATALOG_DOES_NOT_EXIST') @@ -473,7 +453,7 @@ def missing_catalog(*args, **kwargs): def test_a_tracking_start_failure_names_its_cause_and_can_raise( - table_and_snapshot, tmp_path, monkeypatch, caplog): + table_and_snapshot, quick_plan, tmp_path, monkeypatch, caplog): from mlflow.exceptions import MlflowException from mlflow.tracking import MlflowClient @@ -492,7 +472,7 @@ def no_folder(self, name, *args, **kwargs): monkeypatch.setattr(MlflowClient, 'get_experiment_by_name', lambda self, name: None) monkeypatch.setattr(MlflowClient, 'create_experiment', no_folder) - trainer = kumo.Trainer(_plan()) + trainer = kumo.Trainer(quick_plan) with pytest.raises(RuntimeError, match='RESOURCE_DOES_NOT_EXIST') as raised: trainer.fit(table, snapshot=snapshot, mlops=cfg) @@ -529,12 +509,12 @@ def test_a_logged_model_is_named_after_its_model(): def test_a_later_registration_logs_the_model_in_its_fit_run_else_in_none( - table_and_snapshot, tmp_path, tracking_uri): + table_and_snapshot, quick_plan, tmp_path): from kumo.engine.integrations.mlflow import MlflowTracking table, snapshot = table_and_snapshot tracking = MlflowTracking('fits', log_system_metrics=False) - result = kumo.Trainer(_plan()).fit(table, snapshot=snapshot, mlops=( + result = kumo.Trainer(quick_plan).fit(table, snapshot=snapshot, mlops=( kumo.MLOps(str(tmp_path / 'root'), tracking=tracking))) fits = [mlflow.get_experiment_by_name('fits').experiment_id] assert mlflow.MlflowClient().search_logged_models(fits) == [] diff --git a/packages/kumo-engine/test/test_mlflow_target.py b/packages/kumo-engine/test/test_mlflow_target.py index 280e5ce6d..51e3e04bb 100644 --- a/packages/kumo-engine/test/test_mlflow_target.py +++ b/packages/kumo-engine/test/test_mlflow_target.py @@ -24,15 +24,6 @@ ARN = 'arn:aws:sagemaker:us-west-2:account:mlflow-tracking-server/demo' -@pytest.fixture -def tracking_uri(tmp_path): - previous = mlflow.get_tracking_uri() - uri = f'sqlite:///{tmp_path / "mlflow.db"}' - mlflow.set_tracking_uri(uri) - yield uri - mlflow.set_tracking_uri(previous) - - def _connector(identity): return DatabricksConnector(catalog='main', schema_name='ml', identity=identity, warehouse_id='w') diff --git a/packages/kumo-engine/test/test_mlflow_tracking.py b/packages/kumo-engine/test/test_mlflow_tracking.py index 413ea54ab..48c9156ca 100644 --- a/packages/kumo-engine/test/test_mlflow_tracking.py +++ b/packages/kumo-engine/test/test_mlflow_tracking.py @@ -11,8 +11,6 @@ import warnings from pathlib import Path -import numpy as np -import pandas as pd import pytest from kumo_core.artifacts import PROVENANCE_FILE from kumo_core.summaries import Dataset @@ -27,52 +25,12 @@ import kumo.engine as kumo from kumo.engine._manifest import TrackingRef -from kumo.engine._scratch import get_scratch_dir from kumo.engine._tracking import logged_names, run_record from kumo.engine.trainer._run_spec import fit_spec mlflow = pytest.importorskip('mlflow') autologging_utils = pytest.importorskip('mlflow.utils.autologging_utils') - -QUERY = 'PREDICT ents.Y FOR EACH ents.ID' - - -@pytest.fixture -def tracking_uri(tmp_path, monkeypatch): - monkeypatch.chdir(tmp_path) - previous = mlflow.get_tracking_uri() - uri = f'sqlite:///{tmp_path / "mlflow.db"}' - mlflow.set_tracking_uri(uri) - try: - yield uri - finally: - while mlflow.active_run() is not None: - mlflow.end_run() - mlflow.set_tracking_uri(previous) - - -@pytest.fixture -def table_and_snapshot(tmp_path, tracking_uri): - previous_scratch = get_scratch_dir() - kumo.configure(scratch_dir=tmp_path / 'scratch') - try: - rng = np.random.default_rng(0) - (tmp_path / 'data').mkdir() - pd.DataFrame({ - 'ID': np.arange(400), - 'Y': rng.random(400), - 'F1': rng.normal(size=400), - 'amount ($)': rng.normal(size=400), - }).to_parquet(tmp_path / 'data' / 'ents.parquet') - connector = kumo.FileConnector(root_dir=str(tmp_path / 'data')) - ents = kumo.Table.from_source_table(connector['ents.parquet'], - primary_key='ID').infer_metadata() - graph = kumo.Graph(tables={'ents': ents}, edges=[]) - snapshot = graph.snapshot() - query = kumo.PredictiveQuery(graph=graph, query=QUERY) - yield query.generate_training_table(snapshot), snapshot - finally: - kumo.configure(scratch_dir=previous_scratch) +pytestmark = pytest.mark.usefixtures('tracking_uri') def _runs(client): @@ -80,17 +38,20 @@ def _runs(client): [e.experiment_id for e in client.search_experiments()]) -def test_fit_logs_a_run_per_fit_and_per_trial(table_and_snapshot, tmp_path, - caplog): +def _fit_run(): + run, = [run for run in _runs(mlflow.MlflowClient()) + if 'mlflow.parentRunId' not in run.data.tags] + return run + + +def test_fit_logs_a_run_per_fit_and_per_trial(table_and_snapshot, quick_plan, + tmp_path, caplog): from kumo.engine.integrations.mlflow import MlflowTracking table, snapshot = table_and_snapshot client = mlflow.MlflowClient() - plan = ModelPlan() + plan = quick_plan plan.training_job.num_experiments = 2 - plan.training_job.disable_explain = True - plan.optimization.max_epochs = 2 - plan.optimization.max_steps_per_epoch = 2 plan.optimization.base_lr = [0.01, 0.001] plan.column_processing.encoder_overrides = { 'ents.amount ($)': Numerical('standard') @@ -207,8 +168,7 @@ def fail(*args, **kwargs): assert mlflow.active_run() is None -def test_mlflow_error_ends_runs_failed_and_is_not_raised( - tracking_uri, monkeypatch, caplog): +def test_mlflow_error_ends_runs_failed_and_is_not_raised(monkeypatch, caplog): from mlflow.tracking import MlflowClient from kumo.engine.integrations.mlflow import MlflowTracking @@ -229,16 +189,14 @@ def down(self, *args, metrics=(), **kwargs): assert statuses == {'FAILED'} -def test_a_user_run_name_names_every_run(tracking_uri): +def test_a_user_run_name_names_every_run(): from kumo.engine.integrations.mlflow import MlflowTracking tracking = MlflowTracking(run_name='nightly', log_system_metrics=False) replay(tracking) join_tracking_threads() - names = {run.info.run_name for run in _runs(mlflow.MlflowClient()) - if 'mlflow.parentRunId' not in run.data.tags} - assert names == {'fit-nightly'} + assert _fit_run().info.run_name == 'fit-nightly' SAGEMAKER = ('arn:aws:sagemaker:us-west-2:account:' @@ -298,7 +256,7 @@ def refuse(self, name, *args, **kwargs): ('sqlite:///mlflow.db', None), ]) def test_a_tracked_fit_on_sagemaker_says_how_to_open_the_ui( - tracking_uri, uri, command, monkeypatch): + uri, command, monkeypatch): from kumo.engine.integrations import mlflow as integration from kumo.engine.integrations.mlflow import MlflowTracking @@ -321,7 +279,7 @@ def test_a_tracked_fit_on_sagemaker_says_how_to_open_the_ui( assert not [w for w in again if 'Studio' in str(w.message)] -def test_user_disabled_autologging_stays_disabled(tracking_uri): +def test_user_disabled_autologging_stays_disabled(): from kumo.engine.integrations.mlflow import MlflowTracking tracker = MlflowTracking(log_system_metrics=False) @@ -376,8 +334,7 @@ def test_a_hung_server_does_not_stall_the_fit(hung_server, monkeypatch): 'MLFLOW_HTTP_REQUEST_TIMEOUT': '30'} -def test_source_tables_are_logged_as_native_inputs( - tracking_uri, monkeypatch): +def test_source_tables_are_logged_as_native_inputs(monkeypatch): from mlflow.data.delta_dataset_source import DeltaDatasetSource from kumo.engine.integrations.mlflow import MlflowTracking @@ -388,8 +345,7 @@ def test_source_tables_are_logged_as_native_inputs( Dataset(name, source, {}, {}, 'source table') for name, source in INPUTS.items()]) - run, = [run for run in _runs(mlflow.MlflowClient()) - if run.data.tags['kumo.kind'] == 'fit'] + run = _fit_run() inputs = { item.dataset.name: (item.dataset.source_type, json.loads(item.dataset.source)) @@ -409,7 +365,7 @@ def test_source_tables_are_logged_as_native_inputs( assert mlflow.data.get_source(orders).delta_table_id == 'table-id' -def test_link_prediction_metric_names_are_logged(tracking_uri): +def test_link_prediction_metric_names_are_logged(): from kumo.engine._tracking import RunRecord from kumo.engine.integrations.mlflow import MlflowTracking diff --git a/packages/kumo-engine/test/test_snowflake_tracking.py b/packages/kumo-engine/test/test_snowflake_tracking.py index 4ccc5ea0e..dbbd4b33a 100644 --- a/packages/kumo-engine/test/test_snowflake_tracking.py +++ b/packages/kumo-engine/test/test_snowflake_tracking.py @@ -10,8 +10,6 @@ import _snowflake_experiment_fake import pytest from _snowflake_experiment_fake import FakeExperimentTracking, FakeSession -from kumoapi.model_plan import ModelPlan -from test_fit_events import table_and_snapshot # noqa: F401 from test_tracking_contract import CONTEXT, EVENTS, RESULT from kumo.engine._tracking import RunRecord, TrackingSession @@ -110,7 +108,7 @@ def test_failed_fit_ends_the_run_failed(fake): @pytest.mark.parametrize('busy', ['active run', 'other session']) def test_an_instance_in_use_is_left_alone_and_fails_strict_mode( - fake, busy, table_and_snapshot, tmp_path, monkeypatch): + fake, busy, table_and_snapshot, quick_plan, tmp_path, monkeypatch): import kumo.engine as kumo from kumo.engine.integrations.snowflake import SnowflakeTracking @@ -124,12 +122,7 @@ def test_an_instance_in_use_is_left_alone_and_fails_strict_mode( session = FakeSession() calls, run = len(user.calls), user._run table, snapshot = table_and_snapshot - plan = ModelPlan() - plan.training_job.num_experiments = 1 - plan.training_job.disable_explain = True - plan.optimization.max_epochs = 2 - plan.optimization.max_steps_per_epoch = 2 - trainer = kumo.Trainer(plan) + trainer = kumo.Trainer(quick_plan) in_use = 'Snowflake ExperimentTracking is already in use' with pytest.raises(RuntimeError, match=in_use): diff --git a/packages/kumo-engine/test/test_tracking_contract.py b/packages/kumo-engine/test/test_tracking_contract.py index 0b24970a8..92b2fb36e 100644 --- a/packages/kumo-engine/test/test_tracking_contract.py +++ b/packages/kumo-engine/test/test_tracking_contract.py @@ -1,7 +1,6 @@ # 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.""" import hashlib import json @@ -126,7 +125,6 @@ def join_tracking_threads(): def replay(tracker=None, events=EVENTS, result=RESULT, error=None, context=CONTEXT, series=None, summaries=()) -> Dict[str, RunRecord]: - r"""Feeds the recorded fit to `tracker`; returns the expected records.""" recorder = _Recorder(tracker) session = TrackingSession(recorder) session.start(lambda: fit_spec(context)) @@ -223,19 +221,15 @@ def differ(what, have, want): def _trial(run): - r"""An MLflow trial run's record name, ``trial--attempt-``.""" trial = run.info.run_name.split('-')[1] return f"trial-{trial}-attempt-{run.data.params['attempt']}" @pytest.fixture -def mlflow_adapter(tmp_path, monkeypatch): - mlflow = pytest.importorskip('mlflow') - from kumo.engine.integrations.mlflow import MlflowTracking, _mlflow_params +def mlflow_adapter(tracking_uri): + import mlflow - monkeypatch.chdir(tmp_path) - previous = mlflow.get_tracking_uri() - mlflow.set_tracking_uri(f'sqlite:///{tmp_path / "mlflow.db"}') + from kumo.engine.integrations.mlflow import MlflowTracking, _mlflow_params def observe(): client = mlflow.MlflowClient() @@ -269,16 +263,10 @@ def history(run_id, key): for run in runs } - try: - yield Adapter( - make=lambda: MlflowTracking(log_system_metrics=False, - tags=USER_TAGS), observe=observe, - kept_params=_mlflow_params, nested_runs=True, step_metrics=True, - artifacts=True) - finally: - while mlflow.active_run() is not None: - mlflow.end_run() - mlflow.set_tracking_uri(previous) + return Adapter( + make=lambda: MlflowTracking(log_system_metrics=False, tags=USER_TAGS), + observe=observe, kept_params=_mlflow_params, nested_runs=True, + step_metrics=True, artifacts=True) @pytest.fixture @@ -303,7 +291,6 @@ def bucket(key): for key, value in run['params'].items(): runs, name = bucket(key) - # Snowflake logs tags as params. if name.startswith('kumo.') or name in USER_TAGS: runs['tags'][name] = value else: @@ -323,7 +310,6 @@ def bucket(key): return observed def expected(records): - r"""What Snowflake keeps of each trial.""" kept = {} for key, record in records.items(): if record.tags['kumo.kind'] == 'trial': From 4d8795a9b15007720187cfb1931e3ba44566aa71 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 3 Nov 2026 01:04:53 +0000 Subject: [PATCH 268/278] Share the bundle and registry stubs in the Snowflake registry tests Co-Authored-By: Claude Opus 5.5 --- .../kumo-snowflake/tests/test_registry.py | 49 +++++++------------ 1 file changed, 19 insertions(+), 30 deletions(-) diff --git a/packages/kumo-snowflake/tests/test_registry.py b/packages/kumo-snowflake/tests/test_registry.py index 7d0fe7256..ac856e203 100644 --- a/packages/kumo-snowflake/tests/test_registry.py +++ b/packages/kumo-snowflake/tests/test_registry.py @@ -81,6 +81,13 @@ def test_invalid_metrics_preserve_previous_destination(fetch, payload): assert (destination / 'previous').read_bytes() == b'previous bundle' +def _bundle(tmp_path): + bundle = tmp_path / 'bundle' + bundle.mkdir() + (bundle / 'training-spec.json').write_text('{}') + return bundle + + @pytest.mark.parametrize('metadata', [{}, { 'predictive_query': 'x' * 300, 'kumo_bundle_id': 'b1', @@ -88,9 +95,7 @@ def test_invalid_metrics_preserve_previous_destination(fetch, payload): }]) def test_publish_records_pointers_and_holdout_metrics_in_catalog( tmp_path, monkeypatch, metadata): - bundle = tmp_path / 'bundle' - bundle.mkdir() - (bundle / 'training-spec.json').write_text('{}') + bundle = _bundle(tmp_path) calls = [] version = SimpleNamespace(model_name='CHURN', version_name='V1', fully_qualified_model_name='DB.SCHEMA.CHURN') @@ -127,9 +132,7 @@ def registry(_): def test_publish_names_the_next_version_and_retries_a_taken_one( tmp_path, monkeypatch): - bundle = tmp_path / 'bundle' - bundle.mkdir() - (bundle / 'training-spec.json').write_text('{}') + bundle = _bundle(tmp_path) names = ['V1', 'SLIMY_DONKEY_1', 'V3'] tried = [] @@ -143,13 +146,7 @@ def log_model(*args, version_name, **kwargs): model = SimpleNamespace(versions=lambda: [ SimpleNamespace(version_name=name) for name in names]) - - @contextmanager - def registry(_): - yield SimpleNamespace(log_model=log_model, get_model=lambda _: model) - - sink = SnowflakeModelRegistryMixin() - monkeypatch.setattr(sink, '_registry', registry) + sink = _sink(monkeypatch, log_model=log_model, get_model=lambda _: model) monkeypatch.setattr(registry_module, '_bundle_model', lambda _: object()) monkeypatch.setattr(registry_module, '_carrier_signature', lambda: {}) published = sink.publish_model(SimpleNamespace(), str(bundle), @@ -166,13 +163,9 @@ def fail_after_creating(*args, version_name, **kwargs): names.append(version_name) raise RuntimeError('upload failed') - @contextmanager - def failing(_): - yield SimpleNamespace(log_model=fail_after_creating, - get_model=lambda _: model) - tried.clear() - monkeypatch.setattr(sink, '_registry', failing) + sink = _sink(monkeypatch, log_model=fail_after_creating, + get_model=lambda _: model) with pytest.raises(RuntimeError, match='upload failed'): sink.publish_model(SimpleNamespace(), str(bundle), location=ModelRegistryLocation('CHURN')) @@ -244,10 +237,10 @@ def export(directory): }) -def _aliasing_sink(monkeypatch, model): +def _sink(monkeypatch, **registry_methods): @contextmanager def registry(_): - yield SimpleNamespace(get_model=lambda _: model) + yield SimpleNamespace(**registry_methods) sink = SnowflakeModelRegistryMixin() monkeypatch.setattr(sink, '_registry', registry) @@ -260,7 +253,7 @@ def test_set_alias_moves_it_and_refuses_what_snowflake_cannot(monkeypatch): assert isinstance(SnowflakeSinkBackend(), ModelAliasingSink) model = _Model({'V1': ['FIRST'], 'V2': ['LAST', 'DEFAULT'], 'V3': []}) - sink = _aliasing_sink(monkeypatch, model) + sink = _sink(monkeypatch, get_model=lambda _: model) def set_alias(version, alias): sink.set_model_alias(SimpleNamespace(), @@ -284,10 +277,8 @@ def set_alias(version, alias): def test_fetch_by_alias_reads_the_aliased_version(monkeypatch, tmp_path): - sink = _aliasing_sink(monkeypatch, _Model({ - 'V1': ['DEFAULT'], - 'V2': ['PROD'] - })) + model = _Model({'V1': ['DEFAULT'], 'V2': ['PROD']}) + sink = _sink(monkeypatch, get_model=lambda _: model) published = sink.fetch_model( SimpleNamespace(), str(tmp_path / 'out'), location=ModelRegistryLocation('CHURN', alias='prod')) @@ -312,7 +303,7 @@ def failing_version(name): return version model.version = failing_version - sink = _aliasing_sink(monkeypatch, model) + sink = _sink(monkeypatch, get_model=lambda _: model) with pytest.raises(RuntimeError, match='lost connection'): sink.set_model_alias(SimpleNamespace(), ModelRegistryLocation('CHURN', version='V2'), @@ -328,9 +319,7 @@ def test_publish_records_the_task_type_monitors_need(tmp_path, monkeypatch, task, expected): import sys - bundle = tmp_path / 'bundle' - bundle.mkdir() - (bundle / 'training-spec.json').write_text('{}') + bundle = _bundle(tmp_path) (bundle / 'prediction-schema.json').write_text( json.dumps({'task_type': task})) tasks = SimpleNamespace( From e19445a62687630803879c289688523417e1231c Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 3 Nov 2026 01:24:54 +0000 Subject: [PATCH 269/278] Merge the Snowflake run-name tests and reuse FakeSession for schema checks Co-Authored-By: Claude Opus 5.5 --- .../test/_snowflake_experiment_fake.py | 11 ++-- .../test/test_snowflake_tracking.py | 55 ++++++------------- 2 files changed, 21 insertions(+), 45 deletions(-) diff --git a/packages/kumo-engine/test/_snowflake_experiment_fake.py b/packages/kumo-engine/test/_snowflake_experiment_fake.py index ff150009c..ac25b360c 100644 --- a/packages/kumo-engine/test/_snowflake_experiment_fake.py +++ b/packages/kumo-engine/test/_snowflake_experiment_fake.py @@ -1,13 +1,12 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""A recording stand-in for Snowflake's ExperimentTracking.""" import os import re import sys import types -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Dict, List, Optional, Sequence, Tuple from kumo_core.sources.sql import resolve_snowflake_identifier as _identifier @@ -157,7 +156,6 @@ def log_metrics(self, metrics: Dict[str, float], step: int = 0) -> None: def log_artifact(self, local_path: str, artifact_path: Optional[str] = None) -> None: self._record('log_artifact', local_path, artifact_path) - # Artifact URIs do not quote identifiers. if not re.fullmatch(r'[A-Z_][A-Z0-9_$]*', self._run.name): raise RuntimeError(f'Run "{self._run.name}" does not exist in ' f'Experiment {self.experiment}.') @@ -169,13 +167,15 @@ def log_artifact(self, local_path: str, class FakeSession: def __init__(self, database: Optional[str] = 'DB', - schema: Optional[str] = 'PUBLIC') -> None: + schema: Optional[str] = 'PUBLIC', + schemas: Sequence[Tuple[str]] = ()) -> None: self.database, self.schema = database, schema + self.schemas = list(schemas) self.statements: List[str] = [] def sql(self, query: str) -> Any: self.statements.append(query) - return types.SimpleNamespace(collect=lambda: []) + return types.SimpleNamespace(collect=lambda: self.schemas) def get_current_database(self) -> Optional[str]: return self.database @@ -185,7 +185,6 @@ def get_current_schema(self) -> Optional[str]: def install(monkeypatch: Any) -> type: - """Makes `import snowflake.ml.experiment` return a fresh fake.""" FakeExperimentTracking._instance = None module = types.ModuleType('snowflake.ml.experiment') module.ExperimentTracking = FakeExperimentTracking diff --git a/packages/kumo-engine/test/test_snowflake_tracking.py b/packages/kumo-engine/test/test_snowflake_tracking.py index dbbd4b33a..e1c094a8c 100644 --- a/packages/kumo-engine/test/test_snowflake_tracking.py +++ b/packages/kumo-engine/test/test_snowflake_tracking.py @@ -157,7 +157,6 @@ def test_an_existing_run_is_never_resumed(fake, taken, run_name, new, caplog): user.end_run() else: user._run = None - run_name = 'fit' if taken == 'FIT' else 'my_run' tracking = _tracking(session, run_name=run_name) _replay(tracking) @@ -198,35 +197,24 @@ def test_the_users_session_is_not_altered(fake): assert session.statements == [] -class _SchemaSession: - def __init__(self, schemas): - self.schemas = schemas - self.statements = [] - - def get_current_database(self): - return 'DB' - - def get_current_schema(self): - return 'RUNS' +def _connected(): + from kumo.engine.connector.base import Connector - def sql(self, statement): - self.statements.append(statement) - return mock.Mock(collect=lambda: self.schemas) + session = FakeSession(schema='RUNS', schemas=[('RUNS', )]) + connector = mock.Mock(spec=Connector, session=lambda: session) + return _tracking(connector), connector, session def test_a_connector_opens_the_session_and_the_schema_is_checked_first( fake): import kumo.engine as kumo - from kumo.engine.connector.base import Connector - session = _SchemaSession([('RUNS', )]) - connector = mock.Mock(spec=Connector, session=lambda: session) - tracking = _tracking(connector) + tracking, connector, session = _connected() tracking.check() assert tracking.session is session and tracking.connector is connector assert session.statements == ["SHOW SCHEMAS LIKE 'RUNS' IN DATABASE DB"] - tracking.session = _SchemaSession([]) + tracking.session = FakeSession(schema='RUNS') with pytest.raises(ValueError, match='schema DB.RUNS does not exist'): kumo.MLOps('/tmp/root', tracking=tracking) with pytest.raises(ValueError, match='or a Snowpark session'): @@ -236,7 +224,6 @@ def test_a_connector_opens_the_session_and_the_schema_is_checked_first( def test_mlops_loads_a_registered_model_through_the_tracker( fake, monkeypatch, tmp_path): import kumo.engine as kumo - from kumo.engine.connector.base import Connector loaded = [] original = kumo.Trainer.load_from_connector.__func__ @@ -245,9 +232,7 @@ def test_mlops_loads_a_registered_model_through_the_tracker( classmethod(lambda cls, source, **kwargs: loaded.append( (source, kwargs)) if not isinstance(source, kumo.MLOps) else original(cls, source, **kwargs))) - session = _SchemaSession([('RUNS', )]) - connector = mock.Mock(spec=Connector, session=lambda: session) - tracking = _tracking(connector) + tracking, connector, _ = _connected() cfg = kumo.MLOps(str(tmp_path), tracking=tracking) location = kumo.ModelRegistryLocation('CHURN', alias='champion') kumo.Trainer.load_from_connector(cfg, location=location) @@ -259,24 +244,16 @@ def test_mlops_loads_a_registered_model_through_the_tracker( _tracking().load_model(location) -@pytest.mark.parametrize('record, name', [ - (RunRecord(tags={'kumo.kind': 'fit'}, model_label='kumo-demo-churn', - started=datetime(2026, 10, 29, 6, 18, 5, tzinfo=timezone.utc)), +@pytest.mark.parametrize('run_name, record, name', [ + (None, RunRecord(tags={'kumo.kind': 'fit'}, model_label='kumo-demo-churn', + started=datetime(2026, 10, 29, 6, 18, 5, + tzinfo=timezone.utc)), 'FIT_KUMO_DEMO_CHURN_20261029_061805'), + ('nightly', RunRecord(tags={'kumo.kind': 'fit'}), 'FIT_NIGHTLY'), ]) -def test_runs_are_named_by_kind_model_version_and_date(fake, record, name): - tracking = _tracking() - tracking.start(record) - tracking.finish(record, None) - - assert list(fake._instance.runs) == [name] - - -@pytest.mark.parametrize('record, name', [ - (RunRecord(tags={'kumo.kind': 'fit'}), 'FIT_NIGHTLY'), -]) -def test_a_user_run_name_names_every_run(fake, record, name): - tracking = _tracking(run_name='nightly') +def test_runs_are_named_by_kind_model_and_date_or_the_users_name( + fake, run_name, record, name): + tracking = _tracking(run_name=run_name) tracking.start(record) tracking.finish(record, None) From b05fbbb6062f758c3fbec3fe8e132a07a2269328 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 3 Nov 2026 01:24:54 +0000 Subject: [PATCH 270/278] Share the vocab asset and volume stubs in the model-agnostic MLOps tests Co-Authored-By: Claude Opus 5.5 --- .../test/test_model_agnostic_mlops.py | 32 +++++++++---------- 1 file changed, 16 insertions(+), 16 deletions(-) diff --git a/packages/kumo-engine/test/test_model_agnostic_mlops.py b/packages/kumo-engine/test/test_model_agnostic_mlops.py index c6b3de037..113ec0015 100644 --- a/packages/kumo-engine/test/test_model_agnostic_mlops.py +++ b/packages/kumo-engine/test/test_model_agnostic_mlops.py @@ -1,7 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""The MLOps seams take any bundle, run and predictor.""" import io import json @@ -168,10 +167,20 @@ def test_a_bundle_is_scored_by_the_predictor_it_names(tmp_path): load_predictor(unknown) -def test_a_file_asset_is_stored_once_and_resolved_by_its_hash(tmp_path): +def _vocab(tmp_path): (tmp_path / 'vocab').mkdir() (tmp_path / 'vocab' / 'words.txt').write_text('a b c') - asset = FileAsset('toy', 'vocab', tmp_path / 'vocab', version='1') + return tmp_path / 'vocab' + + +def _volume(filesystem): + return SimpleNamespace(sink_backend=SimpleNamespace( + holds_files=lambda connector, path: path.startswith('/Volumes/'), + filesystem=filesystem)) + + +def test_a_file_asset_is_stored_once_and_resolved_by_its_hash(tmp_path): + asset = FileAsset('toy', 'vocab', _vocab(tmp_path), version='1') settings = kumo.MLOps(str(tmp_path / 'root')) recorded = Dependencies(settings=settings) record = recorded.require(asset) @@ -288,7 +297,6 @@ def test_a_model_serves_from_its_bundle_snapshot_and_assets_alone( class _PrefixOnlyS3: - r"""S3 credentials that cover only the root's prefix.""" type_name = 's3' def __init__(self): @@ -321,10 +329,8 @@ def test_an_asset_is_stored_on_s3_with_prefix_scoped_rights( s3 = _PrefixOnlyS3() monkeypatch.setattr(connector, 'file_store', lambda connector, path: (s3, f'bucket/kumo/{path}')) - (tmp_path / 'vocab').mkdir() - (tmp_path / 'vocab' / 'words.txt').write_text('a b c') deps = Dependencies(settings=kumo.MLOps(str(tmp_path / 'root'))) - record = deps.require(FileAsset('toy', 'vocab', tmp_path / 'vocab')) + record = deps.require(FileAsset('toy', 'vocab', _vocab(tmp_path))) assert deps.errors == [] and record.location assert sorted(s3.objects) == [ f'bucket/kumo/assets/toy/{record.sha256}/_provenance.json', @@ -378,13 +384,9 @@ def test_volume_assets_load_through_the_connector_the_loader_was_given( monkeypatch.setenv('XDG_CACHE_HOME', str(tmp_path / 'cache')) (tmp_path / 'volume').mkdir() files = SubTreeFileSystem(str(tmp_path / 'volume'), LocalFileSystem()) - volume = SimpleNamespace(sink_backend=SimpleNamespace( - holds_files=lambda connector, path: path.startswith('/Volumes/'), - filesystem=lambda connector, path: (files, path.lstrip('/')))) - (tmp_path / 'vocab').mkdir() - (tmp_path / 'vocab' / 'words.txt').write_text('a b c') + volume = _volume(lambda connector, path: (files, path.lstrip('/'))) recorded = Dependencies(settings=kumo.MLOps((volume, '/Volumes/c/s/v/a'))) - record = recorded.require(FileAsset('toy', 'vocab', tmp_path / 'vocab')) + record = recorded.require(FileAsset('toy', 'vocab', _vocab(tmp_path))) assert record.location.startswith('/Volumes/c/s/v/a/assets/toy/') manifest = ModelManifest(id='m', dependencies=[record]) @@ -418,9 +420,7 @@ def test_the_pyfunc_reads_a_volume_snapshot_through_assets_from( from kumo.engine.graph import GraphSnapshot from kumo.engine.integrations import _mlflow_model - volume = SimpleNamespace(sink_backend=SimpleNamespace( - holds_files=lambda connector, path: path.startswith('/Volumes/'), - filesystem=lambda connector, path: None)) + volume = _volume(lambda connector, path: None) loads = [] monkeypatch.setattr(GraphSnapshot, 'load_from_connector', lambda connector, *, path: loads.append( From b3a1ddc3009512c3739c0a38e0712e1d79dbf56f Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 3 Nov 2026 01:24:54 +0000 Subject: [PATCH 271/278] Fold the bundle-id pointer test into the publication test and share the specced bundle Co-Authored-By: Claude Opus 5.5 --- .../kumo-engine/test/test_trainer_registry.py | 44 +++++++------------ 1 file changed, 15 insertions(+), 29 deletions(-) diff --git a/packages/kumo-engine/test/test_trainer_registry.py b/packages/kumo-engine/test/test_trainer_registry.py index 13c236981..1d47e25a1 100644 --- a/packages/kumo-engine/test/test_trainer_registry.py +++ b/packages/kumo-engine/test/test_trainer_registry.py @@ -263,7 +263,8 @@ def publish_model(self, connector, bundle_dir, **kwargs): assert calls[0][2]['path'] == path assert calls[0][2]['location'] == location assert calls[0][2]['metrics'] == {'test/auroc': .8} - assert calls[0][2]['metadata']['kumo_bundle_id'] == result.bundle_id + assert {key: value for key, value in calls[0][2]['metadata'].items() + if key.startswith('kumo_')} == {'kumo_bundle_id': result.bundle_id} assert result.published_model is published assert result.metrics == original.metrics assert result.bundle_dir == bundle @@ -580,27 +581,21 @@ def fake_predict(**kwargs) -> None: assert result.frame['TARGET_PRED'].tolist() == [True] -def _export(sink, bundle, location=None, artifact_paths=None): +def _export(sink, bundle, location): return Trainer(ModelPlan())._finish_fit( - TrainingResult({}, artifact_paths or {}, - bundle), predictive_query='PREDICT COUNT(*)', + TrainingResult({}, {}, bundle), predictive_query='PREDICT COUNT(*)', snapshot_id='snap-123', training_table_generation=_generation_plan(), export_config=SimpleNamespace( - connector=SimpleNamespace(sink_backend=sink), - path=None if location else 'm', location=location)) + connector=SimpleNamespace(sink_backend=sink), path=None, + location=location)) -def _recording_sink(calls): - return SimpleNamespace(publish_model=lambda *a, **k: calls.append(k) or - PublishedModel(name='m', version='1', uri='u')) - - -def test_publish_records_only_the_bundle_id_pointer(tmp_path: Path) -> None: - calls = [] - result = _export(_recording_sink(calls), _fake_published_bundle(tmp_path)) - assert {k: v - for k, v in calls[0]['metadata'].items() - if k.startswith('kumo_')} == {'kumo_bundle_id': result.bundle_id} +def _specced_bundle(tmp_path: Path) -> Path: + bundle = _fake_published_bundle(tmp_path) + write_training_spec(bundle, predictive_query='PREDICT COUNT(*)', + model_plan=ModelPlan(), snapshot_id='snap-123', + training_table_generation=_generation_plan()) + return bundle class _AliasingSink(_FakeSink): @@ -647,10 +642,7 @@ def test_export_with_alias_publishes_then_sets_it_best_effort( def test_load_by_alias_warns_on_recorded_bundle_mismatch( tmp_path: Path) -> None: - bundle = _fake_published_bundle(tmp_path) - write_training_spec(bundle, predictive_query='PREDICT COUNT(*)', - model_plan=ModelPlan(), snapshot_id='snap-123', - training_table_generation=_generation_plan()) + bundle = _specced_bundle(tmp_path) sink = _AliasingSink('other') sink.publish_model(None, str(bundle), location=ModelRegistryLocation('CHURN_GNN')) @@ -662,10 +654,7 @@ def test_load_by_alias_warns_on_recorded_bundle_mismatch( def test_training_result_exports_like_its_trainer(tmp_path: Path) -> None: - bundle = _fake_published_bundle(tmp_path) - write_training_spec(bundle, predictive_query='PREDICT COUNT(*)', - model_plan=ModelPlan(), snapshot_id='snap-123', - training_table_generation=_generation_plan()) + bundle = _specced_bundle(tmp_path) sink = _AliasingSink() published = TrainingResult({}, {}, bundle).export( ModelOutputConfig(_FakeConnector(sink), @@ -724,10 +713,7 @@ def test_mlops_registry_registers_after_the_fit(tmp_path: Path, monkeypatch, import kumo.engine as kumo - bundle = _fake_published_bundle(tmp_path) - write_training_spec(bundle, predictive_query='PREDICT COUNT(*)', - model_plan=ModelPlan(), snapshot_id='snap-123', - training_table_generation=_generation_plan()) + bundle = _specced_bundle(tmp_path) monkeypatch.setattr( Trainer, '_fit', lambda self, *a, **k: TrainingResult( {}, {}, bundle, bundle_id='b1')) From 8dddaf9e1a248174218669e99ba721c3532ef7ad Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 3 Nov 2026 01:24:54 +0000 Subject: [PATCH 272/278] Share one recording tracker and session start in the model-logging tests Co-Authored-By: Claude Opus 5.5 --- .../kumo-engine/test/test_mlflow_model.py | 81 +++++++++---------- .../kumo-engine/test/test_mlflow_tracking.py | 2 +- 2 files changed, 37 insertions(+), 46 deletions(-) diff --git a/packages/kumo-engine/test/test_mlflow_model.py b/packages/kumo-engine/test/test_mlflow_model.py index bd5956651..e81734f95 100644 --- a/packages/kumo-engine/test/test_mlflow_model.py +++ b/packages/kumo-engine/test/test_mlflow_model.py @@ -26,6 +26,29 @@ def _sorted(frame): return frame.sort_values('ENTITY').reset_index(drop=True) +class _Tracker: + def __init__(self): + self.calls = [] + + def start(self, record): + self.calls.append('start') + + def on_trial(self, record, event): + pass + + def finish(self, record, error): + self.calls.append(('finish', record, error)) + + +def _started(tracker): + from kumo.engine._tracking import TrackingSession + from kumo.engine.trainer._run_spec import FitContext, fit_spec + + session = TrackingSession(tracker) + session.start(lambda: fit_spec(FitContext('q', {}))) + return session + + def test_fit_logs_registers_and_serves_the_model(table_and_snapshot, quick_plan, tmp_path, caplog): from kumo.engine.integrations.mlflow import MlflowTracking, load_model @@ -240,20 +263,13 @@ def test_fit_waits_for_model_logging_and_fails_the_run_if_it_fails( import time from kumo.engine import _tracking - from kumo.engine._tracking import TrackingSession - from kumo.engine.trainer._run_spec import FitContext, fit_spec monkeypatch.setattr(_tracking, 'FINISH_SECONDS', 0.01) - class Tracker: + class Tracker(_Tracker): def __init__(self, fail): - self.fail, self.calls = fail, [] - - def start(self, record): - self.calls.append('start') - - def on_trial(self, record, event): - pass + super().__init__() + self.fail = fail def log_model(self, bundle_dir, manifest, snapshot, registry): time.sleep(0.2) @@ -261,21 +277,17 @@ def log_model(self, bundle_dir, manifest, snapshot, registry): if self.fail: raise RuntimeError('upload failed') - def finish(self, record, error): - self.calls.append(('finish', error)) - - result = SimpleNamespace(metrics={}, bundle_dir='missing', bundle_id='b') for fail in (False, True): tracker = Tracker(fail) - session = TrackingSession(tracker) - session.start(lambda: fit_spec(FitContext('q', {}))) + session = _started(tracker) with warnings.catch_warnings(record=True) as caught: warnings.simplefilter('always') - _, error = session.log_model(result.bundle_dir, None, None) + _, error = session.log_model('missing', None, None) session.finish(None, error) session._worker.join(5) - assert tracker.calls[:2] == ['start', 'log_model'] - assert tracker.calls[2] == ('finish', error) + started, logged, (finished, _, ended) = tracker.calls + assert (started, logged, finished, ended) == ('start', 'log_model', + 'finish', error) assert (error is not None) is fail assert any('upload failed' in str(w.message) for w in caught) is fail @@ -285,30 +297,16 @@ def test_model_logging_wait_is_bounded(monkeypatch, stall): import threading from kumo.engine import _tracking - from kumo.engine._tracking import TrackingSession - from kumo.engine.trainer._run_spec import FitContext, fit_spec monkeypatch.setattr(_tracking, 'MODEL_SECONDS', 0.5) release = threading.Event() - class Tracker: - calls = [] - - def start(self, record): - pass - - def on_trial(self, record, event): - pass - + class Tracker(_Tracker): def log_model(self, bundle_dir, manifest, snapshot, registry): release.wait(30) - def finish(self, record, error): - self.calls.append((record, error)) - tracker = Tracker() - session = TrackingSession(tracker) - session.start(lambda: fit_spec(FitContext('q', {}))) + session = _started(tracker) if stall == 'dead thread': session._submit(lambda: (_ for _ in ()).throw(SystemExit())) session._worker.join(5) @@ -319,7 +317,7 @@ def finish(self, record, error): assert type(error) is (TimeoutError if stall == 'slow upload' else RuntimeError) - assert tracker.calls == [(None, None)] + assert tracker.calls == ['start', ('finish', None, None)] def test_model_upload_gets_a_longer_http_timeout_and_a_late_one_is_dropped( @@ -380,21 +378,14 @@ def requirement(encoders): def test_a_tracker_dropped_earlier_says_the_model_was_not_logged(): import time - from kumo.engine._tracking import TrackingSession - from kumo.engine.trainer._run_spec import FitContext, fit_spec - - class Tracker: + class Tracker(_Tracker): def start(self, record): raise RuntimeError('server down') def log_model(self, bundle_dir, manifest, snapshot, registry): raise AssertionError('not called') - def finish(self, record, error): - pass - - session = TrackingSession(Tracker()) - session.start(lambda: fit_spec(FitContext('q', {}))) + session = _started(Tracker()) deadline = time.monotonic() + 10 while session._tracker is not None and time.monotonic() < deadline: time.sleep(0.01) diff --git a/packages/kumo-engine/test/test_mlflow_tracking.py b/packages/kumo-engine/test/test_mlflow_tracking.py index 48c9156ca..32e841513 100644 --- a/packages/kumo-engine/test/test_mlflow_tracking.py +++ b/packages/kumo-engine/test/test_mlflow_tracking.py @@ -433,7 +433,7 @@ def test_system_metrics_sample_every_second_unless_the_environment_says( batches = [] with integration._SystemMetrics(batches.append) as metrics: pass - assert batches == [] # fewer than 2 readings + assert batches == [] assert metrics._interval == 1.0 monkeypatch.setenv('MLFLOW_SYSTEM_METRICS_SAMPLING_INTERVAL', '2') assert integration._SystemMetrics(batches.append)._interval == 2.0 From fccff2f05f87c089d2c8c605951d1c445f79a49e Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 3 Nov 2026 01:24:54 +0000 Subject: [PATCH 273/278] Test summary routing beside the summary table tests Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/test/test_summary_tables.py | 6 +++++- packages/kumo-engine/test/test_tracking_contract.py | 6 ------ 2 files changed, 5 insertions(+), 7 deletions(-) diff --git a/packages/kumo-engine/test/test_summary_tables.py b/packages/kumo-engine/test/test_summary_tables.py index afe37e5f5..fdff90149 100644 --- a/packages/kumo-engine/test/test_summary_tables.py +++ b/packages/kumo-engine/test/test_summary_tables.py @@ -6,7 +6,7 @@ from kumo_core.summaries import Table -from kumo.engine._summary_tables import table_json +from kumo.engine._summary_tables import Routed, route, table_json def test_tables_round_measurements_and_keep_definitions_and_counts(): @@ -17,3 +17,7 @@ def test_tables_round_measurements_and_keep_definitions_and_counts(): assert json.loads(table_json(table))['data'] == [ [0.07, 12, 0.487744], [0.3333333333, 7, 27.6108]] + + +def test_a_summary_type_no_tracker_knows_is_skipped(): + assert route([object()]) == Routed() diff --git a/packages/kumo-engine/test/test_tracking_contract.py b/packages/kumo-engine/test/test_tracking_contract.py index 92b2fb36e..0e53d39fa 100644 --- a/packages/kumo-engine/test/test_tracking_contract.py +++ b/packages/kumo-engine/test/test_tracking_contract.py @@ -514,9 +514,3 @@ def start(self, record): assert elapsed < 1.5 assert isinstance(session.error, TimeoutError) - - -def test_a_summary_type_no_tracker_knows_is_skipped(): - from kumo.engine._summary_tables import Routed, route - - assert route([object()]) == Routed() From 6f0974f3ae269e7d55ebcbe148ce334699807345 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 3 Nov 2026 01:54:24 +0000 Subject: [PATCH 274/278] Record the Snowflake experiment on the fit's tracking ref Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/integrations/snowflake.py | 3 ++- packages/kumo-engine/test/test_snowflake_tracking.py | 5 ++++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/integrations/snowflake.py b/packages/kumo-engine/kumo/engine/integrations/snowflake.py index 11196a299..62f7882f3 100644 --- a/packages/kumo-engine/kumo/engine/integrations/snowflake.py +++ b/packages/kumo-engine/kumo/engine/integrations/snowflake.py @@ -127,7 +127,7 @@ def check(self, registry: Any = None) -> None: def start(self, record: RunRecord) -> None: self._reset() - self.run_id = self.logged_run_name = None + self.run_id = self.logged_run_name = self.experiment = None cls = self._module.ExperimentTracking tracking = getattr(cls, '_instance', None) if tracking is None: @@ -143,6 +143,7 @@ def start(self, record: RunRecord) -> None: experiment = tracking.set_experiment(self.experiment_name, database_name=database, schema_name=schema) + self.experiment = str(experiment.name) existing = [ str(row['name']) for row in tracking._sql_client.show_runs_in_experiment( diff --git a/packages/kumo-engine/test/test_snowflake_tracking.py b/packages/kumo-engine/test/test_snowflake_tracking.py index e1c094a8c..5cbcb6feb 100644 --- a/packages/kumo-engine/test/test_snowflake_tracking.py +++ b/packages/kumo-engine/test/test_snowflake_tracking.py @@ -12,7 +12,8 @@ from _snowflake_experiment_fake import FakeExperimentTracking, FakeSession from test_tracking_contract import CONTEXT, EVENTS, RESULT -from kumo.engine._tracking import RunRecord, TrackingSession +from kumo.engine._manifest import TrackingRef +from kumo.engine._tracking import RunRecord, TrackingSession, tracking_ref from kumo.engine.trainer._run_spec import fit_spec, fitted_spec FAKED = ('__init__', 'set_experiment', 'start_run', 'end_run', 'log_params', @@ -162,6 +163,8 @@ def test_an_existing_run_is_never_resumed(fake, taken, run_name, new, caplog): _replay(tracking) assert tracking.run_id == tracking.logged_run_name == new + assert tracking_ref(tracking) == TrackingRef( + platform='snowflake', run_id=new, experiment='ENTS', run_name=new) assert set(user.runs) == {taken, new} assert user.runs[taken]['params'] == {} assert user.runs[new]['status'] == 'FINISHED' From eb9d922d4f3a943850582094fbe576eb13350ac7 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 3 Nov 2026 03:41:27 +0000 Subject: [PATCH 275/278] Check the trained clock before building the prediction table Co-Authored-By: Claude Opus 5.5 --- .../kumo-engine/kumo/engine/trainer/trainer.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 9c7c9ec04..979964962 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -858,6 +858,13 @@ def predict( raise ValueError('A prediction table already holds its entities ' 'and times; pass it or entities=/' 'prediction_time=, not both.') + if self._training_result is None: + raise RuntimeError( + "Trainer.predict() requires a prior Trainer.fit() call on " + "this same Trainer instance.") + check_prediction_precision( + self._timestamp_precision, + index_timestamp_precision(snapshot.index_path)) table = (prediction_table if prediction_table is not None else self._prediction_table(snapshot, entities, prediction_time)) return self._predict( @@ -903,10 +910,6 @@ def _predict( num_workers: int = 1, binary_classification_threshold: Optional[float] = None, ) -> PredictionResult: - if self._training_result is None: - raise RuntimeError( - "Trainer.predict() requires a prior Trainer.fit() call on " - "this same Trainer instance.") from kumo_core.pquery import normalize_graph, parse_and_validate_query from kumo_core.prediction_options import ( prediction_query_contract, @@ -915,9 +918,6 @@ def _predict( ) from kumo_core.types import TaskType - check_prediction_precision( - self._timestamp_precision, - index_timestamp_precision(snapshot.index_path)) graph_config = snapshot.graph_config contract = prediction_query_contract(prediction_table.query, graph_config) From fb03edd867ae1127fdb29561d77dd17626f9244a Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 3 Nov 2026 18:26:59 +0000 Subject: [PATCH 276/278] Bound Databricks volume upload buffering Pass an explicit 16 MiB part size and 4 upload threads to files.upload when the SDK accepts them (databricks-sdk 0.70+), so a large plain object no longer buffers up to parallelism x part_size with the SDK's defaults. Older SDKs upload single-threaded in 10 MiB chunks and get no extra args. Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_databricks/volume.py | 13 +++++++++++-- packages/kumo-databricks/tests/test_volume.py | 17 +++++++++++++++++ 2 files changed, 28 insertions(+), 2 deletions(-) diff --git a/packages/kumo-databricks/src/kumo_databricks/volume.py b/packages/kumo-databricks/src/kumo_databricks/volume.py index 55e486107..02b635197 100644 --- a/packages/kumo-databricks/src/kumo_databricks/volume.py +++ b/packages/kumo-databricks/src/kumo_databricks/volume.py @@ -5,6 +5,7 @@ from __future__ import annotations +import inspect import io import posixpath import queue @@ -22,6 +23,7 @@ _FILES = '/api/2.0/fs/files' _BUFFER = 8 << 20 +_UPLOAD_BOUNDS = {'part_size': 16 << 20, 'parallelism': 4} def is_volume_path(path: str) -> bool: @@ -30,6 +32,13 @@ def is_volume_path(path: str) -> bool: and all(parts[1:]) and posixpath.normpath(path) == path.rstrip('/')) +def _upload(client: WorkspaceClient, path: str, source: BinaryIO) -> None: + upload = client.files.upload + bounds = (_UPLOAD_BOUNDS if 'parallelism' in inspect.signature( + upload).parameters else {}) + upload(path, source, overwrite=True, **bounds) + + def volume_filesystem(client: WorkspaceClient) -> FileSystem: return PyFileSystem(VolumeHandler(client)) @@ -83,7 +92,7 @@ def _list(self, directory: str) -> Dict[str, int]: def _put(self, local: str, path: str) -> None: with _translated(path), open(local, 'rb') as contents: - self._client.files.upload(path, contents, overwrite=True) + _upload(self._client, path, contents) def _stream(self, path: str) -> BinaryIO: return _Upload(self._client, path) @@ -189,7 +198,7 @@ def __init__(self, client: WorkspaceClient, path: str) -> None: def upload() -> None: try: with _translated(path): - client.files.upload(path, source, overwrite=True) + _upload(client, path, source) except BaseException as error: self._error = error diff --git a/packages/kumo-databricks/tests/test_volume.py b/packages/kumo-databricks/tests/test_volume.py index 266403c43..fde116ede 100644 --- a/packages/kumo-databricks/tests/test_volume.py +++ b/packages/kumo-databricks/tests/test_volume.py @@ -72,6 +72,23 @@ def test_unsized_writes_stay_plain_through_a_part(): assert f'{ROOT}/over.csv' not in volume.objects +def test_uploads_bound_the_sdk_part_buffers(): + fs, volume = filesystem(part_bytes=1000) + calls = [] + + def upload(path, contents, *, overwrite, part_size=None, + use_parallel=True, parallelism=None): + calls.append((part_size, parallelism)) + volume.objects[path] = contents.read() + + volume.upload = upload + write(fs, f'{ROOT}/export.csv', BODY[:1000]) + with open_sized(fs, f'{ROOT}/sized.csv', 1000) as writer: + writer.write(BODY[:1000]) + assert calls == [(16 << 20, 4)] * 2 + assert volume.objects[f'{ROOT}/export.csv'] == BODY[:1000] + + def test_a_failed_upload_raises_without_the_signed_url(): fs, volume = filesystem() signed = 'https://bucket.example/part?X-Amz-Signature=secret' From ca72e13b647f94250cb3c346d96290e2246843e9 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 3 Nov 2026 21:04:03 +0000 Subject: [PATCH 277/278] Log a run's inputs in role order and show the snapshot's tables The MLflow runs list shows a run's first input, so a fit logs its training table, then its graph snapshot, then its source tables by name. The snapshot input lists its id and each table's row count as its schema. Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/_data_summaries.py | 18 ++++-- packages/kumo-engine/kumo/engine/_manifest.py | 6 +- .../kumo/engine/integrations/mlflow.py | 4 +- packages/kumo-engine/test/test_artifacts.py | 2 +- .../kumo-engine/test/test_mlflow_tracking.py | 60 +++++++++++++++++++ 5 files changed, 78 insertions(+), 12 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_data_summaries.py b/packages/kumo-engine/kumo/engine/_data_summaries.py index b34f065ee..71067f20e 100644 --- a/packages/kumo-engine/kumo/engine/_data_summaries.py +++ b/packages/kumo-engine/kumo/engine/_data_summaries.py @@ -253,18 +253,24 @@ def table_details(index: Path, snapshot_id: Optional[str] = None) -> Details: rows = (_read(index / 'ingest_stats.json').get('stats') or {}).get('node_counts') or {} profiles = _profiles(_read(index / 'col_stats.json')) - tables = _read(index / 'diskgraph_config.json').get('tables') or () - snapshot = {} if snapshot_id is None else {snapshot_id: ({}, { - 'tables': len(tables), - 'rows': sum(rows.get(table.get('name'), 0) for table in tables)})} + tables = {table.get('kumo_source_table_name') or table.get('name'): table + for table in _read(index / 'diskgraph_config.json').get( + 'tables') or ()} + counts = {name: rows.get(table.get('name'), 0) + for name, table in tables.items()} + snapshot = {} if snapshot_id is None else {snapshot_id: ( + {'snapshot id': {'type': snapshot_id}, + **{name: {'type': f'{count:,} rows'} + for name, count in counts.items()}}, + {'tables': len(tables), 'rows': sum(counts.values())})} return {**snapshot, **{ - table.get('kumo_source_table_name') or table.get('name'): ( + name: ( {column['name']: {'stype': column.get('semantic_type'), 'dtype': column.get('source_dtype')} for column in table.get('columns') or ()}, {'rows': rows.get(table.get('name')), 'columns': profiles.get(table.get('name'), {})}) - for table in tables + for name, table in tables.items() }} diff --git a/packages/kumo-engine/kumo/engine/_manifest.py b/packages/kumo-engine/kumo/engine/_manifest.py index 32e37bd99..2f12ef9ca 100644 --- a/packages/kumo-engine/kumo/engine/_manifest.py +++ b/packages/kumo-engine/kumo/engine/_manifest.py @@ -227,8 +227,8 @@ def artifact_summaries(manifest: Manifest, roles: Mapping[str, str], (item.kind, item.id, str(item.location)) for item in located), brief=True)] if located else []), ] - for item in sorted(located, key=lambda item: _DATASET_ORDER.get( - item.kind, len(_DATASET_ORDER))): + for item in sorted(located, key=lambda item: (_DATASET_ORDER.get( + item.kind, len(_DATASET_ORDER)), item.label or item.id)): if not item.is_dataset: continue location = str(item.location) @@ -242,7 +242,7 @@ def artifact_summaries(manifest: Manifest, roles: Mapping[str, str], return summaries -_DATASET_ORDER = {'snapshot': 0, 'training_table': 1, 'source_table': 2} +_DATASET_ORDER = {'training_table': 0, 'snapshot': 1, 'source_table': 2} def _connector(location: str) -> Optional[str]: diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index 56cc046dc..ca4704667 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -502,8 +502,8 @@ def _dataset_input(self, summary: Dataset) -> Any: source_type=source._get_source_type(), source=source.to_json(), schema=json.dumps({'mlflow_colspec': [ - {'name': name, 'type': _COLUMN_TYPES.get(column['dtype'], - 'string')} + {'name': name, 'type': column.get('type') or + _COLUMN_TYPES.get(column.get('dtype'), 'string')} for name, column in summary.schema.items()]}), profile=json.dumps({ key: value for key, value in summary.profile.items() diff --git a/packages/kumo-engine/test/test_artifacts.py b/packages/kumo-engine/test/test_artifacts.py index bad38920f..01566f2e8 100644 --- a/packages/kumo-engine/test/test_artifacts.py +++ b/packages/kumo-engine/test/test_artifacts.py @@ -152,8 +152,8 @@ def test_a_fits_datasets_are_named_and_labelled_by_role(): assert [(s.name, s.context, s.source['version']) for s in artifact_summaries(model, FIT_ROLES) if isinstance(s, Dataset)] == [ - ('graph snapshot 2026-11-01-5ac245', 'graph context', 'g@5ac245'), ('training table 2026-11-01-fadf04', 'training examples', 'tt-fadf04'), + ('graph snapshot 2026-11-01-5ac245', 'graph context', 'g@5ac245'), ('users', 'source table', '3'), ] diff --git a/packages/kumo-engine/test/test_mlflow_tracking.py b/packages/kumo-engine/test/test_mlflow_tracking.py index 32e841513..0e01ecff7 100644 --- a/packages/kumo-engine/test/test_mlflow_tracking.py +++ b/packages/kumo-engine/test/test_mlflow_tracking.py @@ -365,6 +365,66 @@ def test_source_tables_are_logged_as_native_inputs(monkeypatch): assert mlflow.data.get_source(orders).delta_table_id == 'table-id' +def _fit_manifest(): + from kumo.engine._manifest import ( + ModelManifest, + ModelParents, + SnapshotManifest, + SnapshotParents, + SourceTableEntry, + TrainingTableManifest, + TrainingTableParents, + ) + + return ModelManifest(id='m', parents=ModelParents( + training_table=TrainingTableManifest( + id='tt-fadf04', created_at='2026-11-01T10:00:00+00:00', + location='s3://root/training_tables/tt-fadf04', + parents=TrainingTableParents(snapshot=SnapshotManifest( + id='g@5ac245', created_at='2026-11-01T09:00:00+00:00', + location='s3://root/snapshots/g/5ac245', + parents=SnapshotParents(source_tables={ + name: SourceTableEntry( + connector='S3', location=f's3://data/{name}', + version='3', version_kind='exact') + for name in ('users', 'orders')})))))) + + +def test_a_fits_inputs_are_logged_in_role_order(): + from kumo.engine._manifest import FIT_ROLES, artifact_summaries + from kumo.engine.integrations.mlflow import MlflowTracking + + replay(MlflowTracking(log_system_metrics=False), + summaries=artifact_summaries(_fit_manifest(), FIT_ROLES)) + + assert [item.dataset.name for item in + _fit_run().inputs.dataset_inputs] == [ + 'training table 2026-11-01-fadf04', + 'graph snapshot 2026-11-01-5ac245', 'orders', 'users'] + + +def test_the_snapshot_input_lists_its_tables_and_rows(tmp_path): + from kumo.engine._data_summaries import table_details + from kumo.engine._manifest import FIT_ROLES, artifact_summaries + from kumo.engine.integrations.mlflow import MlflowTracking + + (tmp_path / 'ingest_stats.json').write_text(json.dumps( + {'stats': {'node_counts': {'users': 300, 'orders': 6000}}})) + (tmp_path / 'diskgraph_config.json').write_text(json.dumps( + {'tables': [{'name': 'users'}, {'name': 'orders'}]})) + replay(MlflowTracking(log_system_metrics=False), + summaries=artifact_summaries(_fit_manifest(), FIT_ROLES, + table_details(tmp_path, 'g@5ac245'))) + + snapshot, = [item.dataset for item in _fit_run().inputs.dataset_inputs + if item.dataset.name.startswith('graph snapshot')] + assert json.loads(snapshot.schema) == {'mlflow_colspec': [ + {'name': 'snapshot id', 'type': 'g@5ac245'}, + {'name': 'users', 'type': '300 rows'}, + {'name': 'orders', 'type': '6,000 rows'}]} + assert json.loads(snapshot.profile) == {'tables': 2, 'rows': 6300} + + def test_link_prediction_metric_names_are_logged(): from kumo.engine._tracking import RunRecord from kumo.engine.integrations.mlflow import MlflowTracking From e4d98c031bd1d2cb02de01b028137ef83835fa6d Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 3 Nov 2026 21:14:57 +0000 Subject: [PATCH 278/278] Profile the snapshot input instead of listing its tables as columns Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_data_summaries.py | 9 ++++----- packages/kumo-engine/kumo/engine/integrations/mlflow.py | 4 ++-- packages/kumo-engine/test/test_mlflow_tracking.py | 9 ++++----- 3 files changed, 10 insertions(+), 12 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_data_summaries.py b/packages/kumo-engine/kumo/engine/_data_summaries.py index 71067f20e..a2cadc14f 100644 --- a/packages/kumo-engine/kumo/engine/_data_summaries.py +++ b/packages/kumo-engine/kumo/engine/_data_summaries.py @@ -258,11 +258,10 @@ def table_details(index: Path, snapshot_id: Optional[str] = None) -> Details: 'tables') or ()} counts = {name: rows.get(table.get('name'), 0) for name, table in tables.items()} - snapshot = {} if snapshot_id is None else {snapshot_id: ( - {'snapshot id': {'type': snapshot_id}, - **{name: {'type': f'{count:,} rows'} - for name, count in counts.items()}}, - {'tables': len(tables), 'rows': sum(counts.values())})} + snapshot = {} if snapshot_id is None else {snapshot_id: ({}, { + 'snapshot': snapshot_id, 'tables': len(tables), + 'rows': sum(counts.values()), + **{f'{name} rows': count for name, count in counts.items()}})} return {**snapshot, **{ name: ( {column['name']: {'stype': column.get('semantic_type'), diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index ca4704667..56cc046dc 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -502,8 +502,8 @@ def _dataset_input(self, summary: Dataset) -> Any: source_type=source._get_source_type(), source=source.to_json(), schema=json.dumps({'mlflow_colspec': [ - {'name': name, 'type': column.get('type') or - _COLUMN_TYPES.get(column.get('dtype'), 'string')} + {'name': name, 'type': _COLUMN_TYPES.get(column['dtype'], + 'string')} for name, column in summary.schema.items()]}), profile=json.dumps({ key: value for key, value in summary.profile.items() diff --git a/packages/kumo-engine/test/test_mlflow_tracking.py b/packages/kumo-engine/test/test_mlflow_tracking.py index 0e01ecff7..739354b3b 100644 --- a/packages/kumo-engine/test/test_mlflow_tracking.py +++ b/packages/kumo-engine/test/test_mlflow_tracking.py @@ -418,11 +418,10 @@ def test_the_snapshot_input_lists_its_tables_and_rows(tmp_path): snapshot, = [item.dataset for item in _fit_run().inputs.dataset_inputs if item.dataset.name.startswith('graph snapshot')] - assert json.loads(snapshot.schema) == {'mlflow_colspec': [ - {'name': 'snapshot id', 'type': 'g@5ac245'}, - {'name': 'users', 'type': '300 rows'}, - {'name': 'orders', 'type': '6,000 rows'}]} - assert json.loads(snapshot.profile) == {'tables': 2, 'rows': 6300} + assert json.loads(snapshot.schema) == {'mlflow_colspec': []} + assert json.loads(snapshot.profile) == { + 'snapshot': 'g@5ac245', 'tables': 2, 'rows': 6300, + 'users rows': 300, 'orders rows': 6000} def test_link_prediction_metric_names_are_logged():